How TokenPath works

TL;DR: TokenPath takes a document and an existing AI answer, then finds the source text that supports each claim. It does not rewrite the answer. The first half of this post shows captured TokenPath examples; the second shows the basic calculation in a Google Colab notebook.

Attention already points at the source

TokenPath reads the document and answer together with an open model. It measures which parts of the document are most strongly connected to each part of the answer, then turns those connections into citations.

This model is separate from the model that wrote the answer. The answer can come from GPT-5.5, Claude, a fine-tune, or a human. Nothing is regenerated.

The idea has research behind it. Earlier work showed that a model's internals can attribute generated answers back to their source: hidden-state representations locate the segments an answer copied and where they came from, matching GPT-4 as an attributor without any retraining (Phukan et al., 2024), and raw attention logits give token-, sentence-, and chunk-level attributions with higher fidelity still (Mahapatra et al., 2025). TokenPath applies that approach to citation.

From attention to a citation

Take a four-sentence document about the Oregon Duck and the answer "The Oregon Duck wears green and yellow." In this example, the heatmap endpoint returned a sparse map connecting answer tokens to document tokens:

h = requests.post("https://api.tokenpath.ai/v1/attributions/heatmap",
                  headers=AUTH, json={"document": document,
                                      "question": question,
                                      "answer": answer}).json()

h["shape"]      # [8, 46]  — answer tokens × document tokens
len(h["data"])  # 174 retained weights

Here is that captured response, drawn cell by cell:

You can read the citation straight off the picture: green, and, yellow each put their mass on the exact words green and yellow in sentence three, and wears lands on wears. Summing the columns and pooling by sentence is all "make a citation" is:

[75.5%] The mascot wears a green and yellow costume and a green beanie cap.
[17.4%] The Oregon Duck is the mascot of the University of Oregon.

Look at which sentence wins. The answer says "the Oregon Duck," but the color fact lives in a sentence that says "The mascot." TokenPath points to the fact-bearing sentence even though it uses different words. A keyword or embedding search may prefer the sentence containing the full name. This pattern appears throughout our benchmark results.

The main /v1/attributions endpoint turns that map into a source span:

r = requests.post("https://api.tokenpath.ai/v1/attributions",
                  headers=AUTH, json={"document": document,
                                      "question": question, "answer": answer,
                                      "spans": [[22, 38]]}).json()

{
  "spans": [{
    "answer": {"start": 22, "end": 38, "text": "green and yellow"},
    "source": {
      "start": 115, "end": 131, "text": "green and yellow",
      "confidence": 0.88
    }
  }]
}

Four things string matching can't do

TokenPath can also connect wording that is similar without being identical. These cards show captured API results (code for all four is in the notebook).

Is this really the model's attention?

You should be skeptical of a vendor saying "trust us, it's attention." So here's the same basic calculation with raw 🤗 Transformers — no TokenPath in the loop. Everything from here down runs on Llama-3.1-8B-Instruct. The checkpoint requires access and suitable hardware. This is a standalone demonstration of the calculation, not a copy of the hosted API. The whole thing is one forward pass with attention exposed, then a slice:

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct", attn_implementation="eager")

# one forward pass over [prompt][document][answer], attention kept
out = model(input_ids, output_attentions=True)

attn = mean_over_heads(out.attentions, layers=range(14, 24))
heat = attn[answer_positions][:, document_positions]  # answer-to-document slice

Rolled up to sentences exactly as before, the Oregon Duck example gives:

[67.1%]  The mascot wears a green and yellow costume and a green beanie cap.
[18.3%]  The Oregon Duck is the mascot of the University of Oregon.

The Llama example also ranks "The mascot" sentence first. Its values differ from the API example because these are two separate implementations; they are not expected to match. The complete version (prompt framing, token bookkeeping, sentence rollup) is reproduce_attention.py, about 90 dependency-light lines.

Not all heads — the heads that attribute

Averaging all heads makes the signal mushy: most of the ~1,000 heads in an 8B model do local, syntactic work. A subset — cousins of the induction and retrieval heads from the interpretability literature — track "which source token is this answer drawing from." Finding them is a measurement too: bury a known fact among distractor sentences, ask about it, and score every (layer, head) by how much of the answer's attention lands on the fact.

On this small synthetic probe, the best heads put ~0.94 of the answer's attention on the buried fact, against 0.57 for the average head and 0.20 chance. This helps identify useful candidates; it is not a general citation accuracy score. The probe is find_attribution_heads.py — synthetic data you can read in the source.

The long-context part

The reproduction above uses attn_implementation="eager", and that's what makes it a toy. FlashAttention computes attention without ever materializing the seq × seq score matrix — that's why 100k-token context is affordable at all. Ask Transformers for output_attentions=True and it falls back to building the full matrix: billions of floats per layer on a long document.

The way out is that attribution needs almost none of that matrix — only answer rows (answers are short) against document columns, for the selected heads:

Getting that slice needs no custom kernel. The forward pass runs unmodified — FlashAttention and all — while the hidden states entering the selected layers are captured; those are tokens × hidden, nothing seq × seq about them. Attribution then replays just the attention arithmetic on the side: answer tokens projected to queries, document tokens to keys, a softmax over those few rows for the selected heads. Memory stays proportional to answer-tokens × sequence, not sequence². Generation is untouched; attribution is a cheap second read. (Background on why the full matrix never exists in the first place: the FlashAttention paper.)

How good is it — and what does it measure?

Quality. On LongBench-Cite — a benchmark we didn't write — this pipeline scores F1 0.815, ahead of Anthropic's generation-time Citations API (0.812), within 0.04 of a prompted frontier LLM (0.851), and far ahead of embedding retrieval (0.622). Compared with GPT-5.5, TokenPath was 6× faster and 7× cheaper. The full experiment, including where we still lose: One answer, four ways to cite it.

Run it yourself

• Notebook: how-tokenpath-works.ipynb (Colab) — every example above. The local Llama example requires checkpoint access and suitable hardware.
• Minimal scripts: reproduce_attention.py and find_attribution_heads.py.
• Benchmark: the harness and the results post.
• Papers: attribution from hidden states (Phukan et al., 2024) and from raw attention logits (Mahapatra et al., 2025). Questions, or a document where it breaks? We want to see it: support@tokenpath.ai.