Open-source library · paper · runnable examples

Thinking longer isn’t thinking better.
Thinking with examples is.

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.

Start tutorial → $ pip install "thinkretrieve[faiss]"
<think>
The two rolls give 36 outcomes… so the probability is 1/6. Let me double-check — yes, 1/6. ⤷ retrieved · similar solved problem (cos 0.87) Problem: Two fair dice are rolled…
Solution: count ordered pairs — (1,2) and (2,1) are distinct outcomes…
The example counts ordered pairs — I missed that. There are 2 favourable outcomes, not 1: the answer is 2/36 · 6 = 1/3 ✓
</think>
Results · 5 reasoning models × 4 benchmarks, 3-seed averages

Wins every cell. Costs less than it looks.

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.

20 / 20
(model × benchmark) cells won vs. standard thinking, TTS, static ICL, and random retrieval
▲13.4 pts
AIME 2025, Qwen3-1.7B — 22.2% → 35.6% at the same thinking budget
−18%
output tokens vs. TTS at equal budget (12,913 vs 15,847) — injected examples aren’t generated
~6%
wall-clock overhead — 47 ms/step for FAISS + E5 retrieval, on CPU

More thinking budget helps — only with retrieval

MATH-500 accuracy (%) vs. thinking budget · Qwen3-1.7B · 3 seeds

ThinkRetrieve Sequential TTS
86 88 90 92 8K 16K 22K 92.1 87.5
data table
BudgetSequential TTSThinkRetrieve
8K89.291.5
16K88.891.9
22K87.592.1

Largest gains where reasoning is hardest

AIME 2025 best accuracy (%) · 3 seeds

ThinkRetrieve Sequential TTS
0 20 40 60 80 30.0 35.6 34.4 66.7 71.1 DS-R1-1.5B Qwen3-1.7B Qwen3.5-2B Qwen3-4B Qwen3-8B
data table
ModelTTSThinkRetrieve
DeepSeek-R1-Qwen-1.5B23.330.0
Qwen3-1.7B22.235.6
Qwen3.5-2B31.134.4
Qwen3-4B64.466.7
Qwen3-8B68.971.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.

Evidence · beyond the headline numbers

Why we believe it’s the mechanism, not luck

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).

Beats compute-matched self-consistency

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.

Helps 13× more often than it hurts

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⁻¹²).

Not answer leakage

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.

Robust to the wrong corpus

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.

corpus-mismatch data (accuracy %, 3 seeds)
Benchmark · modelNo retrievalMatched bankCross-domain bank
MATH-500 · Qwen3-1.7B90.292.591.5
MATH-500 · Qwen3-8B93.294.893.8
SciQ · Qwen3-1.7B93.294.594.2
SciQ · Qwen3-8B95.896.896.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 (thinkretrieve/examples/compare_tts_vs_thinkretrieve.py — plain, prompt-RAG, TTS, and ThinkRetrieve over Ollama or Amazon Bedrock).

Method

How a step works

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.

STEP 1 · PROBE

Elicit an interim answer

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.

STEP 2 · RETRIEVE

Search solved examples

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.

STEP 3 · INJECT

Continue, grounded

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.

Don’t take our word for it

Run the comparison on your machine

The comparison downloads GSM8K and an E5 encoder, indexes training solutions, and evaluates four arms with the same model and same thinking budget. It is checkpointed and resumable, and the baseline is just retriever=None, so it is always apples-to-apples.

# Clone once
git clone https://github.com/itsvaibhav01/ThinkRetrieve.git
cd ThinkRetrieve

# Fast offline verification: no model, dataset, GPU, or API key
python -m pip install -e ./thinkretrieve
python thinkretrieve/examples/offline_dummy.py

# Real local smoke run: downloads qwen3:4b, GSM8K, and E5
python -m pip install -e "./thinkretrieve[faiss]" datasets
ollama pull qwen3:4b
python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py \
  --backend openai --model qwen3:4b \
  --limit 2 --bank-size 100 --budget 512

# Full defaults after the smoke run
python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py --backend openai --model qwen3:4b

# Amazon Bedrock alternative (AWS credentials and charges apply)
python thinkretrieve/examples/compare_tts_vs_thinkretrieve.py --backend bedrock --model qwen.qwen3-32b-v1:0

# 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)

① Upgrade your RAG — same index, new injection point

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.

② Procedural memory for agents

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.

③ A free upgrade for small local models

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.

Quickstart

Ten lines, any backend

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)

This local example requires ollama pull qwen3:4b and pip install "thinkretrieve[faiss]". The model is replaceable: use any compatible local server, hosted OpenAI-style API, Anthropic, or Bedrock.

OllamaLM StudioMLXllama.cpp vLLMOpenAIAnthropicAmazon Bedrock Together · Groq · DeepSeek · OpenRouter
Where it helps

Ready-made banks already exist for these domains

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(...):

Competition & school math — validated in the paper

NuminaMath-1.5 (860K problem–solution pairs — the paper’s bank), MATH (12.5K with full LaTeX solutions), GSM8K, MetaMathQA, OpenMathInstruct.

Science QA — validated in the paper

SciQ training split as the bank (+1.0–1.3 pts across five models on the SciQ test set).

Science-agent tool dispatch — runnable case study

On six simulated cell-culture incidents, diagnosis-driven retrieval selected the private SOP and correct tool route in 4/6 cases with a local Qwen 3.5 4B and 5/6 with 9B; the three controls scored 0–2/6. This is a small use-case demo, not a biological safety benchmark. Read the visual article → · runner and raw outputs

Competitive programming — untested, bank exists

open-r1/codeforces: 10K+ Codeforces problems with editorials and ~100K DeepSeek-R1 reasoning traces. If you try it, tell us what you find.

Knowledge-bound factual QA — measured: no gain

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.

Small local models

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.

Private corpora & synthesis

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).

Paper

Citation

ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling. The paper, installable library, guided tutorial, and runnable comparison examples are linked from this repository.

Vaibhav Singh, Soumya Suvra Ghosal, Sarvesh Gharat, Soumyabrata Pal, Ramasuri Narayanam, and Dinesh Manocha. arXiv preprint, 11 August 2026.

Paper abstract and citation details · Download BibTeX · Download RIS (Zotero / EndNote) · DOI

@article{thinkretrieve2026,
  title   = {ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling},
  author  = {Singh, Vaibhav and Ghosal, Soumya Suvra and Gharat, Sarvesh and
             Pal, Soumyabrata and Narayanam, Ramasuri and Manocha, Dinesh},
  journal = {arXiv preprint arXiv:2608.10928},
  year    = {2026},
  doi     = {10.48550/arXiv.2608.10928},
  url     = {https://arxiv.org/abs/2608.10928}
}