Learning to Stop Overthinking at Test Time

TL;DR

This work tackles the problem of wasting computation and the phenomenon of overthinking in test-time reasoning by enabling adaptive per-sample iteration through test-time training. It combines a latent-space visual reasoning framework with a self-supervised rotation task to estimate the optimal number of thinking steps, enabling accurate predictions with minimal compute. The authors introduce Conv-LiGRU, a lightweight, stable recurrent architecture that removes the reset gate and uses ReLU with batch normalization for efficient iterative processing, showing superior robustness to corruption on CIFAR10-C and CIFAR100-C. Overall, the approach demonstrates strong extrapolation capabilities, parameter efficiency, and practical potential for robust visual reasoning under distribution shifts.

Abstract

Test time scaling is currently one of the most active research areas that shows promise after training time scaling has reached its limits. Deep-thinking (DT) models are a class of recurrent models that can perform easy-to-hard generalization by assigning more compute to harder test samples. However, due to their inability to determine the complexity of a test sample, DT models have to use a large amount of computation for both easy and hard test samples. Excessive test time computation is wasteful and can cause the ``overthinking'' problem where more test time computation leads to worse results. In this paper, we introduce a test time training method for determining the optimal amount of computation needed for each sample during test time. We also propose Conv-LiGRU, a novel recurrent architecture for efficient and robust visual reasoning. Extensive experiments demonstrate that Conv-LiGRU is more stable than DT, effectively mitigates the ``overthinking'' phenomenon, and achieves superior accuracy.

Figures (11)

• Figure 1: Deep thinking pipeline with training includes both main (classification) and secondary (self-supervised) tasks. During inference, the iteration with peak self-supervised performance is selected for the halting of the classification task. The reset gate in Conv-GRU is simplified and $\phi$ is replaced from Tanh to ReLu, resulting in Conv-LiGRU.
• Figure 2: Accuracy (%) on CIFAR100-C, at level 5, Resnet, Conv-GRU, Conv-LiGRU
• Figure 3: We evaluate the ability of networks to classify objects on two test sets with higher noise severity than those used during training, measuring accuracy on level 3 inputs (purple) and level 5 inputs (red). Recurrent models are compared against the best feed-forward models.
• Figure 4: Accuracy (%) on CIFAR10-C, at level 5, Conv-LiGRU with and without ACT
• Figure 5: A "Cat" sample input, and outputs from different iterations are shown to illustrate the model’s sequential reasoning process on CIFAR10-C (level 1). We visualize the norm of vector feature $h^{(i, j)}_t$ (row $i$, column $j$) of the feature map $h_t$, demonstrating the model’s feature extraction over iterations. This is a representative example from a Conv-LiGRU model trained on CIFAR10 with $T_{\text{train}} = 30$.