Introduction

Large Language Models (LLMs) are increasingly often expected to handle complex multi-step reasoning directly on the user's device. Yet model's capabilities are largely fixed when it is deployed. When it encounters unfamiliar types of problems or a shift in domain, its reasoning can become unreliable, and on a smartphone we cannot assume access to labelled data, large auxiliary models, or cloud resources to resolve this.

Test-Time Adaptation (TTA) addresses this by letting a model adjust itself at inference time. A particularly promising example is Test-Time Reinforcement Learning (TTRL) [1], which improves a model on unlabelled questions by rewarding it for agreeing with its own majority vote. TTRL can deliver large gains, but its learning signal is only as good as that vote: when the majority is wrong, the model is rewarded for being wrong.

In-Context Learning (ICL) [2], on the other hand, has shown that a few well-chosen worked examples in the prompt can substantially improve reasoning without any weight updates. In our paper, we bring these two ideas together and introduce Context-Guided TTRL (CG-TTRL), the first framework to integrate ICL into an unsupervised test-time fine-tuning pipeline, along with a lightweight context-selection method designed for on-device deployment. CG-TTRL outperforms TTRL and, importantly for devices, reaches strong performance after only a few update steps.

The Challenge: Learning Without Labels

TTRL starts from a user query $q$ and a model $\pi_\theta$ with parameters $\theta$, and proceeds in two phases.

In the exploitation phase, the model generates $M$ predictions $ \{ y ̂_i \}_{(i=1)}^M $ and derives a pseudo-label by majority voting:

$ y \leftarrow mode( \{ y ̂_i \}_{(i=1)}^M ).$

In the exploration phase, $ N < M $ of the generated samples, $ \{ \widehat{y_i^{'}} \}_{(i=1)}^N $, are selected and scored against the pseudo-label with a binary reward:

and $\theta$ is updated via gradient ascent to maximize $\mathbb{E}\left[R(\hat{y}'_i, y)\right] $. The model thus learns to explore diverse reasoning paths that lead to its consensus answer, without any ground truth.

The weakness is that everything relies on $y$. If most samples converge on an incorrect answer, the reward actively reinforces flawed reasoning and penalizes the correct solutions. This is especially likely for small on-device models, whose individual samples are less reliable. Simply adding examples to the prompt is not a complete answer either: poorly chosen examples can hurt performance, and sophisticated selection methods typically require auxiliary models or expensive computation that devices cannot afford.

Our Solution: CG-TTRL

Our key observation is that TTRL's two-phase strategy under-utilizes contextual guidance, which can help in both phases. CG-TTRL augments the query with $C$ retrieved examples ${q^c_i, s^c_i, y^c_i}_{i=1}^{C}$, consisting of example queries $q^c_i$, step-by-step solutions $s^c_i$, and final responses $y^c_i$:

$ q' = \{ q^c_i, s^c_i, y^c_i \}_{i=1}^{C} \cup q,$

and then runs both phases on $q'$ instead of $q$.

Efficient context selection. To find the right examples cheaply, we represent each query with TF-IDF features [3], which highlight terms that are frequent in a given query but rare across the context pool $E_{all}$. These features require only simple counting operations and no additional model on the device. Given a new query $q$, we compute its similarity to all $K$ queries ${q^p_i}_{i=1}^{K}$ in the pool,

$$\text{sim}(q, q^p_i) = \cos\left(f_{TF\text{-}IDF}(q),\ f_{TF\text{-}IDF}(q^p_i)\right),$$

and select the $C$ most similar ones, together with their solutions and responses, as context.

Context-guided exploitation. The $M$ predictions are now drawn from $\pi_\theta(\cdot \mid q')$. Relevant worked examples raise the chance that each sample follows correct reasoning, and therefore the chance that the correct answer wins the majority vote. Since every reward is computed against $y$, a more accurate pseudo-label directly yields a cleaner learning signal.

Context-guided exploration. The $N$ samples used for fine-tuning are also context-augmented, so the policy updates anchor exploration to the reasoning patterns shown in the examples. The same context is reused at final inference, so it benefits every stage of the pipeline. This also distinguishes CG-TTRL from plain ICL, which uses context only for a single inference and does not benefit from reward-driven adaptation.

Experimental Results

We perform test-time training on MATH-500 or AMC and evaluate on five reasoning benchmarks: MATH-500, GSM-100, AMC, AIME 2024, and GPQA. We focus on models suitable for on-device use with 1.5B parameters: Qwen2.5-Math-1.5B-Instruct, Qwen2.5-1.5B-Instruct, and DeepSeek-R1-Distill (1.5B).

CG-TTRL consistently outperforms TTRL. Across all three models and both training datasets, CG-TTRL achieves the best average accuracy. For example, DeepSeek-R1-Distill improves from 37.2% without adaptation to 56.8% with TTRL, and further to 59.4% with CG-TTRL.

Robust under domain shift. The benchmarks span a range of shifts relative to the training data, from mild (AMC, AIME, GSM-100) to substantial (GPQA, which contains graduate-level questions from multiple scientific domains, rather than mathematics only). CG-TTRL is beneficial across this whole spectrum, and the largest relative gains appear on the hardest benchmarks. For Qwen2.5-Math-1.5B-Instruct trained on MATH-500, accuracy on GPQA improves from 27.3% to 38.4% (+40.7% relative to the base model), and on AIME from 10.0% to 23.3% (+133.3%), compared with 37.4% and 20.0% for TTRL.

Strong gains after only a few steps. On a device, running TTRL to convergence is rarely practical, so we studied adaptation with only 1–5 epochs, where one epoch corresponds to three update steps. After a single epoch, CG-TTRL already raises average accuracy from 48.9% to 53.0% (about +8% relative), while TTRL reaches only 49.5% (about +1%). After five epochs, CG-TTRL reaches 54.7% versus 51.3% for TTRL. This fast adaptation makes CG-TTRL especially suitable for on-device use.

Conclusion

Test-time reinforcement learning offers a compelling path towards models that keep improving after deployment, but its self-generated rewards are only as reliable as its majority vote. CG-TTRL shows that a small amount of well-chosen context, retrieved with lightweight TF-IDF features, makes these votes more accurate and guides exploration more effectively. The result is higher accuracy across in-domain and out-of-domain benchmarks and much faster improvement when only a few adaptation steps are possible, making CG-TTRL a practical building block for continuous, context-aware self-improvement of on-device LLMs.