Thinking on the Fly: Test-Time Reasoning Enhancement via Latent Thought Policy Optimization
This addresses robustness issues in reasoning for AI systems, particularly in complex scenarios, though it is an incremental advancement over existing latent reasoning methods.
The paper tackles the brittleness of latent reasoning in large language models on challenging out-of-distribution tasks by introducing Latent Thought Policy Optimization (LTPO), a parameter-free framework that optimizes intermediate latent thought vectors at test time, resulting in substantial accuracy improvements on benchmarks like AIME where baselines fail.
Recent advancements in Large Language Models (LLMs) have shifted from explicit Chain-of-Thought (CoT) reasoning to more efficient latent reasoning, where intermediate thoughts are represented as vectors rather than text. However, latent reasoning can be brittle on challenging, out-of-distribution tasks where robust reasoning is most critical. To overcome these limitations, we introduce Latent Thought Policy Optimization (LTPO), a parameter-free framework that enhances LLM reasoning entirely at test time, without requiring model parameter updates. LTPO treats intermediate latent "thought" vectors as dynamic parameters that are actively optimized for each problem instance. It employs an online policy gradient method guided by an intrinsic, confidence-based reward signal computed directly from the frozen LLM's own output distributions, eliminating the need for external supervision or expensive text generation during optimization. Extensive experiments on five reasoning benchmarks show that LTPO not only matches or surpasses strong baselines on standard tasks but also demonstrates remarkable robustness where others fail. Most notably, on highly challenging AIME benchmarks where existing latent reasoning baselines collapse to near-zero accuracy, LTPO delivers substantial improvements, showcasing a unique capability for complex reasoning.