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

• URL: http://arxiv.org/abs/2603.09221v1
• Date: Tue, 10 Mar 2026 05:42:13 GMT
• Title: Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control
• Abstract summary: We formulate reasoning as optimal control and introduce the Test-Time Control layer, which performs finite-horizon LQR planning over latent states at inference time.<n>We demonstrate that embedding optimal control as an architectural layer provides an effective and scalable mechanism for reasoning beyond test-time training.
• Score: 86.63490309209378
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• ThinkBooster: A Unified Framework for Seamless Test-Time Scaling of LLM Reasoning [110.46431027868954]
  ThinkBooster is a framework for seamless test-time compute scaling of large language model (LLM) reasoning.<n>It consists of a modular Python library implementing state-of-the-art TTC scaling strategy and scorer families.<n>A deployable OpenAI-compatible proxy service enables drop-in integration of adaptive reasoning into real-world applications.
  arXiv  Detail & Related papers  (2026-06-05T05:28:46Z)
• Planner Matters! An Efficient and Unbalanced Multi-agent Collaboration Framework for Long-horizon Planning [32.56971325545821]
  Language model (LM)-based agents have demonstrated promising capabilities in automating complex tasks from natural language instructions.<n>We propose an enhanced multi-agent framework that decomposes automation into three roles: a planner for high-level decision-making, an actor for task execution, and a memory manager for contextual reasoning.
  arXiv  Detail & Related papers  (2026-05-04T02:58:05Z)
• How LLMs Fail and Generalize in RTL Coding for Hardware Design? [56.361436215029045]
  We introduce a new error taxonomy grounded in problem solvability, inspired by cognitive theory.<n>Our taxonomy categorizes failures into syntactic, semantic, solvable functional, and unsolvable functional types.
  arXiv  Detail & Related papers  (2026-04-26T14:34:49Z)
• EdgeReasoning: Characterizing Reasoning LLM Deployment on Edge GPUs [0.36050743818632486]
  Large language models (LLMs) for reasoning tasks on edge GPU face critical challenges from strict latency constraints and limited computational resources.<n>To navigate these constraints, developers must balance reasoning versus non-reasoning architectures, selecting appropriate model sizes, allocating token budgets, and applying test-time scaling strategies.<n>We present EdgeReasoning, a comprehensive study characterizing the deployment of reasoning LLMs on edge GPUs.
  arXiv  Detail & Related papers  (2025-10-21T04:18:25Z)
• Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts [19.518525241726916]
  Encode-Think-Decode (ETD) is a method that enhances the reasoning capabilities of a base model by training it to iterate over a small subset of reasoning-relevant layers during the mid-training stage.<n>ETD models yield substantial gains on 17 reasoning benchmarks, including +28.4% relative accuracy improvement on GSM8K and +36% on MATH with the OLMo-2 1B Base model.
  arXiv  Detail & Related papers  (2025-10-08T15:58:35Z)
• MesaNet: Sequence Modeling by Locally Optimal Test-Time Training [67.45211108321203]
  We introduce a numerically stable, chunkwise parallelizable version of the recently proposed Mesa layer.<n>We show that optimal test-time training enables reaching lower language modeling perplexity and higher downstream benchmark performance than previous RNNs.
  arXiv  Detail & Related papers  (2025-06-05T16:50:23Z)
• PLM: Efficient Peripheral Language Models Hardware-Co-Designed for Ubiquitous Computing [48.30406812516552]
  We introduce the PLM, a Peripheral Language Model, developed through a co-design process that jointly optimize model architecture and edge system constraints.<n>PLM employs a Multi-head Latent Attention mechanism and employs the squared ReLU activation function to encourage sparsity, thereby reducing peak memory footprint.<n> evaluation results demonstrate that PLM outperforms existing small language models trained on publicly available data.
  arXiv  Detail & Related papers  (2025-03-15T15:11:17Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• To Code or not to Code? Adaptive Tool Integration for Math Language Models via Expectation-Maximization [30.057052324461534]
  We propose a novel framework that synergizes structured exploration (E-step) with off-policy optimization (M-step) to create a self-reinforcing cycle between metacognitive tool-use decisions and evolving capabilities.<n>Our 7B model improves over 11% on MATH500 and 9.4% on AIME without o1-like CoT.
  arXiv  Detail & Related papers  (2025-02-02T06:32:23Z)
• Mechanistic Design and Scaling of Hybrid Architectures [114.3129802943915]
  We identify and test new hybrid architectures constructed from a variety of computational primitives. We experimentally validate the resulting architectures via an extensive compute-optimal and a new state-optimal scaling law analysis. We find MAD synthetics to correlate with compute-optimal perplexity, enabling accurate evaluation of new architectures.
  arXiv  Detail & Related papers  (2024-03-26T16:33:12Z)
• PerfRL: A Small Language Model Framework for Efficient Code Optimization [14.18092813639534]
  In this paper, we introduce PerfRL, an innovative framework designed to tackle the problem of code optimization.<n>Our framework leverages the capabilities of small language models (SLMs) and reinforcement learning (RL)<n>Our approach achieves similar or better results compared to state-of-the-art models using shorter training times and smaller pre-trained models.
  arXiv  Detail & Related papers  (2023-12-09T19:50:23Z)
• Energy-efficient Task Adaptation for NLP Edge Inference Leveraging Heterogeneous Memory Architectures [68.91874045918112]
  adapter-ALBERT is an efficient model optimization for maximal data reuse across different tasks. We demonstrate the advantage of mapping the model to a heterogeneous on-chip memory architecture by performing simulations on a validated NLP edge accelerator.
  arXiv  Detail & Related papers  (2023-03-25T14:40:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.