PaperlayerView exact revision on arXiv

Source-linked AI summary

Think Wider: Mitigating Latent Rank Collapse in Implicit Chain-of-Thought Reasoning

Yuwen Hao, Menglin Yang

arXiv:2609.07406v1cs.LG

TL;DR

Explicit CoT can improve reasoning but carries token, latency, and context costs, while implicit CoT can suffer from trajectories that concentrate around a shared direction. WIDER penalizes that shared direction during training without changing inference. Across multiple backbones and reasoning benchmarks, it improves matched implicit-CoT baselines and is associated with higher effective rank and reduced redundancy.

  • Problem

    Explicit rationales incur token, latency, and context costs, while low-rank latent trajectories can bottleneck the input-specific computation passed to the decoder.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    WIDER is a training-time spectral regularizer that estimates the shared direction of each latent trajectory and penalizes projections onto it, leaving inference unchanged.

    Source: §1 · p. 2, §4.2 · p. 4, §4.3 · p. 5

  • Results

    Across multiple backbones and reasoning benchmarks, WIDER improves matched implicit-CoT baselines, with analyses showing higher effective rank, lower dominant-direction energy, and reduced latent-step redundancy.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    The results highlight latent subspace utilization as a factor for efficient implicit reasoning and offer a geometric perspective for analyzing and improving implicit CoT.

    Source: §7 · p. 8

    Experiments cover models up to 8B parameters and a limited set of reasoning benchmarks; benefits at larger scales remain to be explored.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Explicit chain-of-thought can improve reasoning but incurs token, latency, and context costs, motivating implicit reasoning; yet its latent steps may collapse toward a shared direction. The paper introduces WIDER to regularize that geometry during training and reports accuracy gains alongside more diverse latent trajectories.

1 Introduction

  • Explicit CoT externalizes intermediate computation as natural-language tokens, increasing generated tokens, latency, context usage, and deployment cost.These costs motivate implicit reasoning, which keeps intermediate computation latent rather than fully verbalized.Source: §1 · p. 1
  • Implicit CoT can still be unreliable: latent states may become too similar to encode distinct reasoning steps, concentrating energy on a shared direction.The paper calls this failure mode latent rank collapse.Source: §1 · p. 2
  • WIDER penalizes alignment with each trajectory’s shared dominant direction, encouraging a broader latent subspace while leaving inference unchanged.It adds no inference-time modules, tokens, or decoding steps and keeps the backbone and latent schedule unchanged.Source: §1 · p. 2
  • The paper frames latent rank collapse as a geometric failure in which multiple latent states overuse one dominant direction instead of providing distinct computational directions.This framing motivates treating collapse as a training-time regularization problem.Source: §1 · p. 2
  • WIDER improves matched implicit CoT baselines across multiple backbones and reasoning benchmarks, with analyses linking gains to higher effective rank, lower dominant-direction energy, and reduced latent-step redundancy.The introduction attributes the gains to better latent organization rather than surface output shifts.Source: §1 · p. 2

2 Related Work

Prior approaches to reasoning either verbalize intermediate computation or move it into latent states. This paper distinguishes its focus on the geometry of latent trajectories from prior work on how those trajectories are constructed or supervised.

2 Related Work

  • Explicit CoT methods use intermediate natural-language reasoning, but verbalizing computation increases latency and context usage and forces internal states into text.Prior approaches include zero-shot prompting, answer aggregation, decomposition, rationale bootstrapping, program-aided reasoning, and structured thought search.Source: §2 · p. 2
  • Implicit reasoning methods move computation into continuous hidden states or compact latent representations, using feedback, compression, distillation, auxiliary models, or supervision.Examples include Coconut, CCoT, CODI, and SoftCoT.Source: §2 · pp. 2–3
  • Representation-collapse work studies low-dimensional concentration and anticollapse methods that reduce redundancy, preserve variance, or rebalance spectra.WIDER applies spectral regularization to the per-example latent-token matrix without changing architecture or inference.Source: §2 · p. 3

3 Understanding Latent Rank Collapse

The paper defines latent rank collapse as concentration of a trajectory’s normalized latent states around a dominant direction, leaving allocated steps with little directional diversity. Because the decoder receives input-specific computation through this trajectory, low rank can bottleneck its conditional output family.

3 Understanding Latent Rank Collapse

  • The conceptual comparison contrasts healthy trajectories spanning many directions with collapsed trajectories concentrated near one direction.In the collapsed case, hidden states mainly differ by scalar coefficients.Source: Figure 2 · p. 3
  • Latent rank collapse occurs when normalized latent states concentrate around a dominant direction, so successive steps differ mostly by scale and use only a small part of the available subspace.The trajectory need not be exactly rank one.Source: §3 · p. 3
  • When latent steps vary mainly along one direction, the decoder receives more positions but few independent directions of computation, bottlenecking the conditional output family.The paper describes the latent trajectory as the channel carrying input-specific computation from implicit reasoning to the decoder.Source: §3 · p. 3
  • If H(x) lies in a shared r-dimensional subspace, the decoder’s conditional distribution can be expressed through the lower-dimensional coefficients A(x); rank-one collapse reduces this to scalar coefficients.The paper notes that raising effective rank can enlarge per-input conditional capacity while leaving marginal p(y) unchanged.Source: §3 · p. 4
  • A local update that amplifies one direction more than others can make repeated latent states increasingly aligned, motivating measurement and penalization of the dominant shared direction.The paper presents this as a possible mechanism, not an assumption that the model is globally linear.Source: §3 · p. 4

4 Method

WIDER adds a training-time spectral regularizer that estimates the shared direction of a latent trajectory and penalizes projections onto it. Its auxiliary loss controls redundancy while preserving the original answer-training objective and leaving inference unchanged.

4 Method

  • WIDER estimates each example’s shared latent direction and penalizes projections onto it, while leaving the backbone, latent schedule, and decoding procedure unchanged.The method changes only the training objective.Source: §4 · p. 4

4.1 Latent states and row normalization

  • Row normalization makes the auxiliary loss depend on latent-state directions rather than their absolute norms, while the base loss remains responsible for learning the answer.The auxiliary loss targets redundancy among latent steps.Source: §4.1 · p. 4

4.2 Shared-direction proxy

  • The shared-direction proxy sums normalized states and uses the resulting direction to summarize their common component, with stop-gradient applied.The proxy is a cheap operational summary, not an exact top singular-vector estimator; under a spiked latent-step model, it identifies the same shared component as the leading right singular direction.Source: §4.2 · p. 4

4.3 Projection regularization

  • The final objective adds a spectral term to the original latent-reasoning loss, discouraging reuse of the same shared component across latent positions.The added training cost is O(Bmd), and the auxiliary loss is removed at inference.Source: §4.3 · pp. 4–5

5 Experiments

The experiments test WIDER on matched implicit-reasoning baselines across backbones and benchmarks, then examine the regularization coefficient. Accuracy improves broadly, with a peak at λ = 0.10 in the reported sweep.

5.1 Experimental Setup

  • The study evaluates CODI and SIM-CoT on two backbones and four reasoning benchmarks, using matched configurations that differ by the WIDER loss.The benchmarks span binary, multiple-choice, and open-numeric answer formats.Source: §5 · p. 5, §5.1 · p. 5

5.2 Main Results

  • 15 of 16 matched task comparisons improve; the exception is LLaMA-3B SIM-CoT on CommonsenseQA, down 0.49 accuracy points.The main-results table compares matched CODI/SIM-CoT configurations with and without WIDER.Source: §5.2 · p. 6, Table 1 · p. 6
  • LLaMA-3.1-8B CODI improves on all four benchmarks, while effective rank rises from 1.715–2.092 to 3.999–5.523.Accuracy gains are +2.57, +0.32, +4.05, and +1.18 percentage points on StrategyQA, CommonsenseQA, ASDiv-Aug, and AQuA, respectively.Source: §5.2 · p. 6

5.3 Ablation Experiments

  • Accuracy peaks at λ = 0.10, then drops at λ = 0.20 as the spectral penalty over-regularizes the latent trajectory.The peak value is used for the matched comparisons, and the curve is described as robust to moderate λ perturbations.Source: §5.3 · pp. 6–7, Figure 3 · p. 6

6 Mechanistic Analysis

Mechanistic analyses examine repetition, dominant-direction concentration, the content aligned with that direction, and answer-prior shifts. WIDER reduces measured redundancy and dominant-direction energy while increasing effective rank; marginal answer-distribution metrics change little.

6.1 Are Latent Tokens Repetitive?

  • Mean pairwise latent-step NMI falls from 0.472 to 0.144 at k=32 and from 0.306 to 0.142 at k=16.The analysis uses ASDiv-Aug with SIM-CoT; higher off-diagonal NMI indicates stronger repetition.Source: §6.1 · p. 7, Table 2 · p. 7

6.2 Does WIDER Release the Trajectory from a Dominant Shortcut?

  • Effective rank rises from 1.487 to 3.817 as energy shifts from the first singular direction (0.989) into the tail (0.011 to 0.264).Prefix effective rank continues accumulating through s6 with WIDER, while the baseline saturates near 1.5 after the first few steps.Source: §6.2 · p. 7, Figure 6 · p. 8

6.3 What Is Encoded in the Dominant Direction?

  • The dominant direction’s shortcut-like vocabulary share falls from 16.03% to 0.00% after WIDER.The authors describe this as a diagnostic, not a causal statement.Source: §6.3 · p. 8, Table 3 · p. 8

6.4 Is the Gain a Shifted Answer Prior?

  • The four marginal p(y) metrics each move by less than one percentage point, below the +2.22 pp accuracy gap.The authors report that the marginal answer distribution is essentially preserved while per-example accuracy rises.Source: Figure 7 · p. 8, §6.4 · p. 8

7 Conclusion

The paper frames latent rank collapse as a geometric bottleneck in implicit CoT and introduces WIDER, a training-time spectral regularizer. Across multiple backbones and reasoning benchmarks, WIDER improves accuracy over matched implicit-CoT baselines.

7 Conclusion

  • Across multiple backbones and reasoning benchmarks, WIDER delivers consistent accuracy gains over matched implicit-CoT baselines.The paper presents latent subspace utilization as a key factor for efficient implicit reasoning.Source: §7 · p. 8
  • WIDER penalizes projections onto the shared direction to mitigate latent rank collapse during training.The method leaves inference unchanged.Source: §7 · p. 8
  • Latent rank collapse occurs when successive continuous latent states collapse toward a shared direction, leaving additional latent positions underused.The paper characterizes this as a per-example geometric failure that bottlenecks {p(y | x)} by the rank of the latent trajectory.Source: §7 · p. 8

Limitations

The evaluation covers models up to 8B parameters and a limited set of reasoning benchmarks. Experiments use six latent tokens, and the geometric diagnostics do not make individual latent tokens directly human-interpretable or identify their specific reasoning roles.

Limitations

  • The experiments cover models up to 8B parameters and a limited set of reasoning benchmarks.Whether the benefits persist for substantially larger models and larger-scale training remains unexplored.Source: Limitations · p. 9
  • The experiments use a fixed trajectory of six latent tokens.The behavior of latent rank collapse and WIDER with longer or dynamically allocated trajectories remains to be investigated.Source: Limitations · p. 9
  • Geometric diagnostics characterize latent-trajectory organization but do not make individual latent tokens directly human-interpretable or identify their specific reasoning roles.WIDER inherits the interpretability limitations of latent reasoning.Source: Limitations · p. 9

Ethics Statement

The method is a training-time mechanism that encourages broader latent trajectories without changing the backbone or inference procedure. The appendix qualifies its geometric arguments and explains the scope and assumptions of the trajectory diagnostics and projection loss.

Ethics Statement

  • The spectral regularizer is a training-time mechanism that encourages latent trajectories to span a broader subspace while leaving the backbone architecture and inference procedure unchanged.The stated goal is to address latent rank collapse in implicit chain-of-thought reasoning.Source: Ethics Statement · p. 9

A.1 Common-subspace factorization and a local concentration scenario

  • Row normalization preserves the trajectory’s row span and algebraic rank while reducing diagnostic sensitivity to variation in row norms.The normalized trajectory is a positive invertible diagonal rescaling of the original trajectory.Source: §A.1 · p. 12
  • The common-subspace proposition is an exact reduced-coordinate reparameterization only under one fixed subspace shared across inputs, not an information-theoretic bound on real-valued coordinates.The concentrated subspace may vary with x.Source: §A.1 · p. 12
  • Under the stated local linear-update assumptions, repeated application of the same update is one possible route for additional latent steps to occupy few new representation-space directions.The concentration rate is controlled by |λ2/λ1| per step.Source: §A.1 · p. 12

A.2 Justifying the mean-direction proxy and the projection loss

  • The squared norm of the mean direction summarizes average pairwise dot products and motivates using the column-sum direction as a proxy for what latent steps share.The identity is norm-weighted; it coincides with average pairwise cosine alignment for exactly unit-normalized nonzero rows and approximates it when ε is negligible relative to their norms.Source: §A.2 · p. 12
  • The projection loss upper-bounds the squared strength of the shared mean direction, but is not a complete surrogate for the full spectrum when the proxy aligns weakly with the top singular direction.When the proxy aligns with that direction, the loss suppresses the leading spectral component.Source: §A.2 · pp. 12–13

A.3 Signed-mean cancellation as a boundary case

Signed-mean cancellation is an idealized boundary case in which a rank-one trajectory can have a zero or unrepresentative mean direction. Under the stated operational criterion, none of the audited trajectories qualified; the appendix also describes dataset tables and their scope.

A.3 Signed-mean cancellation as a boundary case

  • Exact sign symmetry can make the signed mean zero despite a rank-one trajectory, while near balance with off-axis perturbations can make the mean direction unrepresentative of its leading axis.The construction requires exact sign symmetry; whether an approximate version occurs in practical trajectories is an empirical question.Source: §A.3 · p. 13
  • 0/1,038 trajectories qualified as near-rank-one cancellation in either the baseline or trained WIDER checkpoint under the stated criterion.The audit used ASDiv-Aug trajectories from the LLaMA-3B + SIM-CoT checkpoint pair.Source: §A.3 · p. 13

B Datasets

  • Table 4 summarizes train/test sizes for four reasoning datasets, and the same splits are used across backbones and method variants.Table 5 lists datasets, pretrained models, and their licenses or terms of use; Table 6 gives one representative input-output example per dataset.Source: B Datasets · p. 13, Table 4 · p. 13

C Backbone and Training Configuration

The experiments use matched baseline–WIDER comparisons, changing only whether the spectral auxiliary loss is added. Training configurations are specified across two backbone families, while inference procedures remain unchanged.

C Backbone and Training Configuration

  • Matched CODI and SIM-CoT pairs hold the backbone, data, prompts, latent schedule, and checkpoint selection fixed; the intended difference is the spectral auxiliary loss.No-CoT and SFT-CoT are contextual references, not spectral-objective ablations.Source: C Backbone and Training Configuration · pp. 13–14
  • LLaMA-3.2-3B-Instruct is the default backbone, with Qwen3-1.7B used as a second backbone family and scale.The default backbone is used for main matched comparisons, the coefficient ablation, and the ASDiv-Aug mechanism analysis.Source: C Backbone and Training Configuration · p. 14
  • CODI training uses batch size 128, bfloat16, and 10 epochs; SIM-CoT uses the same training settings, including backbone-specific learning rates.Learning rates are 1e−5 for LLaMA-3.2-3B-Instruct and 1e−4 for Qwen3-1.7B.Source: C Backbone and Training Configuration · p. 14
  • Experiments run on NVIDIA H800 80 GB GPUs with CUDA 12.8.Source: C Backbone and Training Configuration · p. 14

D Results on LLaMA-3.1-8B

On LLaMA-3.1-8B, WIDER improves accuracy across four CODI-matched benchmarks and raises effective rank across datasets. The accompanying vocabulary-projection protocol is a population-level diagnostic, not a token-level interpretation.

D Results on LLaMA-3.1-8B

  • WIDER substantially raises the effective rank of the first six latent states on every dataset.The analysis tests whether the geometric pattern also appears in the larger backbone.Source: D Results on LLaMA-3.1-8B · p. 14, Table 8 · p. 15
  • Latent rank collapse persists in the larger backbone: states can still concentrate along a few directions despite its wider representation space.Higher effective ranks under WIDER indicate broader trajectory subspaces, rather than an effect obtained merely by adding model capacity.Source: D Results on LLaMA-3.1-8B · p. 14

E Dominant-Direction Vocabulary Projection Protocol

  • The vocabulary diagnostic ranks the top 30 tokens by dominant-direction projection magnitude and averages five surface-form categories across held-out prompts.Matched baseline and WIDER rows use the same backbone, dataset, and prompt subset; paired differences reflect shifts in where the dominant direction concentrates.Source: E Dominant-Direction Vocabulary… · pp. 14–15

D Results on LLaMA-3.1-8B

E Dominant-Direction Vocabulary Projection Protocol

  • The diagnostic supports population-level claims about vocabulary regions, not token-level explanations; its coarse categories can mix stop-word collocations with task-specific signal.Individual rankings are not invariant under unit-norm reparameterizations of nearby singular directions.Source: E Dominant-Direction Vocabulary… · p. 15
  • The protocol fixes N=30 and reuses prompts within each backbone–dataset pair; bytepair pieces are categorized independently rather than reassembled into words.The categorizer is released with the code for recomputing share values from released checkpoints.Source: E Dominant-Direction Vocabulary… · pp. 15–16
Loading 2609.07406v1…