ThinkRetrieve is procedural memory for reasoning models: index solved problems once, and the model recalls the most similar worked solution mid-thought — at every reasoning step, right when it's about to go wrong. Higher accuracy, fewer generated tokens, any chat API, no GPU required.
Sequential test-time scaling (the “Wait, let me think again” approach) degrades as the thinking budget grows — models drift, cycle, and compound errors. Grounding each step with a retrieved solved example reverses the curve.
MATH-500 accuracy (%) vs. thinking budget · Qwen3-1.7B · 3 seeds
| Budget | Sequential TTS | ThinkRetrieve |
|---|---|---|
| 8K | 89.2 | 91.5 |
| 16K | 88.8 | 91.9 |
| 22K | 87.5 | 92.1 |
AIME 2025 best accuracy (%) · 3 seeds
| Model | TTS | ThinkRetrieve |
|---|---|---|
| DeepSeek-R1-Qwen-1.5B | 23.3 | 30.0 |
| Qwen3-1.7B | 22.2 | 35.6 |
| Qwen3.5-2B | 31.1 | 34.4 |
| Qwen3-4B | 64.4 | 66.7 |
| Qwen3-8B | 68.9 | 71.1 |
Full tables — GSM-8K, SciQ, cross-domain corpora, Qwen3-32B, ablations (random retrieval, static ICL, encoder and query-formulation choices), and decontamination & answer-leakage audits — are in the paper. A key ablation: input-level RAG with the same retrieved content is ≈ neutral (88.8% vs 89.1% baseline); injecting it inside the trace adds the gain. Where you inject matters more than what you retrieve.
Four controls from the paper separate the retrieval effect from confounds. All numbers are from the paper’s evaluation protocol (3 seeds, LLM-judge scoring).
At the same total token budget, one ThinkRetrieve trajectory outperforms majority voting over k = 2, 4, and 8 independent traces. The gain isn’t “more compute” — it’s where the compute goes.
Paired per-problem analysis (Qwen3-4B, MATH-500, B=12K): retrieval fixes 26.5% of the problems sequential TTS misses while breaking only 2.0% it got right — significant under McNemar’s exact test (p ≈ 6×10⁻¹²).
Two-stage decontamination (exact match + cosine > 0.90 on joint QA embeddings), a top-1 exemplar audit (zero shared answers), and a strict filter blocking any exemplar whose answer equals the gold — accuracy unchanged. The gains are procedural.
Even with a deliberately mismatched bank (science examples for math questions and vice-versa), accuracy stays at or above no-retrieval — so an imperfect domain bank fails soft, not hard.
| Benchmark · model | No retrieval | Matched bank | Cross-domain bank |
|---|---|---|---|
| MATH-500 · Qwen3-1.7B | 90.2 | 92.5 | 91.5 |
| MATH-500 · Qwen3-8B | 93.2 | 94.8 | 93.8 |
| SciQ · Qwen3-1.7B | 93.2 | 94.5 | 94.2 |
| SciQ · Qwen3-8B | 95.8 | 96.8 | 96.6 |
How this differs from Search-o1, RAT, and agentic-search methods: those retrieve documents and facts (what to know). ThinkRetrieve retrieves complete solved examples (how to reason) and injects them inside the live thinking trace. A reproduction script for the API path ships in the repo (scripts/bench_thinkretrieve_bedrock.py — AIME 2025, TTS vs ThinkRetrieve over Amazon Bedrock).
Whenever the model tries to end its reasoning — or thinks too long without stopping — ThinkRetrieve intervenes. The loop needs three chat calls and a vector index; nothing else.
Temporarily close the thinking block and ask for the current best answer. It’s used only as a retrieval signal, then discarded — it captures the model’s reasoning state, not just the question.
Embed query: {question} solution: {interim} (E5) and search a
FAISS bank of worked solutions by cosine similarity. QA-to-QA matching finds
problems with similar solution strategies, not similar phrasing.
Append the retrieved problem + solution into the trace with a reconsideration prompt and let the model keep thinking. Used examples are never repeated; injected tokens count against the thinking budget.
One command indexes GSM8K’s training solutions, answers test questions twice
with the same model and same thinking budget — plain long thinking
vs. ThinkRetrieve — and prints the table. Checkpointed, resumable, and the
baseline is just retriever=None,
so it’s always apples-to-apples.
# Amazon Bedrock: python examples/compare_tts_vs_thinkretrieve.py --backend bedrock --model qwen.qwen3-32b-v1:0 # or a local model on your laptop (Ollama): python examples/compare_tts_vs_thinkretrieve.py --backend openai --model qwen3:4b # Measured output — MacBook (24 GB), Qwen3-1.7B-4bit via MLX, # 60 GSM8K questions, identical 2,048-token thinking budget: # # accuracy thinking tok/q # plain query 38.3% 888 # standard RAG (prepended) 36.7% 921 # long thinking (TTS) 36.7% 1874 # thinkretrieve 45.0% 1794 ← +6.7 vs plain # # paired: fixes 6 plain-query failures, breaks 2 (3:1)
Prompt-level RAG is ≈ neutral on reasoning tasks; injecting the same
content inside the trace is what wins (paper §5). Wrap your existing
vector DB in a 5-line Retriever
and keep everything else.
Facts aren’t the only memory. After a task is solved and verified:
memory.add(task, solution) —
the agent recalls how it solved similar tasks, mid-reasoning,
next time. Measured: a memory of just 23 of the agent’s own
verified solutions rescued as many of its failures (6/37) as a
4,000-example curated bank.
Reasoning drift is worst in the 1–4B models people run on laptops — and that’s where the gains are biggest (+13.4 pts AIME on a 1.7B). Point the backend at Ollama/LM Studio/MLX and reuse any bank.
Full walkthrough of all three recipes: thinkretrieve/TUTORIAL.md in the repo.
from thinkretrieve import ThinkRetrieve, FaissRetriever, OpenAICompatBackend # 1 · an example bank = any solved (question, solution) pairs retriever = FaissRetriever.from_examples([ ("A jacket costs $120 and is discounted 25%. Final price?", "Discount = 0.25*120 = 30. Final = 120-30 = 90. Answer: $90"), # ... datasets, textbooks, runbooks, past tickets ... ]) # 2 · any OpenAI-compatible server or API (Ollama shown) backend = OpenAICompatBackend(model="qwen3:4b", base_url="http://localhost:11434/v1") # 3 · retrieval-augmented reasoning result = ThinkRetrieve(backend, retriever).run( "A phone costs $250 after a 20% discount. Original price?") print(result.answer, result.retrievals)
ThinkRetrieve needs one thing: a corpus of problems paired with worked,
step-by-step solutions. These are public, sizable, and drop straight into
FaissRetriever.from_examples(...):
NuminaMath-1.5 (860K problem–solution pairs — the paper’s bank), MATH (12.5K with full LaTeX solutions), GSM8K, MetaMathQA, OpenMathInstruct.
SciQ training split as the bank (+1.0–1.3 pts across five models on the SciQ test set).
open-r1/codeforces: 10K+ Codeforces problems with editorials and ~100K DeepSeek-R1 reasoning traces. If you try it, tell us what you find.
We tested MedReason medical MCQ (n=40, same setup as the GSM8K run): plain 35%, long thinking 30%, ThinkRetrieve 30%. When a question needs a missing fact rather than a reasoning procedure, exemplars can’t supply it — exactly the method’s stated scope (“how to reason, not what facts”). Know this boundary before adopting.
Reasoning drift is worst in 1–4B models — exactly what runs on a laptop via Ollama/MLX. A 1.7B model gains 13 points on AIME 2025.
Runbooks, solved tickets, graded homework, worked financial analyses — or generate question–solution pairs once with a strong model and retrieve from them forever. Mismatched banks fail soft (see Evidence).
ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling. The research harness — all baselines, trigger studies, graders, and audits — lives in the same repository as the library.
@article{thinkretrieve2026,
title = {ThinkRetrieve: Retrieval-Augmented Reasoning Traces
for Test-Time Scaling},
author = {Singh, Vaibhav and others},
journal = {arXiv preprint arXiv:XXXX.XXXXX},
year = {2026}
}