Source-linked AI summary

Do Latent-CoT Models Think Step-by-Step? A Mechanistic Study on Sequential Reasoning Tasks

Jia Liang, Liangming Pan

arXiv:2602.00449v1cs.AIcs.LG

TL;DR

The paper asks whether latent-CoT models genuinely execute sequential reasoning or use compressed shortcuts, and studies CODI on controlled polynomial-iteration tasks. Mechanistic analyses show faithful bridge-state computation on short horizons but partial late-stage reasoning on longer or less compressible tasks, clarifying when latent-CoT succeeds and where it fails.

  • Problem

    Latent-CoT hides reasoning in continuous vectors, making it difficult to verify genuine multi-step computation rather than sophisticated heuristics.

  • Method

    The study analyzes CODI on sequential polynomial-iteration tasks using logit-lens decoding, probes, attention analysis, and activation patching.

  • Results

    CODI forms bridge states on 2–3 hop tasks but shifts to partial late-bottleneck reasoning for longer horizons, with signatures largely absent under prime moduli.

  • Takeaways & Limitations

    Latent-CoT yields faithful iterative computation when the task is compressible, but otherwise favors compressed late-stage solutions or fails to stabilize.

  • Takeaways & Limitations

    With fixed latent compute, CODI struggles on effectively incompressible sequential computations, where performance is better under explicit CoT.

Abstract

from arXiv · show

Latent Chain-of-Thought (Latent-CoT) aims to enable step-by-step computation without emitting long rationales, yet its mechanisms remain unclear. We study CODI, a continuous-thought teacher-student distillation model, on strictly sequential polynomial-iteration tasks. Using logit-lens decoding, linear probes, attention analysis, and activation patching, we localize intermediate-state representations and trace their routing to the final readout. On two- and three-hop tasks, CODI forms the full set of bridge states that become decodable across latent-thought positions, while the final input follows a separate near-direct route; predictions arise via late fusion at the end-of-thought boundary. For longer hop lengths, CODI does not reliably execute a full latent rollout, instead exhibiting a partial latent reasoning path that concentrates on late intermediates and fuses them with the last input at the answer readout position. Ablations show that this partial pathway can collapse under regime shifts, including harder optimization. Overall, we delineate when CODI-style latent-CoT yields faithful iterative computation versus compressed or shortcut strategies, and highlight challenges in designing robust latent-CoT objectives for sequential reasoning.

1. Introduction

The paper investigates whether latent-CoT models genuinely perform sequential computation or rely on shortcuts, using CODI and controlled polynomial-iteration tasks. Its analyses reveal faithful bridge-state computation on short tasks but partial, late-stage reasoning as depth increases, with behavior shaped by modulus compressibility.

  • Motivation: Latent-CoT adds internal computation through continuous thought vectors or recurrent updates, but its hidden representations make genuine reasoning difficult to verify.Accuracy alone cannot distinguish multi-step computation from sophisticated heuristics.
  • Research approach: The study uses CODI and strictly sequential polynomial tasks with ground-truth intermediates to test whether latent states represent and propagate step-by-step computation.Logit-lens decoding, linear probes, attention analysis, and activation patching localize and trace the computation.
  • Main findings: On 2–3 hop tasks, CODI constructs intermediate bridge states while routing the final input through a separate direct pathway to the answer readout.The computation therefore combines latent intermediate-state construction with late-stage input delivery.
  • Main findings: For n ≥4, CODI rarely completes a full latent rollout and instead retains only the final one or two intermediates in a late-bottleneck trace.This pattern indicates a partial latent reasoning route rather than faithful execution of the entire chain.
  • Modulus dependence: Mechanistic signatures persist across many composite moduli but largely vanish for prime moduli, where accuracy and intermediate-state decodability decline.The paper links this split to many-to-one contractions in composite rings versus full-history dependence under prime moduli.
  • Interpretation: Teacher-guided compression appears to drive the partial late-bottleneck strategy under composite moduli, while full-history computations destabilize latent rollouts under prime moduli.The comparison highlights when latent-CoT favors compressed solutions over faithful iterative computation.

2. Related Work

Related work has characterized mechanistic structures in implicit and explicit reasoning, while latent-CoT studies have mainly shown that hidden representations can be decoded or steered. This paper addresses the remaining mechanistic gap with causal analysis of CODI on strictly sequential tasks.

  • Implicit reasoning: Mechanistic studies of implicit reasoning identify layer roles, internal structures, and depth-related bottlenecks in multi-hop inference.Some work also reports abrupt training transitions in reasoning-like behavior.
  • Explicit reasoning: Explicit-CoT research finds attention heads, circuits, and neuron activations that propagate or encode intermediate variables.Generated CoT tokens can serve as external memory while internal states are also updated.
  • Latent reasoning: Prior latent-CoT work shows that representations can be decoded or steered, but these findings are primarily correlational rather than mechanistic.Related work also identifies conditions such as representation superposition under which latent CoT can outperform explicit CoT on graph reasoning.
  • Present work: This paper distinguishes itself by giving a mechanistic account of CODI using strictly sequential tasks with ground-truth intermediates.The controlled setting supports direct tests of representation, propagation, and causal pathway use.

3. Approach

The paper trains CODI with an explicit-CoT teacher and latent-thought student, then probes the student’s internal computation on polynomial-iteration tasks. Feature-space distillation aligns hidden states at the pre-answer boundary, while mechanistic tools distinguish iterative updates from shortcut routing.

  • CODI framework: CODI compresses an explicit CoT into a short sequence of continuous thought vectors while retaining a teacher and student mode within one model.The teacher emits supervised CoT steps; the student emits fixed-length latent thoughts and predicts the final answer.
  • CODI framework: Feature-space self-distillation aligns teacher and student hidden representations at a designated pre-answer boundary without requiring emitted rationales.The alignment uses an ℓ1 loss with stop-gradient on the teacher.
  • Task design: Polynomial-iteration tasks provide a controlled sequential testbed with explicit state updates and ground-truth intermediate states.The task generates states over integers modulo m using s_t = s_{t−1}x_t + b (mod m).
  • Training procedure: The student predicts the final state after a latent-thought trajectory and uses feature distillation at the [Ans] boundary.Training combines final-answer cross-entropy with teacher-student representation matching.
  • Mechanistic analysis: Logit-lens decoding, attention analysis, and activation patching localize latent states and identify which pathways are necessary for correct answers.These tools track state decodability, information routing, and causal dependence on inputs or latent positions.

4. Empirical Experiments

CODI forms and maintains required bridge states for short sequential tasks, while routing final inputs through a separate direct pathway and combining streams at the answer readout. As hop depth increases, intermediate visibility collapses toward late states, and prime moduli eliminate the observed partial-rollout signature.

  • Two-hop tasks: The bridge state s2 becomes decodable throughout six latent steps, indicating that CODI forms and maintains it in the latent channel.Mean decoded probability ranges from 0.359 to 0.709, while linear-probe confidence approaches 1.
  • Two-hop tasks: Attention analysis shows a two-stream mechanism: the latent trajectory stores the bridge state, while x3 follows a direct copy-like route to [Ans].The final prediction mixes both streams in the [Ans] residual stream before unembedding.
  • Two-hop tasks: 100% accuracy is restored by patching x3-corrupted runs at [Ans], whereas latent-position patches yield 0% recovery, supporting direct x3-to-output routing.For x2 corruption, latent-position patching instead produces substantial recovery, implicating the latent channel in storing the intermediate state.
  • Three-hop tasks: CODI forms both bridge states on three-hop tasks while copying the final input x4 directly to [Ans].This extends the short-task pattern beyond the two-hop setting.
  • Longer-hop tasks: For n ≥4, only late intermediates are reliably decodable, while the final input x_n+1 remains on a separate direct pathway and fuses with the latent readout at [Ans].The latent stream may expose s_n alone or the final pair s_n−1 and s_n, rather than a full intermediate chain.
  • Task-definition effects: Under prime moduli, performance drops sharply and the late-intermediate partial-rollout signature disappears, unlike the largely stable behavior under composite moduli.In prime-modulus runs, s_n is no longer decodable during latent steps.

5. Theoretical Analysis

The theoretical analysis explains why polynomial-iteration tasks differ sharply between composite and prime moduli: composite updates can contract state information, whereas prime-modulus updates preserve full-history dependence.

  • Task setup: The paper analyzes how the algebraic structure of R_m changes the intrinsic dependence of the final state on earlier inputs.The task uses sequential affine updates over the ring R_m = Z/mZ.
  • Bijections and contractions: An affine update is bijective exactly when its multiplier is a unit, meaning gcd(x, m) = 1.This bijection criterion determines whether an update preserves or contracts state information.
  • Bijections and contractions: When d = gcd(x, m) > 1, multiplication maps m states to m/d outputs, with every output having d preimages.The corresponding affine map is therefore d-to-1 and erases distinctions between prior states.
  • Empirical implication: CODI and the non-CoT baseline degrade sharply when m is prime, whereas composite moduli support the compression patterns predicted by the theory.Table 1 aggregates accuracy across sequence lengths for 3-layer, 2-head transformers.
  • Prime-field regime: For prime moduli, every update is a permutation, so no state information is irreversibly discarded and the final state typically depends on the full input history.The unrolled recurrence retains multiplicative chains of length Θ(T), requiring sustained state propagation for long horizons.
  • Composite-ring regime: For composite moduli, a recent non-unit update typically creates a short effective memory, so the final label can depend mainly on a low-entropy history summary and a short terminal suffix.The terminal all-unit suffix satisfies Pr(L ≥ k) = u(m)^k, and contraction events shrink distinguishable states by gcd-based factors.

6. Ablation Study

The ablation study finds that CODI’s long-horizon mechanism remains a late-intermediate pathway rather than becoming a deeper latent rollout, while attention directly routes the final input to the answer.

  • Long-horizon ablations: Across long-horizon settings, CODI primarily encodes late intermediates and increasing latent steps does not reliably induce a deeper rollout.The robustness evaluation uses n = 31 with a 3-layer, 2-head student and m = 50 across multiple latent-step counts and architectures.
  • Loss ablations: Removing feature distillation alone preserves the late-bottleneck signatures, but removing both distillation and the teacher objective eliminates them.This identifies the teacher loss as the key driver of the observed late-bottleneck latent mechanism.
  • Attention routing: The [Ans] token attends strongly to the final input token x_n+1, consistent with a direct pathway routing that input to the answer readout.The specialized attention head provides evidence for direct final-input routing alongside the latent pathway.

7. Discussion

The discussion contrasts CODI with explicit CoT and standard non-CoT models, emphasizing that CODI benefits from compressible task structure but struggles when sequential computation preserves long-range dependence.

  • Comparison with Non-CoT: A standard non-CoT transformer can learn a full internal rollout, with intermediate state s_i decodable at the next input position x_i+1.This implicit trace is nevertheless described as brittle under some modes.
  • Comparison with CoT: Explicit CoT training reliably recovers the full intermediate trajectory across model scales and optimization settings.The discussion presents explicit CoT as more robust for recovering sequential state trajectories.
  • Task dependence: On prime-modulus tasks, CoT training substantially outperforms both CODI and the standard non-CoT transformer.The comparison exposes a limitation of CODI-style latent CoT when the underlying computation is effectively incompressible.
  • Task dependence: CODI recovers much of the performance gap in more compressible regimes, such as composite moduli.Composite-ring contractions provide a mechanistic lens for why late-bottleneck strategies can remain viable.

8. Conclusion

CODI’s latent-CoT behavior depends on task compressibility: it can perform faithful iterative computation on some settings but shifts toward compressed late-stage solutions or fails when sequential information cannot be stably compressed.

  • Composite-modulus tasks support step-by-step computation on 2–3 hops but induce a late-intermediate partial rollout on longer horizons.
  • Prime-modulus tasks largely lack these mechanistic signatures, with reduced accuracy and disappearing intermediate-state decodability.
  • Many-to-one contractions under composite moduli make updates more compressible and bias CODI toward late updates within a fixed latent budget.
  • Future work should test whether compressibility dependence persists across latent-CoT objectives and architectures and develop compute allocation matched to task demands.

Impact Statement

The paper frames mechanistic interpretability as a way to determine whether latent reasoning uses faithful computation or shortcuts, using logit-lens decoding and attention analysis to inspect hidden representations and routing.

  • Latent chain-of-thought moves reasoning into hidden activations, making mechanistic analysis important for distinguishing faithful multi-step computation from shortcuts.
  • Logit lens: The logit lens applies the model’s output decoder to earlier hidden states to locate when correct intermediate values become decodable.
  • Attention maps: Attention maps represent routing as normalized weights from query positions to key/value positions, with rows indexing queries and columns indexing attended positions.
  • In the two-hop task, [Ans] attends strongly to x3, supporting a copy-like pathway that routes the final input directly into the answer residual stream.
  • Linear probing: Probing measures whether task-relevant variables are linearly decodable from frozen hidden activations using supervised predictors.

C.1. Linear Probing Implementation Details

The implementation uses linear probes, attention interventions, and activation patching to measure which task variables are represented at selected positions and depths and which components support correct predictions.

  • Linear probes are trained on frozen activations from correct runs and evaluated with held-out classification accuracy.
  • Each probe is a bias-free linear classifier mapping 256-dimensional hidden vectors to 50 discrete classes.
  • Representations are sampled from input, boundary, six latent-thought, EoT, and answer positions across four residual-stream depths.
  • The dataset retains approximately 5,000 test activations per sequence length but filters to examples with correct final predictions before probe training.
  • Probe evaluation uses a 64%/16%/20% train/validation/test split, Adam optimization, and a 1/50 = 0.02 chance baseline.
  • Activation patching: Activation patching replaces corrupted-run activations with clean-run counterparts to test which components are causally necessary for correct outputs.
  • Results: For x2 corruption, the largest recovery occurs at latent positions [ℓ1] and [ℓ2], whereas x3 corruption recovers mainly when patching [Ans].

F. Mechanistic Analysis on Three-hop Polynomial task.

Three-hop analyses find temporally ordered intermediate-state representations and causal support for early latent computation, while longer tasks often show a two-step partial rollout focused on the final intermediates.

  • In the three-hop task, s1 becomes decodable early across the latent trajectory, while s2 becomes decodable at [EoT], consistent with sequential computation.
  • Patching x1- or x2-corrupted runs into early latent steps substantially restores accuracy, supporting a causally relevant latent representation for s2.
  • For x3 corruption, patching has its strongest effect at [EoT], consistent with s3 being formed or consolidated near the latent boundary.
  • The full step-by-step pattern is not uniform: some runs expose only s3 in latent steps without clear evidence of s2 formation.
  • For longer composite-modulus tasks, s_n−1 becomes decodable earlier and s_n later, producing a two-step partial latent rollout.
  • Compressibility: The polynomial update map is d-to-1 when d = gcd(x, m), so composite-modulus updates can contract multiple inputs into the same output.

I. Ablation Study

Ablations show that CODI’s late-intermediate pathway is robust to latent-step and architecture changes but depends on teacher supervision. Removing only distillation preserves the mechanism, whereas student-only training eliminates it.

  • Latent-step ablations: Across p ∈ {1, 2, 3, 6, 9, 12, 20}, CODI consistently represents late intermediate states on long n-hop tasks.The final state s_n appears most reliably, with occasional representation of (s_{n−1}, s_n).
  • Depth and width ablations: Across 2–7 layers and 2–8 attention heads, CODI primarily routes late intermediates on long n-hop tasks.s_n is most reliable, while (s_{n−1}, s_n) sometimes appear.
  • Loss ablations: CODI combines teacher, student, and feature-space distillation losses during training.The teacher predicts an explicit state trace, the student predicts the final answer, and distillation aligns representations near the answer boundary.
  • Loss ablations: Removing distillation alone preserves late-state formation and the attention route from the final input token to the answer boundary.The late trace still contains s_n, and occasionally both s_{n−1} and s_n.
  • Loss ablations: Removing both distillation and teacher supervision eliminates reliable late-step encoding of s_n or (s_{n−1}, s_n).The result suggests teacher supervision induces state representations that the student can compress into a late-bottleneck computation.
  • Non-CoT comparison: Standard Non-CoT Transformers provide a contrast: their intermediate-state rollout is brittle across the compared configurations.The 5-hop visualization contrasts a clear rollout in one model with little or no trace in another, unlike latent-CoT’s late-bottleneck tendency.
Loading 2602.00449v1…