Voorbij training tijdens testtijd: Leren redeneren via hardware-efficiënte optimale regeling
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
Samenvatting
Associatief geheugen heeft lange tijd de basis gevormd voor het ontwerp van sequentiële modellen. Naast het oproepen van informatie, redeneren mensen door toekomstige toestanden te projecteren en doelgerichte acties te selecteren, een capaciteit die moderne taalmodelen steeds meer nodig hebben maar niet intrinsiek bezitten. Terwijl eerder werk reinforcement learning of training tijdens testtijd gebruikt, blijft plannen extern aan de modelarchitectuur. Wij formuleren redeneren als optimale controle en introduceren de Test-Time Control (TTC) laag, die tijdens inferentie eindig-horizon LQR-planning uitvoert over latente toestanden, een waardefunctie binnen neurale architecturen representeert, en deze gebruikt als een genest doel om planning vóór voorspelling mogelijk te maken. Om schaalbaarheid te waarborgen, leiden we een hardware-efficiënte LQR-oplosser af, gebaseerd op een symplectische formulering, en implementeren deze als een gefuseerde CUDA-kernel, waardoor parallelle uitvoering met minimale overhead mogelijk is. Geïntegreerd als adapter in voorgetrainde LLM's, verbeteren TTC-lagen de wiskundige redeneerprestaties met tot +27.8% op MATH-500 en 2-3x Pass@8 verbeteringen op AMC en AIME. Dit toont aan dat het inbedden van optimale controle als een architecturale component een effectief en schaalbaar mechanisme biedt voor redeneren dat verder gaat dan training tijdens testtijd.
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.