Jenseits des Test-Time Trainings: Lernen zu Schlussfolgern durch hardwareeffiziente optimale Steuerung
Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
March 10, 2026
Autoren: Peihao Wang, Shan Yang, Xijun Wang, Tesi Xiao, Xin Liu, Changlong Yu, Yu Lou, Pan Li, Zhangyang Wang, Ming Lin, René Vidal
cs.AI
Zusammenfassung
Assoziatives Gedächtnis bildet seit langem die Grundlage für den Entwurf sequenzieller Modelle. Über reine Erinnerung hinaus schließen Menschen durch die Projektion zukünftiger Zustände und die Auswahl zielgerichteter Aktionen – eine Fähigkeit, die moderne Sprachmodelle zunehmend benötigen, aber nicht von Haus aus besitzen. Während frühere Arbeiten Verstärkendes Lernen oder Test-Time-Training nutzen, bleibt die Planung extern zur Modellarchitektur. Wir formulieren Schlussfolgern als optimale Steuerung und führen die Test-Time-Control (TTC)-Schicht ein, die zur Inferenzzeit eine Finite-Horizon-LQR-Planung über latente Zustände durchführt, eine Wertfunktion innerhalb neuronaler Architekturen repräsentiert und diese als verschachteltes Ziel nutzt, um Planung vor der Vorhersage zu ermöglichen. Um Skalierbarkeit zu gewährleisten, leiten wir einen hardwareeffizienten LQR-Löser basierend auf einer symplektischen Formulierung ab und implementieren ihn als fused CUDA-Kernel, was parallele Ausführung mit minimalem Overhead ermöglicht. Als Adapter in vortrainierte LLMs integriert, verbessern TTC-Schichten die mathematische Reasoning-Leistung um bis zu +27,8 % auf MATH-500 und 2-3-fache Pass@8-Verbesserungen auf AMC und AIME. Dies zeigt, dass die Einbettung optimaler Steuerung als architektonische Komponente einen effektiven und skalierbaren Mechanismus für Schlussfolgern jenseits von Test-Time-Training bereitstellt.
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.