WIDER method improves latent reasoning diversity in language models
Think Wider: Mitigating Latent Rank Collapse in Implicit Chain-of-Thought Reasoning
Machine Learning
Summary
Large language models can think step-by-step to solve problems, but writing out each step takes time and space. The authors found that when these reasoning steps are kept inside the model's memory (latent states) instead of written out, these steps can become too similar and less useful. To fix this, they created WIDER, a training trick that encourages these memory states to be more varied and spread out. This makes the model’s internal thinking clearer and improves its problem-solving without slowing it down.
What this means in practice
- •For language model developers: Train large language models to maintain diverse internal reasoning steps without changing decoding methods, improving reasoning quality efficiently.
- •For ai hardware engineers: Design hardware and systems that better support training with spectral regularization to avoid latent collapse in large-scale language models.
Authors
Yuwen Hao, Menglin Yang
Abstract
Chain-of-thought (CoT) reasoning improves the reasoning ability of large language models by introducing intermediate computation, but explicit rationales increase decoding length, latency, and context cost. Implicit CoT offers a more efficient alternative by moving intermediate reasoning into continuous latent states. However, latent reasoning can be unstable: successive latent states may become overly similar and collapse toward a shared dominant direction, reducing the diversity of the reasoning trajectory. In this work, we identify $\textit{latent rank collapse}$ and propose $\textbf{WIDER}$, a lightweight spectral regularizer for implicit CoT. During training, WIDER estimates the shared direction of each latent trajectory and penalizes projections onto this direction, encouraging latent states to span a broader representational subspace. The method is plug-and-play and leaves the backbone model, latent schedule, and inference-time decoding procedure unchanged. We further formulate this collapse as a geometric bottleneck in implicit reasoning, casting its mitigation as a training-time regularization problem rather than an inference-time decoding change. Extensive experiments show that WIDER improves matched implicit CoT baselines, while mechanistic analyses reveal higher effective rank, lower dominant-direction energy, and reduced redundancy among latent steps. These results highlight latent subspace utilization as an important factor for efficient continuous reasoning, providing a geometric perspective for analyzing and improving implicit CoT. Code is available at https://github.com/whitesweater/WIDER.