[论文解读] Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control

Peihao Wang, Shan Yang+9 more|arXiv (Cornell University)|Mar 10, 2026

Reinforcement Learning in Robotics被引用 0

一句话总结

论文将最优控制的 TTC 层嵌入到预训练的大语言模型中，以实现通过一个硬件高效的 LQR 求解器在测试时进行规划，从而提升在数学与符号任务上的推理能力。

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.

研究动机与目标

• 通过将规划融入模型架构，激发超越基于记忆的预测的推理能力。
• 在推理过程中引入一个 Test-Time Control (TTC) 层，对潜在状态执行有限视窗的 LQR 规划。
• 开发一个可微分的 TTC 框架，使控制层能够端到端学习并进行梯度传播。
• 设计一个硬件高效的共形 LQR 求解器，融入 CUDA 以实现推理时的高吞吐规划。
• 在将 TTC 层作为适配器并入到预训练的大语言模型中时，在具有挑战性的推理基准上展示收益。

提出的方法

• 将推理形式化为一个带有线性动态和二次代价的递推视窗最优控制问题（ receding-horizon LQR ）。
• 定义 TTC(h0, A_t,B_t,Q_t,R_t) 将上下文潜在状态映射到第一步的最优控制动作。
• 推导一个可微分的基于 KKT 的公式，用于对 TTC 层进行反向传播（内部 LQR 和外部世界建模更新）。
• 开发一个硬件高效的共形迭代求解器，用并行矩阵乘法和单次密集求逆来替代顺序 Riccati 求解。
• 将 A_t 和 R_t 参数化为对角矩阵以实现常数时间的求逆，并将共形迭代融合到 CUDA 内核中以提升吞吐量。
• 展示 TTC 作为一个适配器，与基于记忆的模块交错，形成 TTC-Net 架构。

实验结果

研究问题

• RQ1 是否可以通过在语言模型内部实现一个最优控制问题的规划来实现推理？
• RQ2 测试时控制层如何在大型语言模型中实现可微分和端到端可训练？
• RQ3 哪些硬件-算法协同设计策略能够在推理时实现可扩展、低开销的规划？
• RQ4 在数学和符号推理基准上，TTC 增强的模型是否优于仅记忆的架构？

主要发现

• TTC 层在 MATH-500 上将数学推理能力提升高达 27.8%。
• 将 TTC 与预训练的 LLMs 结合，在 AMC 与 AIME 上实现 2–3× 的 Pass@8 提升。
• 一个共形 LQR 求解器实现并行化、硬件高效的规划，开销极小。
• TTC-Net 在具有挑战性的推理任务上始终优于基于记忆的模型。
• 结构化对角参数化与融合的 CUDA 内核为推理和训练中的规划提供高吞吐量。
• 该方法提供一个统一视角，即训练与推理通过一个硬件高效的最优控制层实现整合。