RTTC : Reward-Guided Collaborative Test-Time Compute (2025)

Abstract: Test-Time Compute (TTC) has emerged as a powerful paradigm for enhancing the performance of Large Language Models (LLMs) at inference, leveraging strategies such as Test-Time Training (TTT) and Retrieval-Augmented Generation (RAG). However, the optimal adaptation strategy varies across queries, and indiscriminate application of TTC strategy incurs substantial computational overhead. In this work, we introduce Reward-Guided Test-Time Compute (RTTC), a novel framework that adaptively selects the most effective TTC strategy for each query via a pretrained reward model, maximizing downstream accuracy across diverse domains and tasks. RTTC operates in a distributed server-client architecture, retrieving relevant samples from a remote knowledge base and applying RAG or lightweight fine-tuning on client devices only when necessary. To further mitigate redundant computation, we propose Query-State Caching, which enables the efficient reuse of historical query states at both retrieval and adaptation levels. Extensive experiments across multiple LLMs and benchmarks demonstrate that RTTC consistently achieves superior accuracy compared to vanilla RAG or TTT, validating the necessity of adaptive, reward-guided TTC selection and the potential of RTTC for scalable, high-performance language model adaptation.

## introduction large language models have achieved remarkable performance across a wide range of tasks. however, their robustness and adaptability to new domains or distribution shifts at inference time remain open challenges. traditionally, two major paradigms have been explored to enhance llm performance at test time: retrieval-augmented generation (rag) @xcite , which augments model input with retrieved knowledge, and test-time training (ttt) @xcite ttc strategy latency memory accuracy no adaptation (direct inference) + + + retrieval-augmented generation (rag) ++ ++ +++ test-time training (ttt) +++ +++ +++ aiming to tackle the above challenges, we propose reward-guided test-time computing (rttc). this framework dynamically selects among three strategies for each query: no adaptation (i.e., returning the model's response without adaptation), retrieval-augmented generation (rag), and test-time training (ttt). at the core of rttc is a pretrained reward model that evaluates candidate responses and guides the system to choose the most effective adaptation strategy in a query-adaptive manner (see figure 2 ). this reward-guided collaboration enables rttc to adaptively exploit the most suitable strategy for each query, achieving robust downstream performance improvements across diverse domains and tasks. unlike prior work that statically applies either rag or ttt, our approach introduces a principled decision-making mechanism that maximizes performance. additionally, rttc also introduces the query-state caching (qsc) mechanism to further optimize the test-time efficiency. qsc leverages historical query embeddings and their associated retrieved samples or fine-tuned model state (e.g., @xcite ) to potentially bypass the need for repeated retrieval and fine-tuning, thus reducing computational overhead and latency. overall, our main contributions are: in summary, our work is the first to unify direct inference, rag, and ttt within a reward-guided, query-adaptive framework. the remainder of the paper is organized as follows. we discuss related work in §2. then, §3 describes the rttc system, while §4 presents results on various downstream tasks. our final thoughts are in §5. ## related work test-time compute. test-time compute techniques have been proposed as an alternative to scaling model parameters for improving model performance @xcite . notable strategies include chain-of-thought (cot) prompting @xcite and few-shot learning @xcite . cot prompting guides the model through intermediate steps to break down complex tasks, enhancing its ability to handle intricate queries. few-shot learning helps the model adapt to new tasks with just a few examples. other test-time compute methods include verifying the model's results, for example, through code execution @xcite . retrieval-augmented generation (rag). rag has emerged as a prominent paradigm for enhancing model performance by incorporating external knowledge at inference time @xcite . by retrieving relevant documents and augmenting the model's input, rag enables llms to access up-to-date or domain-specific information beyond their pretraining corpus. this approach has demonstrated strong results across open-domain question answering, fact verification, and knowledge-intensive tasks. rag introduces additional computational overhead due to retrieval and longer input sequences, and its effectiveness can vary depending on the quality of retrieved content and the task. test-time training (ttt). ttt @xcite @xcite has proven effective for adapting models to distribution shifts by fine-tuning on retrieved samples during inference. recent work, such as sift @xcite , has improved retrieval strategies for ttt, while omni-arc (iron-bararc24, 2024) leveraged ttt to achieve stateof-the-art results in the arc-agi challenge (arc-agi, 2025; chollet, 2019). despite these advances, ttt can incur significant computational and memory costs. ## rttc system this section outlines the architecture and workflow of reward-guided collaborative test-time compute (rttc), a system designed to optimize the performance of large language models (llms). rttc unifies direct inference, retrieval-augmented generation (rag), and test-time training (ttt) within a reward-driven, query-adaptive framework, leveraging a remote multi-domain knowledge base for robust adaptation. the overall workflow is shown in figure 2 and formally summarized in algorithm 1. ## experiments we implement a prototype of the rttc system as a testbed, demonstrating the potential benefits in a larger deployment. next, we discuss the resources utilized in our experimentation, followed by results demonstrating the benefits of enabling rttc. ## conclusion test-time compute (ttc) is an effective paradigm for enhancing model performance at the expense of increased computation during inference. we present rttc, a system that adaptively selects the optimal ttc strategy for each query at test-time. the results of the rttc prototype, which utilizes supervisory signals from a knowledge base, serve as a call to action to investigate further improvements to learning at test-time methods. rttc is guided by a reward model that assists in the decision of the ttc method to apply at inference time. importantly, the reward-guided approach in rttc is theoretically extensible to any ttc strategy, providing an open and generalizable direction for future research in adaptive test-time compute. the current version of rttc works on text queries. with increased sophistication, future versions of rttc must handle more complex tasks and multimodal queries.