Model selects easier problems to speed up test-time learning
SelectTTRL: Representation-Guided Data Selection for Efficient Test-Time Reinforcement Learning
Machine Learning
Summary
Some language models learn to solve new problems by trying many answers and choosing the best one, but this can be slow for all problems. The authors found a way to guess which problems will benefit from extra learning by measuring how the model’s initial answers relate to known correct patterns, without needing the correct answers. Using this guess, their method saves time by only doing extra learning on problems likely to improve, keeping overall accuracy high. This approach works across different models and question types, cutting down on expensive computations.
What this means in practice
- •For machine learning engineers: Cut down the time and computation needed for models to adapt on new test problems by selectively applying reinforcement learning only on promising cases.
- •For ai system operators: Improve efficiency of deployed language models by avoiding costly on-the-fly training on tasks unlikely to benefit from it.
Authors
Yi Li, Linan Yue, Weibo Gao
Abstract
Test-Time Reinforcement Learning (TTRL) enables large language models to adapt directly on unlabeled test problems by constructing pseudo-labels from multiple sampled responses. However, existing TTRL methods typically apply the same expensive sampling and optimization procedure to all test problems, despite substantial differences in their training utility. In this paper, we study whether such problems can be identified before costly test-time training, without access to ground-truth answers. We uncover a stable relationship between internal response representations and problem-level empirical accuracy. Specifically, we define COS as the cosine similarity between a model's response representation and a correctness-related direction constructed from an independent question-answering dataset. COS consistently correlates with problem solvability across different models, datasets, and experimental settings, and we further provide a theoretical analysis explaining this relationship. Building on this finding, we propose SelectTTRL, a representation-guided data selection method that ranks test problems using COS and applies TTRL only to a selected percentile interval. SelectTTRL leaves the standard TTRL pipeline unchanged for the selected problems, while avoiding unnecessary multi-response sampling and reinforcement learning updates on the remaining ones. Experiments across multiple reasoning benchmarks demonstrate that SelectTTRL substantially reduces test-time computation while maintaining or improving the performance of standard TTRL. Code is available at https://github.com/heyiliyi/SelectTTRL-public.