PaperlayerView exact revision on arXiv

Source-linked AI summary

Structural Process Supervision for Latent Chain-of-Thought Reasoning

Yiqi Li, Xu Chen, Chen Ju, Jiangchao Yao, Zhaoyang Li, Jinsong Lan, Xiaoyong Zhu, Bo Zheng, Yu Wang

arXiv:2609.09928v1cs.AI

TL;DR

Existing latent reasoning methods lack direct process supervision for latent embeddings, risking representation collapse and uneven information distribution. PMPS uses learnable prototypes and progressive alignment to supervise latent and explicit CoT representations, and it consistently improves latent-reasoning accuracy with outputs below half the explicit-CoT length. These results include outperforming CoT-SFT on GPT-2 and leading latent baselines across larger models and a harder task.

  • Problem

    Existing latent reasoning approaches lack direct process supervision for latent embeddings, which can cause representation collapse and uneven information distribution.

  • Method

    PMPS projects latent and explicit CoT embeddings into a shared prototype space for many-to-many soft alignment, while PSA schedules positional guidance from sequential to adaptive matching.

  • Results

    PMPS consistently outperforms latent reasoning baselines across model families, scales, and task difficulties while reducing output token length by more than 50% relative to explicit CoT.

  • Takeaways & Limitations

    Structural process supervision enables latent reasoning to achieve strong generalization and robust performance with substantially lower output length than explicit CoT.

Abstract

from arXiv · show

Latent reasoning approaches enhance token-level efficiency and robustness by replacing verbose, explicit chain-of-thought (CoT) tokens with compact continuous-space embeddings. However, existing methods lack direct process supervision over these latent embeddings, which often leads to representation collapse and uneven information distribution. To address this, we propose Prototype-Mediated Process Supervision (PMPS), which introduces learnable reasoning prototypes as semantic anchors to provide structural process-level supervision for latent reasoning. PMPS projects latent embeddings and explicit CoT embeddings into a shared prototype space, achieving many-to-many soft alignment between unequal-length representations through prototype assignment. Meanwhile, we introduce a Progressive Sequential Alignment (PSA) module to further guide training: positional priors initially encourage sequential alignment structure, then gradually relax to permit adaptive matching. Experimental results show that PMPS compresses output token length to under 50% of explicit CoT on GSM8K-Aug. Compared to leading baseline SIM-CoT, our method achieves average accuracy gains of 2.08% across different model families. On GPT-2, PMPS even surpasses CoT-SFT. On larger models and a more challenging task, PMPS consistently attains the highest accuracy among all latent reasoning methods with comparable output length.

1 INTRODUCTION

Explicit CoT improves reasoning but incurs efficiency and expressiveness limitations, motivating latent reasoning methods. PMPS addresses missing process supervision by structuring latent representations and achieves shorter outputs with stronger accuracy across evaluations.

  • Explicit CoT reasoning has low inference efficiency and constrained expressiveness, motivating more efficient latent reasoning paradigms.
  • Existing latent methods can leave intermediate reasoning unguided, causing representation collapse or imposing rigid alignment constraints.
  • PMPS uses learnable prototypes, shared embedding projections, soft many-to-many alignment, and bidirectional cross-prediction to supervise latent reasoning structure.
  • PSA initially imposes sequential structure with a Gaussian positional prior, then relaxes it through cosine annealing for adaptive alignment.
  • Over 50% shorter reasoning length and 3.96% higher accuracy than explicit CoT were achieved on GSM8K-Aug with GPT-2.
  • PMPS is reported to outperform leading latent reasoning baselines across model families while preserving inference efficiency.

2 RELATED WORK

Prior work replaces explicit reasoning with latent embeddings, but existing process supervision requires discrete reasoning steps or rigid alignment. PMPS instead uses prototype-mediated soft alignment to supervise continuous latent representations without an auxiliary decoder.

  • iCoT and Coconut replace explicit CoT tokens with latent embeddings through curriculum learning, while CODI distills supervision at the answer position.
  • SIM-CoT uses an auxiliary decoder to enforce rigid one-to-one assignment between explicit CoT steps and latent embeddings.
  • Figure 2 depicts PMPS training with a parallel teacher-student architecture that transfers explicit CoT supervision to latent embeddings through prototype-space cross-prediction.
  • PMPS addresses prior alignment limitations through structural prototype-mediated soft alignment without an auxiliary decoder or hard CoT partition.
  • Existing process supervision uses explicit discrete reasoning steps, whereas learnable prototypes can structure continuous embedding spaces.

3 METHODOLOGY

The methodology supervises latent reasoning through prototype-mediated alignment with explicit CoT embeddings, then adds a curriculum-scheduled positional prior to preserve sequential structure while allowing adaptive matching. Training-only components provide process supervision without adding inference cost.

  • Prototype-Mediated Process Supervision: PMPS uses teacher-generated explicit CoT embeddings as supervision signals because answer-only training lacks process-level guidance and can cause latent homogenization.The explicit and latent embeddings come from the same model, facilitating knowledge transfer.
  • Prototype-Mediated Process Supervision: PMPS introduces learnable prototypes as shared semantic anchors because direct one-to-one distillation is infeasible when latent embeddings are fewer than CoT tokens.Both representation sets are projected into a structured prototype space for soft alignment.
  • Prototype-Mediated Process Supervision: Sinkhorn-Knopp assignments produce soft prototype codes with per-token distributions and uniform prototype utilization, supporting many-to-many matching between latent and explicit CoT representations.The resulting bridge allows one latent embedding to attend to multiple CoT embeddings and one CoT embedding to contribute to multiple latent embeddings.
  • Prototype-Mediated Process Supervision: Bidirectional cross-prediction trains latent and CoT representations to explain each other’s prototype assignments, encouraging diverse and semantically meaningful reasoning information.The CoT-to-latent direction explains explicit CoT embeddings collectively, while the reverse direction encourages latent embeddings to capture diverse CoT semantics.
  • Progressive Sequential Alignment: PSA adds a Gaussian positional prior that favors sequential alignment early, then cosine annealing relaxes it to permit adaptive matching as training progresses.The prior operates only on matching scores, adds no learnable parameters, and incurs zero inference overhead.
  • Overall Training Objective: The teacher task, prototype layer, projection head, and PSA module are used only during training, leaving inference architecture and computational cost identical to the student task.The overall objective combines the base loss with a weighted process-supervision loss.

4 EXPERIMENTS

Across diverse backbones and benchmarks, PMPS improves latent-reasoning accuracy while substantially shortening outputs, generalizing to larger models and harder tasks. Ablations indicate that prototype-mediated alignment and PSA contribute to its performance and prevent latent-representation collapse.

  • Experimental setup: PMPS is evaluated on GSM8K-Aug, three OOD benchmarks, and three backbone models, with accuracy and output token length reported.The experiments compare explicit CoT, curriculum, token-compression, and self-distillation baselines.
  • Overall performance: PMPS exceeds CoT-SFT by 3.96% in average accuracy on GPT-2 while using 47.1% of its output tokens.It also surpasses CODI and SIM-CoT by 4.33% and 2.86% in average accuracy, respectively.
  • OOD generalization: PMPS surpasses SIM-CoT by 3.06%, 1.88%, and 1.46% in average OOD accuracy across the three backbone models.The corresponding PMPS averages are 51.82%, 53.23%, and 58.10%.
  • Inference efficiency: PMPS reduces output length to 47.1%, 37.5%, and 47.2% of explicit CoT across the three models while maintaining comparable or superior accuracy.Its projection head and prototypes add less than 2M parameters, under 1% of the base model.
  • Scaling and harder tasks: On LLaMA-3.2-3B-Instruct, PMPS beats CODI and SIM-CoT by 3.74% and 1.22% in average accuracy while maintaining comparable output length.On MATH, it gains 2.1% over CODI and compresses output to 13% of CoT-SFT.
  • Representation analysis: PMPS avoids the odd-even collapse observed in CODI and the tail collapse observed in SIM-CoT.CODI shows within-group similarities of 0.693 and 0.800 versus 0.250 cross-group similarity; SIM-CoT’s tail similarity averages 0.805 versus 0.591 for the head.
  • Ablation study: Removing L_l2c, prototypes, L_c2l, or PSA reduces average performance by 3.78%, 3.49%, 2.05%, and 0.76%, respectively.Moderate PSA enforcement performs best: α_max = 2 for GPT-2 and α_max = 6 for Llama-3.2-1B, while excessive enforcement hurts performance.

5 CONCLUSION

PMPS provides structural process-level supervision for latent reasoning through soft alignment and achieves strong accuracy with substantially shorter outputs. Its current evaluation, however, is mainly limited to mathematical reasoning.

  • PMPS applies structural process-level supervision to latent embeddings through soft alignment.
  • Across diverse model families, model sizes, and task difficulties, PMPS achieves state-of-the-art performance among latent reasoning approaches.
  • PMPS reduces output token length by more than 50% compared with explicit CoT.
  • The evaluation is mainly limited to mathematical reasoning, leaving broader reasoning tasks for future study.

A.1 BENCHMARKS

The benchmark suite combines GSM8K-Aug with harder and out-of-domain arithmetic tasks, while MATH provides a broader competition-level mathematics testbed. These datasets assess accuracy across arithmetic difficulty, structural variation, and multi-step decomposition.

  • GSM8K-Aug expands GSM8K to approximately 385K training examples while using structured arithmetic-expression CoT annotations.
  • GSM-Hard replaces original problem values with larger numbers requiring more complex arithmetic operations.
  • SVAMP tests sensitivity to structural variations in elementary mathematics word problems.
  • MultiArith evaluates decomposition of multi-step arithmetic narratives requiring at least two mathematical operations.
  • MATH contains 12,500 competition-level problems across seven mathematical subjects with difficulty levels from 1 to 5.

A.2 BASELINES

The baselines cover explicit reasoning, curriculum-based latent reasoning, compression, superposition, distillation, and auxiliary step-level supervision. Together, they represent varied strategies for moving reasoning from explicit tokens into latent representations.

  • CoT-SFT fine-tunes models on Question, Chain-of-Thought, and Answer triples to generate explicit reasoning tokens.
  • Coconut progressively replaces explicit language reasoning steps with continuous thoughts through multi-stage curriculum learning.
  • CoLaR-2 compresses explicit reasoning using compressed-embedding prediction followed by reinforcement learning over diverse reasoning paths.
  • Latent-SFT models latent reasoning as a chain of superposition in which each latent token encodes multiple reasoning trajectories.
  • CODI jointly trains explicit and implicit CoT tasks while aligning hidden states at a designated token position.
  • Sim-CoT uses an auxiliary decoder to align each implicit token with its corresponding explicit reasoning step during training.

B IMPLEMENTATION DETAILS

The experiments use standardized optimizer and LoRA settings across baselines, with method-specific hyperparameters for each training approach. PMPS additionally uses prototype, temperature, assignment, and matching settings, while Figure 4 visualizes latent-token interpretations against CoT steps.

  • Implementation details: All experiments use AdamW with weight decay 0.1, warmup ratio 0.03, and LoRA fine-tuning for most compared methods.
  • Baseline configurations: CoT-SFT, Coconut, CoLaR-2, Latent-SFT, and CODI or Sim-CoT use distinct batch sizes, learning rates, training schedules, or distillation settings.
  • Interpretability analysis: Figure 4 decodes latent embeddings and aligns their top-five highest-probability tokens with corresponding CoT steps.
  • PMPS configuration: PMPS sets prototype dimension d to 128, matching temperature τm to 0.3, prediction temperature τ to 0.1, and batch-level Sinkhorn assignment for three iterations.

C INTERPRETABILITY EXPERIMENT

The interpretability experiment decodes each latent reasoning embedding into its top-5 vocabulary tokens to inspect implicitly encoded linguistic or symbolic content. Figure 5 examines how latent-token and prototype counts affect performance on GPT-2 with GSM8K-Aug against CODI.

  • Token-Level Decoding: Each latent embedding is projected through the language-model head, and its five highest-probability vocabulary tokens are extracted for token-level decoding analysis.This reveals linguistic or symbolic content implicitly encoded at each latent reasoning step.
  • Ablation Setup: Figure 5 studies the effects of varying latent-token and prototype counts in GPT-2 experiments on GSM8K-Aug, using CODI as the comparison baseline.

D ABLATION STUDY

The ablation shows that PMPS is generally more robust than CODI across latent-token counts, with six or seven latent tokens performing best and a 4× prototype-to-latent ratio yielding the strongest overall results. Too few prototypes limit semantic granularity, while too many introduce redundancy.

  • Ablation on Number of Latent Tokens: PMPS outperforms CODI at every tested latent-token count except 2, where limited capacity constrains prototype-based soft matching.The best PMPS performance occurs with 6 and 7 latent tokens.
  • Ablation on Number of Latent Tokens: At 8 latent tokens, CODI drops sharply while PMPS remains competitive, indicating more stable supervision as the number of latent representations increases.The passage attributes CODI’s decline to answer-level distillation being insufficient for supervising a larger latent set.
  • Ablation on Number of Prototypes: 4× is the best prototype-to-latent ratio across latent counts, whereas 2× and 3× provide insufficient semantic granularity and 6× introduces redundant prototypes.
Loading 2609.09928v1…