RAG over Thinking Traces Can Improve Reasoning Tasks
Negar Arabzadeh, Wenjie Ma, Sewon Min, Matei Zaharia
cs.IR, cs.AI, cs.CL
2026-05-05
Berkeley RAG over thinking traces (T³) lifts Gemini-2.5-Flash on AIME 2025–2026 from 53.3 to 83.3 (+56.3%), beating web-corpus retrieval.
RAG is the default way to feed facts into LLMs. On math and code, the same recipe has been unreliable. Prior work that retrieves from general web corpora mostly helps weaker models, and noisy retrieved context can derail a solution. That pattern has fed a working assumption: retrieval is for knowledge, not for reasoning.
The Berkeley group argues the failure is the corpus, not the retriever. Web pages and textbooks supply facts. Reasoning needs process: how a solution was derived. They store thinking traces from strong reasoners as the retrieval collection, then rewrite those long, noisy traces offline into compact scaffolds (T³). The solver and the retrieve-then-generate loop stay unchanged. Only the documents change.
Two stages, with a vanilla RAG loop at inference.
Offline, a thinker (Gemini-2-thinking in the main setup) solves an auxiliary problem set and dumps raw traces. Gemini-2-Flash-Lite then rewrites each trace into one of three retrieval units:
All three rewrites are query-independent, so they run once. On the 59K corpus, each variant costs at most $26.74; all three together cost $64.74.
At inference, e5-base-v2 retrieves the top-3 units. Raw traces are split into 512-token chunks. Transformed traces average under 1,000 words and are retrieved as 2,048-token units. The hits are concatenated as hints in front of the question. The thinker, the rewriter, and the solver can be three different models.
Two corpora. T³-59k: about 59K problems with Gemini-2-thinking traces from the S1 pipeline, roughly 53K of them math. T³-114K: 114K problems with QwQ-32B traces from OpenThoughts (89K math, 20K code, 4K science, 1K puzzles). Both are decontaminated against the eval sets with a 13-gram Jaccard filter, dropping about 1.8%. AIME 2025–2026 also post-dates the source collections.
Baselines include no retrieval, random retrieval, OpenWebMath (6.4M), StackExchange (29.8M), two Wikipedia dumps, GitHub, ArXiv, the full CompactDS store (639M+ documents via DS-Serve and Contriever), and Tavily live search. General-purpose corpora use the same 512-token chunks and the same encoder.
Three reasoning benchmarks: AIME 2025–2026 (30 questions per year, Average@8), GPQA-Diamond (198 questions, Average@4), and LiveCodeBench (202 problems, pass@1 Average@4). Solvers: GPT-5, GPT-OSS-120B, Gemini-2.5-Flash. Generation cap 16K tokens, temperature 0.6.
AIME is where the gap is largest:
| Setting | GPT-5 | GPT-OSS-120B | Gemini-2.5-Flash |
| No RAG | 86.7 | 78.3 | 53.3 |
| Chunked raw traces | 91.7 (+5.8%) | 85.0 (+8.6%) | 80.0 (+50.1%) |
| Best T³ variant | Reflect 93.3 (+7.6%) | Semantic 83.3 (+6.4%) | Semantic 83.3 (+56.3%) |
| Best web corpus | GitHub 91.7 | CompactDS 80.0 | GitHub / Tavily 60.0 |
Full unchunked traces often lose to chunks: GPT-OSS-120B drops from 78.3 to 73.3 on AIME. Web corpora are unstable. OpenWebMath sends GPT-OSS-120B to 63.3 on AIME (-19.2%) and 37.6 on LiveCodeBench (-35.1%). CompactDS at 639M+ documents still does not dominate a 59K-trace index.
On GPQA-Diamond, T³-Struct lifts GPT-5 from 83.8 to 87.4 and Gemini-2.5-Flash from 77.3 to 80.8; T³-Semantic lifts GPT-OSS-120B from 70.7 to 74.7. On LiveCodeBench, T³-Struct lifts GPT-OSS-120B from 57.9 to 61.4. Science and code gains are smaller than math, matching the math-heavy corpus.
Two ablations. Retrieving final answers is weaker than retrieving traces on AIME for all three solvers; Gemini-2.5-Flash prefers answers on GPQA and LiveCodeBench. Holding the 59K problem set fixed, Gemini-2-thinking is the best thinker. GPT-OSS-120B traces crash GPT-5 on GPQA from 83.8 to 52.0 and drop Gemini-2.5-Flash on AIME from 53.3 to 45.0. How the trace is written matters more than which problems were used.
The larger T³-114K set, after transformation, underperforms T³-59K on most AIME and GPQA settings. Size is not the driver; alignment between source problems and the target benchmark is.
Averaged across the three benchmarks, GPT-5 goes from 76.14 to 80.53 accuracy while cost falls from 1.22 to 1.04 cents per query (14.8% cheaper). GPT-OSS-120B goes from 68.99 to 74.82, 0.10 to 0.09 cents. Gemini-2.5-Flash spends more with T³ than with no retrieval, but still less than with full traces, at higher accuracy. In one AIME 2026 case, Gemini-2.5-Flash finished 0 of 8 attempts without retrieval and 7 of 8 with T³-Reflect.
For anyone already running RAG, this is a corpus swap, not a new serving topology. Thinking tokens that were paid for once can be stored, rewritten once, and reused as hints. Weaker solvers gain the most in relative terms; GPT-5 still moves from 86.7 to 93.3 on AIME.
This sits on a different axis from Buffer of Thoughts, Retrieval-of-Thought, and Retrieval-Augmented Thoughts. Those systems instantiate templates, traverse thought graphs, or revise a chain of thought online. T³ pushes the work into offline indexing. Inference stays top-k retrieve-then-generate. Operationally that is closer to replacing an index than to replacing an agent controller.
This is a correction to corpus choice, not a new reasoning algorithm. The math evidence is solid. Code and scientific QA are trailing.
The paper lists three. Only vanilla RAG was tested; iterative, adaptive, and mid-reasoning retrieval are left open. The corpora are heavily math-skewed, which likely caps GPQA and LiveCodeBench. The thinker effect is not fully isolated: there is no full cross of many thinkers on the same problem set with the same solver.
A few more caveats sit in the tables. The best rewrite (Struct vs Semantic vs Reflect) depends on the task and the solver, so a production system still needs a validation split. Full traces can hurt GPT-OSS-120B, and the wrong thinker can collapse scores, so a trace dump is not free lunch. Decontamination is 13-gram Jaccard; semantically close but lexically different items can slip through. AIME is 60 questions with Average@8; the +56.3% relative jump is on a 53.3 baseline, an absolute +30.0. Cost is reported as a three-benchmark average, not per task.