Amortized Reasoning Index

Reasoning models are excellent judges of relevance. As long as your agent gets to see the relevant documents somehow, it will probably identify them and use them well.

But what if there are simply too many documents to feed to your reasoning model? That’s usually the case when you work over a massive code repository or have your agents search the web. To this day, for hard enough queries, this search bottleneck results in poor recall [1], even when a frontier LLM agent is allowed to repeatedly query the best existing retrievers.

In this work, we ask: Is there a way to scale the judgment of a reasoning model to all documents in a large corpus? We think the answer might be yes.

We introduce the Amortized Reasoning Index (ARI), a new kind of retriever that updates its weights during search, using gradients derived at inference time from the judgments of a reasoning model.

Diagram of ARI over three rounds for the search query “Find users happy with the new iOS update.” In round 1, the retriever’s top three tweets are all irrelevant. The reasoning model ranks them anyway, and that ranking updates the retriever’s weights. In round 2, the updated retriever surfaces two tweets the reasoning model never ranked, one of them relevant. After a second update, all three top tweets in round 3 are relevant. Later queries start from the updated weights. — **Figure 1.** ARI conducts retrieval using test-time training. Even if the first results are all irrelevant, a reasoning model can still rank them, and training the retriever to match that local ranking allows it, at inference time, to find relevant tweets that were never ranked before.

At a high level, ARI works like this:

When a search query q arrives, a small retriever works over the full corpus and returns its top-K matches. If the query is hard, these top-K matches will be very poor, and may even contain no relevant documents!
ARI then asks a reasoning model to re-rank those top-K documents according to its relative preferences, and then uses a simple form of reinforcement learning with a contrastive reward to update the retrieval model’s weights at test time to encourage its rankings to align with the reasoning model.
Immediately, and still during the processing of the same query q, ARI uses the updated weights to search again. As a result of the gradient step, the retriever may now surface relevant documents that the reasoning model prefers but never actually saw before.

This process can then repeat, and finally the self-improved retriever is retained and further tuned live on future queries.

Even when the initial candidates are poor, our early results suggest that relative ranking between them can still teach the small retriever to surface better unjudged documents that the reasoning model never saw. And because the updated weights persist, reasoning spent on past queries improves retrieval for future ones without any relevance labels.

Figure 2. Reasoning about incoming queries improves retrieval for queries the reasoning model never sees. Recall@100 on held-out queries as ARI processes the query stream, with no relevance labels during adaptation and no reasoning-model calls at evaluation. ARI’s Recall@100 grows 3.2× on Twitter-Conflict, 6.7× on Writing and 7× on BRIGHT AoPS, ending above Qwen3-Embedding-8B on all three. Dashed lines mark Qwen3-Embedding-0.6B and Qwen3-Embedding-8B; the dash-dot line marks Gemini-Embedding-2; BRIGHT AoPS shows the mean of four train/test splits within a fixed 49,987-document pool.

Figure 2 summarizes some of our early results with ARI. On two “oblique” search tasks that guided the development of ARI and one fully held-out reasoning-intensive task, ARI allows a comparatively tiny 400M-parameter cross-encoder to self-improve in a very consistent manner on unseen test queries, from nothing but a stream of unlabeled queries. Even when the retriever trained by ARI starts extremely weak (e.g., below 10% on AoPS), it still has no trouble climbing rapidly and by up to 7× its initial Recall@100.

ARI can be understood as an extreme form of pseudo-relevance feedback [2, 3], in which we update the weights of the retrieval model itself at inference time, and not just the text or the vector of the query [4, 5, 6, 7]. Our experience developing ARI so far suggests that this allows for much richer learning, allows the learning to compose and persist across queries, and allows all the learning in ARI to happen at inference time over the (initially very weak!) pool of documents that the small retriever can surface by itself. We expect both of these aspects (updating the weights; learning at inference time and without labels) to become key to the future of retrieval going forward.

Going forward, we’re excited about turning ARI into a persistent memory for agents. Due to their limited context window, today’s agents go through redundant reasoning and issue similar queries each session. With ARI, each investigation by an agent over your data should leave the retriever better prepared to guide the next session.

A First Implementation

We now describe our first implementation of the Amortized Reasoning Index (ARI). For extracting relevance judgments for training, we use Kimi-K3 as the reasoning model. For retrieval, we start from Ettin-0.4B [8], a small cross-encoder, i.e., a Transformer that jointly encodes each query–document pair and maps them to a score for ranking.

We make this choice so the architectural and learning bottlenecks of bi-encoders [9] do not confound the ARI learning signal we want to study. Cross-encoders are usually considered too slow for retrieval, but we find them quite manageable for long-tail collections, up to roughly 100,000 documents.

With only basic optimizations (bf16 weights and compilation) and by restricting search to the top 10,000 candidates from a cheap first-stage retriever, the trained retriever takes on the order of a couple of seconds per query on one GPU, with essentially unchanged quality (Appendix A). Given that we envision ARI in an LLM agent pipeline, a couple of seconds might already be feasible for some applications. Nonetheless, future work will investigate ARI over sub-linear search, with paradigms like late interaction.

To simplify our experiments, we assume that queries arrive in batches of 20 (e.g., concurrent independent searches), and we give each batch three rounds. In each round, we use the current retriever to score every document against each query and keep its top 100. We then ask Kimi-K3 to rank these 100 listwise [10], in random groups of 20. We combine the rankings into preference pairs with a Bradley–Terry model [11] and take gradient steps on the retriever with a pairwise RankNet loss [12], which rewards it for scoring the preferred document higher.

This operates as a simple form of reinforcement learning with a contrastive reward [13], which updates the retrieval model’s weights at test time [14] to align its rankings with the reasoning model. After the three rounds, we keep the updated weights for the next batch. Alongside each batch, we also revisit earlier queries, some with fresh retrievals and judgments, to spread training across the query history. Appendix B gives the full recipe.

Preliminary Results

We report ARI’s performance on three tasks. These tasks seek: tweets that implicitly express a given stance (Twitter-Conflict), passages by the same author as a query passage (Writing), and math problems solved with the same technique as a query problem (BRIGHT AoPS [15]). For Twitter-Conflict, we report “pooled” judgments, whereby we ask a different LLM (Opus 5.5) to operate as a judge of the top 100 results of ARI and all other reference models. In each experiment, ARI learns from a stream of unlabeled queries, and we score it with relevance labels on held-out queries, on which the retrieval system cannot invoke the reasoning model anymore. We developed ARI on the first two tasks. For AoPS, we adapted only the instructions given to the reasoning model.

We report Qwen3-Embedding-0.6B, Qwen3-Embedding-8B and Gemini-Embedding-2 for reference. We do not consider these as baselines for various reasons; for one, they are larger than Ettin-0.4B and trained on far more data. Moreover, they are bi-encoders with a different speed–quality trade-off, and we do not apply ARI to them. Whether learning during search would improve models like these is interesting but beyond the scope of this post.

Figure 2 reported Recall@100, which shows ARI learning very quickly to bring hard-to-reach relevant documents into its top 100. These gains are broad, not concentrated in a few queries. Across Twitter-Conflict’s held-out queries, the median number of relevant tweets in a query’s top 100 grows nearly fivefold after 100 production queries, three quarters of queries at least double their count, and every query that started with none gains at least one. The top 10 follows suit: the share of queries with no relevant tweet there falls from 63% to 14%.

Takeaway 1: A small retriever can learn to search much better without any relevance labels, just from the judgments a reasoning model makes while answering queries. After 100 unlabeled queries (60 on BRIGHT AoPS), ARI’s Recall@100 on held-out queries grows 3.2× to 7×.

Figure 3 tracks where relevant documents sit in the whole ranking. Their typical rank improves four- to fourteenfold within the first 20 queries. By the end, ARI ranks relevant documents higher than every other retriever on Twitter-Conflict and Writing, Gemini-Embedding-2 included; on BRIGHT AoPS it passes Qwen3-Embedding-0.6B, but not Qwen3-Embedding-8B or Gemini-Embedding-2. Many of the relevant tweets that rise were never shown to the reasoning model for any query, and those it did judge while answering other queries rise further still.

Takeaway 2: ARI is not just memorizing the reasoning model’s past judgments; what it learns carries over to new queries and even to documents the reasoning model never saw. On Twitter-Conflict, after 100 queries, 29% of the relevant tweets in held-out top-10 lists were neither in the starting retriever’s top 10 nor ever judged by the reasoning model, for any query.

Figure 3. ARI moves relevant documents up from deep in the ranking. Typical rank of relevant documents on held-out queries (geometric mean, log scale, lower is better) as ARI processes the query stream. ARI ends below both Qwen embeddings on Twitter-Conflict and Writing. On BRIGHT AoPS, Qwen3-Embedding-8B ends lower (314 vs. 368): ARI puts more relevant problems near the top but leaves more beyond rank 10,000, and a geometric mean weighs those heavily. Per-task NDCG@10 and Recall@10/100 curves appear in Appendix B.

Serving the Reasoning Index

Because ARI’s index is a cross-encoder, nothing about a document is computed ahead of time: answering a query means running the model over every document in the corpus. That is why cross-encoders are usually considered too slow for retrieval. But the collections where oblique search matters most, such as a team’s documents, support tickets or research notes, often hold tens or hundreds of thousands of documents, not billions, and at that scale a full scan is feasible. Standard inference optimizations alone make it three to four times faster without hurting quality (Appendix A). That is fine for batch jobs, such as working through a backlog of queries offline, but far too slow for someone waiting on a search box. Recent work on making cross-encoders efficient enough for end-to-end retrieval [16] suggests this gap will keep shrinking.

Narrowing the search first closes most of the remaining gap (Figure 4). When ARI scores only Qwen3-Embedding-4B’s top 10,000 candidates on Twitter-Conflict, a query takes about two seconds, 28 times faster than the original full scan, with NDCG@10 within 0.01 of it. On Writing, scoring BM25’s top 2,000 is 13 times faster with no loss on either metric. Smaller candidate pools trade quality for speed, down to a few hundred milliseconds per query, and answering every held-out Twitter-Conflict query drops from 2.6 hours with the original scan to under three minutes with 5,000 candidates.

A weak final ranker can still be a good filter. Qwen3-Embedding-4B’s NDCG@10 (0.109 on pooled labels) is less than half of ARI’s (0.273), yet its top 10,000 candidates contain 90% of the relevant tweets that full-scan ARI places in its top 100, which suggests a division of labor in which an embedding model narrows the search and ARI makes the finer relevance distinctions. Appendix A gives the setup and full tables.

Takeaway 3: ARI need not scan the whole corpus. Where a typical reranker only sees a retriever’s top few hundred results, ARI can afford to score a much deeper candidate pool, so the first-stage retriever just has to be decent enough to get relevant documents somewhere in it.

Figure 4. Median latency per query against Recall@100 on held-out queries, on one NVIDIA H200. Circles score only a first stage’s top K candidates (Qwen3-Embedding-4B for Twitter-Conflict, BM25 for Writing), squares are optimized full scans, and hollow circles the original scan; the dashed line marks full-scan recall. On Twitter-Conflict’s pooled labels, filtered recall is a slight underestimate (Appendix A).

Where We Think This Goes: Self-Improving Search and Memory

We’re extremely excited about ARI as a search paradigm for retrieval that gets better with each use. Instead of trying to learn some singular representation upfront and keep throwing compute and parameters at the size of our encoders, ARI asks: what if the retriever’s main job is to scale a reasoning model’s judgment over a large corpus?

As this early investigation suggests, ARI’s answer is that this can be achieved by having the retriever update weights, live at inference time, to align with the reasoning model’s relative ranking preferences. This not only teaches the retriever to adjust its fine-grained rankings, but allows it to find numerous gold nuggets it didn’t even surface to the top 100 before and to improve rapidly on unseen queries from the same distribution.

Importantly, we’re specifically not claiming that ARI as-is builds general retrieval models. In fact, our entire argument is that it doesn’t need to do so. Our most exciting finding is that small retrievers, even with a relatively poor start holding them back, might soon be able to overcome that lack of “generality” not just through bigger models or more training data, but through experience and by discovering where a reasoning model wants to look.

This is, in our view, a better way to spend (and amortize [17]) inference-time compute. And it has implications well beyond search and into agent memory, as agents increasingly have to work over massive collections of documents. Today, due to limited context windows and techniques like compaction and recursion, agents quite literally have to re-discover the corpus they’re in repeatedly over long-horizon work.

But much of memory is knowing where to look [18, 19, 20]. A coding agent rereads the same files and decides again which search results matter, then discards those judgments when the session ends. ARI offers a path for agents like this to steer their retrieval tools towards successful findings much more efficiently, by essentially reinforcing their behavior in good search trajectories.

Our bet is that some of the reasoning that search requires can become reusable infrastructure, so that the effort spent answering today’s questions changes what tomorrow’s searches can find. How much of it can be reused, and for which workloads, we do not yet know. ARI is a first step toward finding out.

Acknowledgments

This work was supported in part by MIT OCP project 6954031 and cmpnd.ai, and used compute credits offered by Sail Research and Mixedbread. Diane also wants to thank Ben Clavié, Madeleine Pollack, and Soyheon Kim for reading an initial draft of this post and providing feedback, and her OASYS labmates for support.

Citation

If you found this blog helpful, please cite it as:

@misc{tchuindjo2026ari,
  title        = {{Amortized Reasoning Index: Scaling the Judgment of Reasoning Models Beyond Their Context Window}},
  author       = {Tchuindjo, Diane and Shah, Devavrat and Khattab, Omar},
  year         = {2026},
  month        = nov,
  howpublished = {\url{https://dianetc.com/musings/ari/}},
  note         = {Blog post}
}