Oltre l'Addestramento al Momento del Test: Imparare a Ragionare Tramite Controllo Ottimo Efficiente a Livello Hardware
Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
March 10, 2026
Autori: Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, Zhangyang Wang, Ming Lin, René Vidal
cs.AI
Abstract
La memoria associativa ha a lungo costituito la base per la progettazione di modelli sequenziali. Oltre al richiamo, gli esseri umani ragionano proiettando stati futuri e selezionando azioni orientate a uno scopo, una capacità che i moderni modelli linguistici richiedono sempre più ma non codificano nativamente. Mentre i lavori precedenti utilizzano l'apprendimento per rinforzo o l'addestramento al momento del test, la pianificazione rimane esterna all'architettura del modello. Noi formuliamo il ragionamento come un problema di controllo ottimo e introduciamo il layer Test-Time Control (TTC), che esegue una pianificazione LQR a orizzonte finito sugli stati latenti al momento dell'inferenza, rappresenta una funzione di valore all'interno delle architetture neurali e la utilizza come obiettivo annidato per abilitare la pianificazione prima della previsione. Per garantire la scalabilità, deriviamo un risolutore LQR efficiente in termini hardware basato su una formulazione simplettica e lo implementiamo come un kernel CUDA fuso, consentendo un'esecuzione parallela con overhead minimo. Integrati come adattatori in LLM preaddestrati, i layer TTC migliorano le prestazioni nel ragionamento matematico fino al +27,8% su MATH-500 e incrementi di 2-3x in Pass@8 su AMC e AIME, dimostrando che l'incorporamento del controllo ottimo come componente architetturale fornisce un meccanismo efficace e scalabile per il ragionamento che va oltre l'addestramento al momento del 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.