Além do Treinamento no Momento do Teste: Aprendendo a Raciocinar por Meio de Controle Ótimo com Eficiência em 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
Resumo
A memória associativa tem sustentado há muito tempo o projeto de modelos sequenciais. Para além da recordação, os seres humanos raciocinam projetando estados futuros e selecionando ações orientadas a objetivos, uma capacidade que os modelos de linguagem modernos exigem cada vez mais, mas não codificam nativamente. Embora trabalhos anteriores utilizem aprendizagem por reforço ou treino no momento do teste, o planeamento permanece externo à arquitetura do modelo. Nós formulamos o raciocínio como controlo ótimo e introduzimos a camada de Controlo no Momento do Teste (TTC), que executa planeamento LQR de horizonte finito sobre estados latentes durante a inferência, representa uma função de valor dentro de arquiteturas neurais e a utiliza como objetivo aninhado para permitir o planeamento antes da previsão. Para garantir escalabilidade, derivamos um resolvedor LQR eficiente em hardware baseado numa formulação simplética e implementamo-lo como um kernel CUDA fundido, permitindo execução paralela com sobrecarga mínima. Integradas como adaptador em LLMs pré-treinados, as camadas TTC melhoram o desempenho do raciocínio matemático em até +27,8% no MATH-500 e melhorias de 2-3x em Pass@8 no AMC e AIME, demonstrando que incorporar o controlo ótimo como um componente arquitetónico fornece um mecanismo eficaz e escalável para raciocínio além do treino no momento do teste.
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.