TLDR: RTTC (Reward-Guided Test-Time Compute) is a new framework that uses a reward model to adaptively select the best strategy (no adaptation, RAG, or TTT) for each LLM query at inference time, significantly improving accuracy and efficiency compared to applying single methods. It also features Query-State Caching to reuse past computations, further reducing overhead.
Large Language Models (LLMs) have become incredibly powerful, but making them robust and adaptable to new information or different scenarios during their use (known as inference time) remains a challenge. Two popular methods to boost LLM performance at this stage are Retrieval-Augmented Generation (RAG) and Test-Time Training (TTT).
RAG works by giving the LLM additional, relevant information retrieved from a knowledge base alongside the original query. This helps the model access up-to-date or specific domain knowledge it might not have been trained on. TTT, on the other hand, involves briefly fine-tuning the LLM’s parameters using relevant examples at the time of inference, helping it adapt to new data distributions.
While both RAG and TTT are effective, they have drawbacks. Their effectiveness can vary greatly depending on the specific question asked. Sometimes RAG is better, sometimes TTT, and sometimes the model’s initial answer is already good enough. Applying these methods indiscriminately to every query can also be very costly in terms of computation and time, leading to unnecessary delays and resource use.
To address these issues, researchers have introduced a new framework called Reward-Guided Test-Time Compute (RTTC). This innovative system intelligently decides the best way to handle each query, choosing between no adaptation, RAG, or TTT. At its core, RTTC uses a special “reward model” that evaluates potential answers and guides the system to pick the most effective strategy for each specific question. This adaptive approach helps maximize accuracy across various topics and tasks.
RTTC operates using a distributed server-client setup. A remote server holds a vast knowledge base, handling the retrieval of relevant information. Meanwhile, the actual processing, reward evaluation, and model adaptation happen on the client devices. This design helps protect user privacy by keeping sensitive inference data local and reduces the need for clients to store large knowledge bases.
The workflow of RTTC is quite clever. When a query comes in, the LLM first generates an initial response. A pretrained reward model then assesses the quality of this response. If the initial response is good enough (exceeds a certain quality threshold), it’s returned immediately, saving computational resources. If not, the system moves to retrieve relevant information from the knowledge base. With this retrieved data, it tries RAG. The reward model evaluates the RAG-augmented response. If it’s better than the initial one, that response is used. If still not satisfactory, RTTC then performs a lightweight Test-Time Training (TTT) using the retrieved samples to adapt the model, generating a final response.
An optional “joint” strategy also exists where both RAG and TTT are run in parallel for challenging queries, and the reward model simply picks the best outcome. This can further boost performance but comes with higher computational costs.
To further enhance efficiency, RTTC incorporates a feature called Query-State Caching (QSC). This mechanism stores historical query information, including retrieved samples for RAG and fine-tuned model states for TTT. If a new query is similar enough to a past one, QSC allows the system to reuse previously computed results, significantly reducing redundant retrieval and fine-tuning operations, thereby speeding up the process and lowering latency.
Extensive experiments have shown that RTTC consistently achieves higher accuracy compared to using RAG or TTT alone across various LLMs (like Llama-3, Mistral, and Qwen) and different tasks in coding, math, and medical domains. The results validate the importance of adaptively selecting the right test-time compute strategy, guided by a reward model. QSC also proved effective in maintaining performance while significantly improving efficiency.
Also Read:
- Boosting LLM Search Efficiency Through Parallel Query Decomposition
- Nested-ReFT: Boosting LLM Fine-Tuning Efficiency with Layer Skipping
While RTTC offers compelling advancements, the researchers acknowledge some limitations. These include the need for manual tuning of certain parameters, the importance of a well-curated and extensive knowledge base for optimal performance, and ongoing privacy considerations for user data on the server. Despite these, RTTC represents a significant step towards more scalable and high-performing language model adaptation. You can read the full research paper here.
