[논문 리뷰] Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control

Peihao Wang, Shan Yang외 9명|arXiv (Cornell University)|2026. 03. 10.

Reinforcement Learning in Robotics인용 수 0

한 줄 요약

논문은 최적 제어 TTC 계층을 사전 학습된 LLM에 내장하여 테스트 시 계획을 가능하게 하는 하드웨어 효율적인 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.

연구 동기 및 목표

• 메모리 기반 예측을 넘어서는 추론을 모델 아키텍처에 계획을 통합하여 촉진한다.
• 추론 중 잠재 상태에 대해 유한-수평 LQR 계획을 수행하는 Test-Time Control (TTC) 계층을 도입한다.
• 제어 계층을 통해 엔드투엔드 학습과 그래디언트 전파를 가능하게 하는 differentiable TTC 프레임워크를 개발한다.
• 추론 시 높은 처리량을 위한 CUDA에 융합된 하드웨어 효율적인 심플릭틱 LQR 솔버를 설계한다.
• 미리 학습된 LLM에 어댑터로 TTC 계층을 통합했을 때 도전적인 추론 벤치마크에서 이득을 입증한다.

제안 방법

• 선형 역학 및 이차 비용을 가진 재발 수평의 LQR를 사용한 최적 제어 문제로 추론을 형식화한다.
• TTC(h0, A_t,B_t,Q_t,R_t)를 정의하여 컨텍스트 잠재 상태를 첫 번째 최적 제어 행동으로 매핑한다.
• TTC 계층을 통해 역전파를 가능하게 하는 differentiable KKT 기반 형식을 도출한다(내부 LQR 및 외부 세계 모델 업데이트 포함).
• 시퀀스 리깅 대신 병렬 행렬 곱과 하나의 밀집 역전을 사용하는 하드웨어 효율적인 심플릭틱 반복 솔버를 개발한다.
• 상수 시간 역전을 가능하게 하기 위해 A_t와 R_t를 대각선으로 매개화하고 심플릭틱 반복을 CUDA 커널에 융합하여 처리량을 높인다.
• TTC를 메모리 기반 모듈과 인터리브하는 어댑터로 TTC-Net 아키텍처를 형성한다.

실험 결과

연구 질문

• RQ1 언어 모델 내에서 계획을 최적 제어 문제로 내재화함으로써 추론을 달성할 수 있는가?
• RQ2 테스트 시 제어 계층을 대형 언어 모델 내에서 differentiable하고 엔드투엔드 학습 가능하게 만들 수 있는가?
• RQ3 추론 시에 확장 가능하고 저오버헤드인 계획을 가능하게 하는 하드웨어-알고리즘 공동 설계 전략은 무엇인가?
• RQ4 TTC-강화 모델이 수학적 및 기호적 추론 벤치마크에서 메모리 기반 아키텍처를 능가하는가?

주요 결과

• TTC 계층은 MATH-500에서 수학적 추론 성능을 최대 27.8% 향상시킨다.
• 사전 학습된 LLM과의 통합은 AMC와 AIME에서 Pass@8 성능을 2~3배 향상시킨다.
• 심플릭틱 LQR 솔버는 병렬화 가능하고 하드웨어 효율적 계획을 거의 오버헤드 없이 가능하게 한다.
• TTC-Net은 도전적인 추론 과제에서 일관되게 메모리 기반 모델보다 우수한 성능을 보인다.
• 구조화된 대각 매개변수화 및 융합 CUDA 커널은 추론 및 학습 시 계획의 처리량을 높인다.
• 이 접근 방식은 학습과 추론이 하드웨어 효율적인 최적 제어 계층을 통해 통합된 통합된 관점을 제공한다.