RTTC: Reward-Guided Collaborative Test-Time Compute

TL;DR

RTTC introduces an adaptive framework that intelligently selects the optimal test-time compute strategy for large language models, significantly improving accuracy while reducing unnecessary computation across diverse tasks.

Contribution

This work presents RTTC, a reward-guided, adaptive approach for selecting test-time compute strategies, incorporating distributed retrieval, lightweight fine-tuning, and query-state caching.

Findings

RTTC outperforms vanilla RAG and TTT in accuracy across multiple benchmarks.

Adaptive strategy selection reduces computational overhead.

Query-State Caching improves efficiency by reusing historical query states.

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.