Más allá del entrenamiento en tiempo de prueba: Aprendizaje del razonamiento mediante control óptimo eficiente en hardware
Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
March 10, 2026
Autores: Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, Zhangyang Wang, Ming Lin, René Vidal
cs.AI
Resumen
La memoria asociativa ha sido durante mucho tiempo la base del diseño de modelos secuenciales. Más allá de la recuperación, los humanos razonan proyectando estados futuros y seleccionando acciones dirigidas a objetivos, una capacidad que los modelos de lenguaje modernos requieren cada vez más pero que no codifican de forma nativa. Si bien trabajos previos utilizan aprendizaje por refuerzo o entrenamiento en tiempo de prueba, la planificación sigue siendo externa a la arquitectura del modelo. Formulamos el razonamiento como un problema de control óptimo e introducimos la capa de Control en Tiempo de Prueba (TTC), que realiza una planificación LQR de horizonte finito sobre estados latentes durante la inferencia, representa una función de valor dentro de arquitecturas neuronales y la aprovecha como objetivo anidado para permitir la planificación antes de la predicción. Para garantizar escalabilidad, derivamos un solucionador LQR eficiente en hardware basado en una formulación simpléctica y lo implementamos como un núcleo CUDA fusionado, permitiendo ejecución paralela con sobrecarga mínima. Integradas como adaptadores en LLMs preentrenados, las capas TTC mejoran el rendimiento en razonamiento matemático hasta en un +27.8% en MATH-500 y mejoras de 2-3x en Pass@8 en AMC y AIME, demostrando que la incorporación del control óptimo como componente arquitectónico proporciona un mecanismo efectivo y escalable para el razonamiento que va más allá del entrenamiento en tiempo de prueba.
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.