OLMoTrace: Tracing Language Model Outputs Back to Trillions of Training Tokens | Summary

This article explains the key points of OLMoTrace: Tracing Language Model Outputs Back to Trillions of Training Tokens.

• 2025-04-09 (arXiv)
• Liu, Jiacheng, Blanton, Taylor, Elazar, Yanai, Min, Sewon, Chen, YenSung, Chheda-Kothary, Arnavi, Tran, Huy, Bischoff, Byron, Marsh, Eric, Schmitz, Michael, et al.
• Allen Institute for AI, University of Washington, UC Berkeley, Stanford University
• Paper
• Github
• Project Page

Read this article in Korean

Summary

• OLMoTrace addresses the problem of locating exact passages shared by a language model’s generated response and its full training corpus. On 98 internal conversations averaging 458 response tokens, its matching and document-retrieval stages averaged 4.46 seconds against a multi-trillion-token corpus; this is a measured runtime under the reported deployment conditions, with no measured comparison against another tracing system.
• The system extends infini-gram, a text search engine for trillion-token corpora, with a parallel algorithm that finds maximal matching spans: exact passages not contained in another valid match under its word and sentence boundary constraints. It selects long, uncommon passages, retrieves training documents containing them, and ranks those documents for topical relevance.
• The final configuration achieved document-relevance scores of 1.82 for the first displayed document and 1.50 across the top five on a 0–3 scale, using a language model as an evaluator. Case studies demonstrate factual-source inspection, reused creative expressions, and matching mathematical calculations, but lexical overlap establishes neither causal influence nor factual support.

1 Introduction

OLMoTrace makes training-data overlap inspectable after a response has been generated. Its motivation is the computational difficulty of applying existing behavior-tracing methods to full training corpora containing trillions of tokens, the text units used for indexing and matching.

• The deployed system supports three flagship OLMo models, the paper’s open language models, including OLMo-2-32B-Instruct. It searches their pre-training, mid-training, and post-training data, covering initial training and subsequent training stages.
• Applying the system to another model requires access to its full training data. The interface is available at https://playground.allenai.org, and the core implementation is released under Apache 2.0 at https://github.com/allenai/infinigram-api.

The Celine Dion example connects highlighted response passages to training-document cards and an expanded document view. The relevance colors represent topical ranking rather than causal confidence.

2 System Description

The interface highlights matching response passages and displays retrieved training documents, using color saturation to indicate topical relevance. Users can select a passage to inspect its matching documents or select a document to locate its matching passages; document cards show 80-token snippets, with an extended 500-token context available for inspection.

• For OLMo-2-32B-Instruct, Table 1 reports 3164 M documents and 4611 B tokens counted with the Llama-2 tokenizer, the text-to-token converter used for the index. Mid-training sources already present in pre-training are excluded from both the reported mid-training statistics and the index.
• The three supported models share pre-training and mid-training data and have slightly different post-training data.

The index covers all three training stages: pre-training contributes 3081 M documents and 4575 B tokens, mid-training 81 M and 34 B, and post-training 1.7 M and 1.6 B. The reported total is 3164 M documents and 4611 B tokens, showing the scale and pre-training predominance of the searchable corpus.

3 The Inference Pipeline

The pipeline first tokenizes the response with the Llama-2 tokenizer and identifies maximal exact matches. Each match must occur in the training corpus, begin and end at complete-word boundaries, contain a period or newline only at its end, and be uncontained by another span satisfying those conditions.

• It then keeps K = ⌈0.05 × L⌉ spans, where L is the response length in tokens, ranked by the smallest span unigram probability: the product of the individual token probabilities estimated from the full training corpus.
• This selection favors longer passages containing uncommon tokens. Token probabilities are cached; the authors report that ranking by length alone produced less relevant retrieved documents.

For each retained span, the system retrieves up to 10 enclosing document snippets, randomly sampling 10 when more matches exist. It merges overlapping response spans and snippets belonging to the same document, then reranks documents using BM25, a lexical retrieval score, with the user prompt and response concatenated as the query.

• The retrieved document collection serves as the BM25 corpus. Documents receive high, medium, or low relevance labels, and a highlighted span inherits the highest relevance level among its enclosing documents.
• Span filtering and the per-span retrieval limit mean that displayed documents are a selected subset of exact matches.

The diagram separates exact-match discovery from span selection, document retrieval, merging, and relevance ranking. Its highlights and relevance levels were adjusted for illustration, so the example is a pipeline schematic.

3.1 Fast Span Computation

Fast span computation uses a suffix array, an index that orders all corpus suffixes lexicographically, stored on solid-state drives rather than loaded into memory. A suffix is the text from a chosen starting position to the end; for every response suffix, a single infini-gram FIND query identifies its longest matching prefix, the initial sequence shared with corpus text. When the full suffix is absent, only the two neighboring corpus suffixes at its insertion position need inspection.

• Queries for response suffixes run in parallel. A subsequent pass suppresses non-maximal spans by retaining spans whose ending positions exceed all previously encountered endings.
• The paper reports \(O(L\log N)\) time complexity and \(O(\log N)\) latency under full parallelization, where N is the corpus token count. Actual throughput is limited by disk input/output.
• Each index shard, a partition of the index, is limited to 500 B tokens. Queries run across shards in parallel, and a second FIND query locates occurrences of the selected longest prefix, benefiting from cached disk reads.

The absent response suffix falls between two indexed corpus suffixes, allowing the longest common prefix to be recovered by inspecting those neighbors. Parallel processing across response suffixes and suppression of contained spans explain the algorithm’s efficiency.

3.2 Benchmarking Inference Latency

The latency benchmark uses a Google Cloud Platform node running on central processing units, with 64 virtual CPUs, 256 GB of random-access memory, and 40 TB of solid-state-drive storage. Across 98 conversations collected from internal Ai2 Playground usage, responses averaged 458 tokens and pipeline steps 1–3 averaged 4.46 seconds per query.

• The measured stages include maximal-span computation, span filtering, and document-snippet retrieval. The benchmark does not separately quantify the remaining merging, reranking, or interface-rendering stages.
• The evaluation reports a mean on this internal sample, without a latency distribution, a controlled competing-system baseline, or a concurrent-load benchmark.

4 Analyses

On the same 98 conversations, selected spans before merging have a mean length of 10.4 tokens and a median of 10. BM25 scores are normalized by an empirical soft upper bound of 0.18 times the response’s character count, with normalized scores ≥ 0.7 labeled high relevance, scores between 0.5 and 0.7 medium relevance, and scores < 0.5 low relevance.

• Under these thresholds, 14% of retrieved documents and 19% of spans receive high-relevance labels.
• Retrieved documents come predominantly from pre-training: 96.7%, compared with 0.9% from mid-training and 2.4% from post-training. Of all retrieved documents, 0.9% come from supervised fine-tuning, training on example responses; 1.5% from direct preference optimization, preference-based post-training; and none from reinforcement learning with verifiable rewards, post-training using checkable rewards.
• The authors emphasize that these proportions depend on conversation topic; mathematical responses may retrieve more post-training documents.

A human expert scored the top five documents using a 0–3 topical-relevance rubric under an early system configuration, obtaining averages of 1.90 for the first document and 1.43 for the top five. LLM-as-a-Judge, evaluation by a large language model, using gpt-4o achieved a Spearman rank correlation of 0.73 with those human judgments and scored the same configuration at 1.73 and 1.28; this correlation measures agreement between rankings.

• After configuration tuning guided by the model evaluator, the final setting scored 1.82 for first documents and 1.50 for top-five documents; 63.3% and 55.1%, respectively, received scores of 2 or 3.
• Table 2 defines the rubric from unrelated material to a direct topic-and-scope match. Table 3 presents successive configuration changes rather than independent one-factor comparisons; the final configuration was not separately evaluated by the human expert.
• The evaluation measures topical fit to the prompt and response, rather than causal contribution or the truth of the response.

The histogram centers selected spans near 10 tokens, with a mean of 10.4. The score plots motivate response-length-dependent relevance thresholds and show that high-relevance labels cover 14% of documents and 19% of spans; these labels derive from BM25 scoring.

5 Case Studies

The three case studies illustrate what exact-match inspection can reveal: a training document repeats the statement that the Space Needle was built for the 1962 World Fair; a Tolkien-style story shares “I’m going on an adventure” with Hobbit fan fiction; and a solution to AIME 2024 I Problem 4, a combinatorics exam problem, contains a calculation also present in post-training data.

• The mathematical match is \(\binom{10}{4} = \frac{10!}{4!(10-4)!} = 210\). This shows that the calculation string occurs in training data, without establishing how the model acquired the broader problem-solving capability.
• For factual inspection, users must examine document context and source reliability. The screenshots demonstrate individual uses rather than a measured improvement in fact-checking accuracy.
• Figure 6 shows how users inspect individual matches: selecting a highlighted passage filters the document panel, and “Locate Span” restricts response highlights to those found in a selected document. Clicking again or using “Clear Selection” restores the full view.

The three panels show factual-source inspection, overlap between generated fiction and training text, and a mathematical calculation shared with post-training data. They provide examples of inspectable overlap without measuring causal attribution or fact-checking performance.

The paired screenshots demonstrate reciprocal navigation: selecting a response passage narrows the documents, and selecting a document narrows the response highlights. This interaction supports inspection of individual matches within a response containing several overlaps.

The rubric distinguishes unrelated material, broad or insufficiently specified relevance, correct-topic material with a context mismatch, and a direct topic-and-scope match. The corresponding model-evaluation prompt measures relevance to the user prompt and response rather than factual correctness.

The final configuration obtains 1.82 and 1.50 average relevance scores for first and top-five documents, with 63.3% and 55.1% scored relevant. Rows progressively reverse configuration choices, so their differences are sequential comparisons; the human-evaluation row concerns the early configuration.

6 Related Work

Retrieval-augmented generation, which conditions generation on retrieved documents, differs from OLMoTrace’s retrieval after generation. Web search searches a changing web index, whereas OLMoTrace searches the model’s training corpus.

• Influence functions use gradient information to identify influential training examples, but the paper describes their computational cost as prohibitive at trillion-token scale. OLMoTrace instead retrieves examples by exact lexical overlap.
• Other cited approaches train models to cite pre-training documents or supplied context. OLMoTrace requires neither a change to model training nor intervention in generation.

Limitations

Exact lexical matches must not be interpreted as causal explanations, supporting evidence, or citations for a model response. The system’s matching criterion leaves semantic relationships without verbatim overlap outside its tracing scope.

• Internal adversarial testing exposed news articles and song lyrics, although the authors did not encounter copyrighted books. They provide a takedown process and document removal without rebuilding the full index.
• The authors did not find personally identifiable information in their testing and implemented a regular-expression filter, matching specified text patterns, to block documents containing it. Toxicity mitigation relies on existing Ai2 Playground prompt moderation, with no additional filtering reported.

Author Contributions

The paper names Jiacheng Liu and Taylor Blanton as core contributors and separately credits research, engineering, design, user experience, project management, legal, communications, and advising roles.

Appendix

• A More Screenshots of OLMoTrace extends the interface demonstration with passage-to-document and document-to-passage selection. Clicking a highlighted passage filters the document panel, while “Locate Span” restricts highlights to passages found in the selected document; clicking again or using “Clear Selection” restores the full view.
• B Production System Setup describes pd-balanced storage supporting up to 80,000 read input/output operations per second per virtual machine, two N2 virtual-machine replicas, shared disk mounts, and separate processing workers. Prefetching, speculative reading of data before it is requested, is disabled because extra reads worsen the disk-throughput bottleneck, and document retrieval is batched. With 12 shards, the disk analysis estimates 40 × 2 × 12 = 960 reads per response token and approximately 1.2 seconds for 100 tokens; this is an analytical estimate rather than an additional measured benchmark.
• C More Details on Document Relevance Evaluation specifies the 0–3 rubric and gpt-4o-2024-08-06 evaluator. Table 3 traces successive changes: retaining spans occurring more than 10 times, ranking spans by unigram probability, increasing document context from 100 to 500 tokens, and including the user prompt in BM25 queries. The final top-five relevance score rises from the early configuration’s 1.28 to 1.50, while the first-document score rises from 1.73 to 1.82 under the same model evaluator.

Brief Thoughts

OLMoTrace makes exact training-data overlap interactively inspectable at 4.6-trillion-token scale, supported by a concrete indexing algorithm and a measured 4.46-second mean for matching and retrieval. The 98-conversation internal sample and evaluator-guided tuning limit conclusions about relevance beyond the evaluated setting, while exact matches remain evidence of overlap rather than causal explanations.