За пределами обучения во время тестирования: обучение логическому мышлению с помощью энергоэффективного оптимального управления

Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control

March 10, 2026

Авторы: Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, Zhangyang Wang, Ming Lin, René Vidal

cs.AI

Аннотация

Ассоциативная память долгое время лежала в основе проектирования последовательных моделей. Однако помимо воспроизведения, люди рассуждают, проецируя будущие состояния и выбирая целенаправленные действия — способность, которая все более необходима современным языковым моделям, но не заложена в них изначально. В то время как предыдущие работы используют обучение с подкреплением или обучение во время тестирования, планирование остается внешним по отношению к архитектуре модели. Мы формулируем рассуждение как задачу оптимального управления и представляем слой управления во время тестирования (Test-Time Control, TTC), который выполняет планирование по методу LQR с конечным горизонтом над латентными состояниями во время вывода, представляет функцию ценности внутри нейросетевых архитектур и использует ее как вложенную цель, чтобы обеспечить планирование до предсказания. Для обеспечения масштабируемости мы выводим аппаратно-эффективный решатель LQR на основе симплектической формулировки и реализуем его в виде слитого CUDA-ядра, что позволяет параллельное выполнение с минимальными накладными расходами. Интегрированные в качестве адаптера в предварительно обученные большие языковые модели, слои TTC улучшают результаты математических рассуждений до +27.8% на MATH-500 и в 2-3 раза повышают показатель Pass@8 на AMC и AIME, демонстрируя, что встраивание оптимального управления в качестве архитектурного компонента предоставляет эффективный и масштабируемый механизм для рассуждений, превосходящий обучение во время тестирования.

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.