Au-delà de l'entraînement au moment du test : Apprendre à raisonner via un contrôle optimal efficace en matériel
Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
March 10, 2026
Auteurs: Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, Zhangyang Wang, Ming Lin, René Vidal
cs.AI
Résumé
La mémoire associative a longtemps soutenu la conception des modèles séquentiels. Au-delà du rappel, les humains raisonnent en projetant des états futurs et en sélectionnant des actions orientées vers un but, une capacité que les modèles de langage modernes nécessitent de plus en plus mais n'encodent pas nativement. Si les travaux antérieurs utilisent l'apprentissage par renforcement ou l'entraînement au moment du test, la planification reste externe à l'architecture du modèle. Nous formulons le raisonnement comme un contrôle optimal et introduisons la couche de contrôle au moment du test (TTC), qui effectue une planification LQR à horizon fini sur les états latents lors de l'inférence, représente une fonction de valeur dans les architectures neuronales et l'utilise comme objectif imbriqué pour permettre la planification avant la prédiction. Pour garantir l'évolutivité, nous dérivons un solveur LQR efficace matériellement basé sur une formulation symplectique et l'implémentons comme un noyau CUDA fusionné, permettant une exécution parallèle avec une surcharge minimale. Intégrées comme adaptateur dans les LLM pré-entraînés, les couches TTC améliorent les performances en raisonnement mathématique jusqu'à +27,8 % sur MATH-500 et des améliorations de 2 à 3x du Pass@8 sur AMC et AIME, démontrant qu'intégrer le contrôle optimal comme composant architectural fournit un mécanisme efficace et évolutif pour le raisonnement au-delà de l'entraînement au moment du test.
English
Associative memory has long underpinned the design of sequential models. Beyond recall, humans reason by projecting future states and selecting goal-directed actions, a capability that modern language models increasingly require but do not natively encode. While prior work uses reinforcement learning or test-time training, planning remains external to the model architecture. We formulate reasoning as optimal control and introduce the Test-Time Control (TTC) layer, which performs finite-horizon LQR planning over latent states at inference time, represents a value function within neural architectures, and leverages it as the nested objective to enable planning before prediction. To ensure scalability, we derive a hardware-efficient LQR solver based on a symplectic formulation and implement it as a fused CUDA kernel, enabling parallel execution with minimal overhead. Integrated as an adapter into pretrained LLMs, TTC layers improve mathematical reasoning performance by up to +27.8% on MATH-500 and 2-3x Pass@8 improvements on AMC and AIME, demonstrating that embedding optimal control as an architectural component provides an effective and scalable mechanism for reasoning beyond test-time training.