Efficient Post-Training Refinement of Latent Reasoning in Large Language Models

TL;DR

This paper tackles the inefficiencies and rigidity of Chain-of-Thought reasoning by proposing a training-free post-training refinement framework that operates entirely in the latent space of a Coconut-based reasoning backbone. It introduces two complementary modules—Contrastive Reasoning Feedback Search and Residual Embedding Refinement—to steer latent trajectories toward more accurate solutions while stabilizing updates. Empirical results across five benchmarks show consistent accuracy gains, with notable improvements on MathQA (+5%+) and substantial token-efficiency (over 92% fewer generated tokens) compared to text-based reasoning. The approach is lightweight, model-agnostic, and deployable on frozen LLMs, offering practical benefits for resource-limited settings and real-time reasoning tasks.

Abstract

Reasoning is a key component of language understanding in Large Language Models. While Chain-of-Thought prompting enhances performance via explicit intermediate steps, it suffers from sufficient token overhead and a fixed reasoning trajectory, preventing step-wise refinement. Recent advances in latent reasoning address these limitations by refining internal reasoning processes directly in the model's latent space, without producing explicit outputs. However, a key challenge remains: how to effectively update reasoning embeddings during post-training to guide the model toward more accurate solutions. To overcome this challenge, we propose a lightweight post-training framework that refines latent reasoning trajectories using two novel strategies: 1) Contrastive reasoning feedback, which compares reasoning embeddings against strong and weak baselines to infer effective update directions via embedding enhancement; 2) Residual embedding refinement, which stabilizes updates by progressively integrating current and historical gradients, enabling fast yet controlled convergence. Extensive experiments and case studies are conducted on five reasoning benchmarks to demonstrate the effectiveness of the proposed framework. Notably, a 5\% accuracy gain on MathQA without additional training.

Figures (9)

• Figure 1: Overview of our reasoning framework. The bottom path shows how contrastive feedback first identifies the update direction for the reasoning embedding, which is then integrated through residual refinement to produce the current reasoning state.
• Figure 2: Contrastive Reasoning Feedback Search. We compare two models with various reasoning abilities: one stronger ("good") and one weaker ("bad"). The direction from bad to good indicates the path to move.
• Figure 3: Accuracy (%) of different reasoning methods across five benchmarks.
• Figure 4: Latent Reasoning in Training vs. Inference.
• Figure 5: Case study: our method adjusts the latent embedding to reach the correct answer.