Structural Process Supervision for Latent Chain-of-Thought Reasoning

Chain-of-thought (CoT) prompting (Wei et al., 2022) improves LLM reasoning, but the inference cost increases substantially with verbose reasoning length, motivating reasoning in continuous hidden-state space. iCoT (Deng et al., 2024) and Coconut (Hao et al., 2024) progressively replace explicit CoT tokens with latent embeddings via curriculum learning. CODI (Shen et al., 2025) stabilizes training through teacher–student distillation at the answer position. SIM-CoT (Wei et al., 2025) enforces rigid one-to-one assignment between CoT steps and latent embeddings with an auxiliary decoder. CoLaR (Tan et al., 2025) and Latent-SFT (Deng et al., 2026) compress a contiguous sequence of CoT tokens into a single latent embedding. However, these methods either lack direct reasoning supervision or enforce an overly rigid correspondence between the CoT and latent embeddings, which limits the expressiveness of latent representations. We propose PMPS to address these limitations via th structural prototype-mediated soft alignment, requiring no auxiliary decoder and no hard CoT partition.

Process Reward Models (Lightman et al., 2024) show that step-level correctness labels yield better reasoning verifiers than outcome-only rewards, with follow-up work automating such supervision via Monte Carlo tree search (Wang et al., 2024). However, existing process supervision methods require explicit, discrete reasoning steps, which are incompatible with continuous latent embeddings. In parallel, SwAV (Caron et al., 2020) and PCL (Li et al., 2020), originally from visual self-supervised learning, show that learnable prototypes can effectively structure embedding spaces without negative pairs. Knowledge distillation approaches (Sun et al., 2019; Wang et al., 2020) align teacher–student representations but typically assume matched sequence lengths. Our work bridges these by leveraging prototypes to provide many-to-many structural process-level supervision across an unequal number of latent embeddings and CoT embeddings, complemented by answer-position distillation from the teacher to ensure outcome-level alignment.

To obtain high-quality explicit Chain-of-Thought (CoT) embeddings that align with the model’s intrinsic distribution as supervision signals, we adopt the widely used self-distillation framework (Zhang et al., 2019; Liao et al., 2023; Liu et al., 2026) as the foundational architecture. Within this framework, the model is jointly optimized for two tasks: a teacher task that performs explicit CoT reasoning, and a student task that executes the target latent reasoning paradigm. Inspired by CODI (Shen et al., 2025), we align the hidden states of teacher and student task at the answer token positions. The objective function for this base architecture is formulated as:

where $\mathcal{L}_{ce}^{stu}$ and $\mathcal{L}_{ce}^{tea}$ are cross-entropy losses for the student’s answer prediction and the teacher’s full-sequence prediction, and $\mathcal{L}_{dist}$ is a Smooth L1 loss aligning teacher and student hidden states at the answer token. However, relying solely on this base architecture lacks process-level supervision for the latent embeddings, which can lead to latent homogenization and representation collapse.

To mitigate the representation collapse induced by the absence of process supervision, we propose to leverage the CoT embeddings from the teacher task as supervision signals. These explicit CoT embeddings are well-suited targets for two reasons: (1) the $\mathcal{L}_{ce}^{tea}$ loss ensures that they encode high-quality, semantically meaningful reasoning information, and (2) since both the explicit CoT embeddings and latent embeddings are generated by the same model, they are inherently co-distributed, facilitating more effective knowledge transfer.

However, direct one-to-one distillation is infeasible because the number of latent embeddings $N_{l}$ is far smaller than the number of CoT tokens $N_{c}$. We therefore propose Prototype-Mediated Process Supervision (PMPS), which introduces $K$ learnable prototype vectors $\mathbf{P}=\{\mathbf{p}_{1},\ldots,\mathbf{p}_{K}\}\in\mathbb{R}^{K\times d}$ as shared semantic anchors to mediate a soft many-to-many alignment between the two sets of representations, imposing structured semantic organization on the latent space. Let $\mathbf{H}_{l}=\{\mathbf{l}_{1},\ldots,\mathbf{l}_{N_{l}}\}\in\mathbb{R}^{N_{l}\times D}$ and $\mathbf{H}_{c}=\{\mathbf{c}_{1},\ldots,\mathbf{c}_{N_{c}}\}\in\mathbb{R}^{N_{c}\times D}$ denote the latent embeddings and the explicit CoT embeddings, respectively, where $D$ is the hidden dimension and $d$ is the prototype dimension.

Projection and Scoring. A shared two-layer MLP with batch normalization serves as the projection head $g(\cdot)$, projecting both $\mathbf{H}_{l}$ and $\mathbf{H}_{c}$ from the original hidden space into the prototype space, followed by $\ell_{2}$ normalization:

where $\mathbf{h}$ denotes either the latent embedding or the explicit CoT embedding. Each unit feature vector $\mathbf{z}$ is scored against the $\ell_{2}$-normalized prototypes via inner product $s_{k}=\mathbf{z}^{\top}\mathbf{p}_{k}$, yielding score matrices $\mathbf{S}_{l}\in\mathbb{R}^{N_{l}\times K}$ and $\mathbf{S}_{c}\in\mathbb{R}^{N_{c}\times K}$.

Prototype Assignment. To obtain soft assignments while preventing prototype collapse, we apply the Sinkhorn-Knopp algorithm (Cuturi, 2013) to enforce an equipartition constraint. Given a score matrix $\mathbf{S}\in\mathbb{R}^{M\times K}$, where $M$ denotes the number of tokens, we initialize $\mathbf{Q}^{(0)}\propto\exp(\mathbf{S}/\epsilon)$, with $\epsilon$ serving as the regularization parameter. We then iteratively alternate row and column normalization:

The converged $\mathbf{Q}\in\mathbb{R}^{M\times K}$ satisfies that each row sums to $1$ (a valid per-token distribution) and each column sums to $M/K$ (uniform prototype utilization). This procedure is applied independently to both $\mathbf{S}_{l}$ and $\mathbf{S}_{c}$ with stop-gradient operation following SwAV (Caron et al., 2020), producing the final soft assignment codes $\mathbf{Q}_{l}\in\mathbb{R}^{N_{l}\times K}$ and $\mathbf{Q}_{c}\in\mathbb{R}^{N_{c}\times K}$.

Many-to-Many Matching. Leveraging the prototype assignments as a bridge, we compute the soft matching matrix between latent embeddings and explicit CoT embeddings:

where temperature $\tau_{m}$ controls matching sharpness. We then apply softmax normalization along the last dimensions to obtain bidirectional matching weights:

This formulation naturally induces a many-to-many alignment: a single latent embedding can attend to multiple explicit CoT embeddings for compression, while a single explicit CoT embedding can contribute to multiple latent embeddings for semantic sharing.

Bidirectional Cross-Prediction Loss. Alongside the stop-gradient assignment codes $\mathbf{Q}$, we compute gradient-carrying prototype predictions $\bm{\phi}_{l}=\text{softmax}(\mathbf{S}_{l}/\tau)$ and $\bm{\phi}_{c}=\text{softmax}(\mathbf{S}_{c}/\tau)$, where $\tau$ is a prediction temperature. For each explicit CoT embedding $j$, we aggregate the matched latent predictions weighted by $\mathbf{A}_{c2l}$ and minimize the cross-entropy with its assignment code:

where $Q_{c}^{(j,k)}$ is the $(j,k)$-th entry of $\mathbf{Q}_{c}$, $A_{c2l}^{(j,i)}$ is the matching weight from the $j$-th explicit CoT embedding to the $i$-th latent embedding, and $\phi_{l,i}^{(k)}$ is the predicted probability of the $i$-th latent embedding for the $k$-th prototype. Symmetrically, the latent-to-CoT direction yields:

The overall process supervision loss is $\mathcal{L}_{proto}=\frac{1}{2}(\mathcal{L}_{c2l}+\mathcal{L}_{l2c})$. The CoT-to-latent direction ($\mathcal{L}_{c2l}$) ensures that the latent embeddings collectively explain each explicit CoT embedding, while the latent-to-CoT direction ($\mathcal{L}_{l2c}$) encourages each latent embedding to capture diverse CoT semantics.

Explicit CoT inherently follows a progressive process: early tokens build up the reasoning chain, while final tokens converge toward the answer. However, the unconstrained matching in the base PMPS allows any latent token to align with any CoT position, which risks failing to capture the progressive reasoning structure. Conversely, enforcing a rigid segmented mapping (e.g., uniformly partitioning the CoT into $N_{l}$ segments for strict assignment) is overly restrictive: uniform partitioning introduces noisy supervision that undermines the flexibility of semantic-driven matching. To navigate between these two extremes, we introduce Progressive Sequential Alignment (PSA), which injects sequential structure into the matching through a soft Gaussian positional prior and gradually relaxes it via curriculum-scheduled decay.

Concretely, we augment $\mathbf{M}_{c2l}$ with a bias penalizing positionally distant matches:

where $j/(N_{c}{-}1)$ and $i/(N_{l}{-}1)$ are normalized positions in $[0,1]$, and $\alpha\geq 0$ controls the enforcement strength.

To balance structural guidance with representational flexibility, we adopt cosine annealing to decay $\alpha$ from $\alpha_{\max}$ to $0$ over training:

where $t$ and $T$ are the current and total training steps. Early in training ($\alpha$ large), the strong prior helps latent embeddings establish sequential structure and stabilizes matching. Late in training ($\alpha\to 0$), the constraint is fully released, allowing adaptive alignment where reasoning steps may span or consolidate across latent tokens. PSA modifies only the matching scores before softmax, introducing no learnable parameters and incurring zero inference overhead.

The overall training objective combines the base loss with our proposed process supervision loss:

where $\lambda_{proto}$ is the weighting coefficient for the process supervision loss.

The teacher task, prototype layer, projection head, and PSA module are used exclusively during training. At inference time, the model architecture and computational cost are identical to the single student task. This design ensures that all accuracy improvements from process supervision come at zero additional inference cost.

Table 1 presents the main results on one in-domain benchmark (GSM8K-Aug) and three out-of-domain benchmarks (GSM-Hard, SVAMP, MultiArith), with three backbone models from different model families: GPT-2, Qwen2.5-0.5B-Instruct, and LLaMA-3.2-1B-Instruct.

To further validate the scalability of PMPS, we conduct two additional groups of experiments: scaling to a larger backbone model and training on a more challenging dataset.

To further investigate the reasons behind PMPS’s superior performance, we directly analyze the pairwise cosine similarity among the $N_{l}{=}6$ latent embeddings on randomly sampled 100 GSM8K-Aug test instances using GPT-2. The analysis reveals that CODI and SIM-CoT suffer from distinct patterns of representation collapse, whereas PMPS maintains well-diversified representations free from either collapse mode. Additionally, we conduct the latent embedding interpretability experiments, which are provided in Appendix C.

In this section, we conduct ablation studies to investigate the effects of different PMPS components, the maximum enforcement strength of PSA, and the number of latent embeddings and prototypes. Results on the number of latent embeddings and prototypes are provided in Appendix D.