Method
3. 方法
3.1 问题定义与概述
给定一个推理问题 (x)，传统的思维链（Chain-of-Thought, CoT）方法要求模型显式生成一系列中间推理步骤 (c_1, c_2, \dots, c_T)，然后基于这些步骤生成最终答案 (y)。然而，显式生成推理链会引入显著的推理时延和计算开销，严重限制了大语言模型在实际场景中的部署效率。
本文提出 CODI-PMPS（Compressed Distilled reasoning with Swapped Assignments between Views），一种基于自监督簇对齐的隐式推理框架。其核心思想是：将显式思维链的推理过程压缩到少量隐向量的迭代变换中，并通过一种新颖的PMPS风格损失函数在共享原型空间中对隐式推理步骤施加结构化监督。与以往方法仅在答案位置进行特征蒸馏不同，本方法通过原型簇对齐机制实现了隐式推理过程与显式思维链之间细粒度的语义对应。
3.2 模型架构
CODI-PMPS 的整体架构如图 X 所示，由以下核心模块组成：一个共享的因果语言模型骨干网络 (f_\theta)（如 GPT-2 或 LLaMA），一个可选的投影层 (g_\phi)，以及一个PMPS监督模块 (\mathcal{S}_\psi)。训练过程中，模型同时执行两条路径：
- 学生路径（隐式推理）：将问题编码为初始隐向量后，通过多轮迭代在隐空间中进行推理，最终从隐向量解码出答案。
- 教师路径（显式推理）：将问题、完整思维链和答案拼接为一个序列，进行标准的因果语言建模。
两条路径共享同一个语言模型骨干 (f_\theta)，教师路径的隐藏状态作为PMPS损失的监督信号，指导学生路径中隐向量的学习。
3.3 隐式推理机制
隐式推理的核心在于用少量隐向量的迭代变换替代冗长的文本推理链。具体过程分为三个阶段：
问题编码。 给定输入问题 (x = (x_1, x_2, \dots, x_n))，我们首先将其通过语言模型 (f_\theta) 进行编码：
[ \mathbf{h}0 = f_\theta(x_1, \dots, x_n){[:, -1, :]} ]
其中 (\mathbf{h}_0 \in \mathbb{R}^{d}) 为最后一个token位置的隐藏状态，作为初始隐向量。同时，我们保留编码过程中的KV缓存 (\mathcal{K}_0) 以加速后续计算。
迭代隐式推理。 初始隐向量 (\mathbf{h}_0) 经过可选的投影层变换后，被送入语言模型进行 (N_l) 轮迭代推理。在第 (i) 轮迭代中（(i = 1, \dots, N_l)）：
[ \mathbf{h}i = g_\phi\left(f_\theta\left(\mathbf{h}{i-1}; \mathcal{K}{i-1}\right){[:, -1, :]}\right) ]
其中 (\mathcal{K}_{i-1}) 是从前 (i-1) 步累积的KV缓存，(g_\phi) 为投影层（包含Dropout、线性变换、GELU激活和可选的LayerNorm）。每一步隐向量 (\mathbf{h}_i) 作为一个"隐式推理token"输入模型，使得模型能够在连续空间中进行多步推理，而无需将中间状态解码为离散文本。
答案解码。 经过 (N_l) 步隐式推理后，我们将特殊标记 (\langle\text{eot}\rangle)（end-of-thought）的嵌入拼接至累积KV缓存之后，然后通过自回归方式生成最终答案 (y)。
与显式思维链相比，隐式推理具有以下优势：（1）推理步数 (N_l)（如6步）远少于思维链token数量（通常数十至数百），大幅降低了推理延迟；（2）每个隐向量可以在连续空间中编码比单个离散token更丰富的语义信息，实现了推理过程的有效压缩。
3.4 PMPS风格的隐式推理监督
隐式推理面临的核心挑战在于：如何在缺乏显式中间监督信号的情况下，引导隐向量学习到有意义的推理能力？现有方法通常仅在答案位置对齐学生和教师的隐藏状态，这种"端到端"的监督方式无法约束中间推理步骤的语义，容易导致隐向量退化为冗余的重复表示。
为此，我们提出一种受PMPS（Swapping Assignments between Views）启发的结构化监督方法，将隐式推理的隐藏状态与教师思维链的隐藏状态在共享的原型簇空间中进行对齐。
3.4.1 投影与原型表示
设隐式推理步骤的隐藏状态集合为 (\mathbf{H}_\text{lat} = [\mathbf{h}1, \dots, \mathbf{h}{N_l}] \in \mathbb{R}^{N_l \times d})，教师思维链token对应的隐藏状态为 (\mathbf{H}_\text{cot} = [\mathbf{c}1, \dots, \mathbf{c}{N_c}] \in \mathbb{R}^{N_c \times d})。
我们引入一个共享的MLP投影头 (\pi: \mathbb{R}^d \to \mathbb{R}^{d'})，以及一组可学习的原型向量 (\mathbf{P} = [\mathbf{p}_1, \dots, \mathbf{p}_K] \in \mathbb{R}^{K \times d'})。投影头将高维隐藏状态映射至低维空间，原型向量定义了一组共享的"推理概念簇心"。
对于每个隐向量或思维链向量 (\mathbf{v})，其原型得分计算为：
[ \mathbf{s}(\mathbf{v}) = \mathbf{P} \cdot \bar{\pi}(\mathbf{v})^\top \in \mathbb{R}^K ]
其中 (\bar{\pi}(\mathbf{v}) = \pi(\mathbf{v}) / |\pi(\mathbf{v})|_2) 表示L2归一化后的投影向量。
3.4.2 Sinkhorn-Knopp 簇分配
遵循PMPS的设计，我们通过Sinkhorn-Knopp算法为每个token计算软簇分配码（code），而非直接使用softmax。Sinkhorn-Knopp算法在保持每个token获得有效簇分配的同时，强制所有token在原型簇上的分布满足均分约束（equipartition constraint），确保所有 (K) 个原型被均匀利用，从而有效防止原型坍缩问题。
具体而言，对于原型得分矩阵 (\mathbf{S} \in \mathbb{R}^{M \times K})（(M) 为token数量），初始化 (\mathbf{Q}^{(0)} = \exp(\mathbf{S} / \epsilon)^\top) 并执行以下迭代归一化：
[ \mathbf{Q}^{(t+1)} = \text{ColNorm}\left(\text{RowNorm}\left(\mathbf{Q}^{(t)}\right)\right) ]
其中 (\epsilon) 为正则化参数，经过 (T_\text{sk}) 次迭代后得到最终的分配码矩阵 (\mathbf{Q})。每一行 (\mathbf{q}_i \in \Delta^{K-1}) 表示第 (i) 个token在 (K) 个原型上的软分配概率。
我们分别对隐式推理步骤和思维链token计算分配码：
[ \mathbf{Q}_\text{lat} = \text{Sinkhorn}(\mathbf{S}_\text{lat}) \in \mathbb{R}^{N_l \times K}, \quad \mathbf{Q}_\text{cot} = \text{Sinkhorn}(\mathbf{S}_\text{cot}) \in \mathbb{R}^{N_c \times K} ]
3.4.3 簇空间跨视图匹配
PMPS损失的核心思想是"交换预测"：用一个视图的分配码作为目标，训练另一个视图的原型预测去匹配。然而，隐式推理与思维链之间存在严重的长度不对称（(N_l \ll N_c)），无法直接进行逐token对齐。
为此，我们提出簇空间软匹配（Cluster-space Soft Matching）机制，通过分配码的相似性自动发现隐向量与思维链token之间的对应关系。具体而言，我们计算跨视图匹配矩阵：
[ \mathbf{M}_{ij} = \frac{\mathbf{q}_i^\text{cot} \cdot (\mathbf{q}_j^\text{lat})^\top}{\tau_m} ]
其中 (\tau_m) 为匹配温度参数。随后通过softmax归一化得到软分配矩阵：
[ \mathbf{A}_\text{c2l} = \text{softmax}(\mathbf{M}, \text{dim}=N_l) \in \mathbb{R}^{N_c \times N_l} ]
(\mathbf{A}_\text{c2l}[j, i]) 表示第 (j) 个思维链token被分配给第 (i) 个隐向量的概率。类似地，(\mathbf{A}_\text{l2c} = \text{softmax}(\mathbf{M}^\top, \text{dim}=N_c)) 表示反向分配。
这一设计允许一个隐向量匹配多个思维链token，自然地实现了推理信息的压缩；同时不同的隐向量倾向于匹配不同的思维链片段，实现了推理步骤的功能分工。
3.4.4 渐进式序列位置偏置
在自由匹配的基础上，我们进一步引入高斯位置偏置（Gaussian Positional Bias），鼓励隐式推理步骤与思维链在序列位置上保持大致的对应关系。其直觉在于：第一个隐向量应主要关注思维链的前部信息，最后一个隐向量应主要关注思维链的后部信息。
对于第 (j) 个思维链token和第 (i) 个隐向量，位置偏置定义为：
[ \text{bias}(j, i) = -\lambda \left(\frac{j}{N_c - 1} - \frac{i}{N_l - 1}\right)^2 ]
其中 (\lambda) 控制偏置强度。该偏置被加入匹配矩阵 (\mathbf{M}) 中：
[ \mathbf{M}'{ji} = \mathbf{M}{ji} + \text{bias}(j, i) ]
偏置强度 (\lambda) 支持多种退火策略（余弦退火、线性衰减或恒定），实现从严格序列对应到自由匹配的渐进过渡（课程学习）。在训练初期，较强的位置偏置引导模型快速建立粗粒度的序列对应结构；随着训练推进，偏置逐步减弱，允许模型学习更灵活的跨位置匹配模式。
3.4.5 PMPS损失计算
基于软分配矩阵，我们构造混合原型预测，并计算双向交叉熵损失。
从思维链到隐向量方向（C2L）：对于每个思维链token (j)，其混合预测为分配给它的隐向量原型分布的加权平均：
[ \tilde{\mathbf{p}}j^\text{lat} = \sum{i=1}^{N_l} \mathbf{A}_\text{c2l}[j, i] \cdot \mathbf{p}_i^\text{lat} ]
其中 (\mathbf{p}_i^\text{lat} = \text{softmax}(\mathbf{s}(\mathbf{h}_i) / \tau)) 为第 (i) 个隐向量的原型概率分布。损失为：
[ \mathcal{L}\text{c2l} = -\frac{1}{N_c} \sum{j=1}^{N_c} \sum_{k=1}^{K} \mathbf{q}{jk}^\text{cot} \log \tilde{p}{jk}^\text{lat} ]
从隐向量到思维链方向（L2C）：对称地计算反向损失 (\mathcal{L}_\text{l2c})。最终PMPS损失为：
[ \mathcal{L}_\text{PMPS} = \frac{1}{2} \left(\mathcal{L}_\text{c2l} + \mathcal{L}_\text{l2c}\right) ]
3.5 训练目标
CODI-PMPS的总训练目标由四个损失项组成：
[ \mathcal{L} = \mathcal{L}_\text{CE}^\text{student} + \alpha \cdot \mathcal{L}_\text{distill} + \beta \cdot \mathcal{L}_\text{CE}^\text{teacher} + \gamma \cdot \mathcal{L}_\text{PMPS} ]
其中：
学生交叉熵损失 (\mathcal{L}_\text{CE}^\text{student})：隐式推理后生成答案的标准语言建模损失，确保模型能够从隐向量中正确解码答案。
蒸馏损失 (\mathcal{L}_\text{distill})：在答案起始位置对齐学生路径和教师路径的隐藏状态（使用Smooth L1损失），提供答案级别的跨路径监督。
教师交叉熵损失 (\mathcal{L}_\text{CE}^\text{teacher})：对完整思维链序列的标准语言建模损失，维持共享骨干网络对显式推理的建模能力。
PMPS对齐损失 (\mathcal{L}_\text{PMPS})：如3.4节所述，在原型簇空间中对隐式推理步骤与思维链token进行结构化对齐。
训练过程中，我们采用LoRA进行参数高效微调，仅更新语言模型的低秩适配器参数、投影层参数以及PMPS模块参数（投影头和原型向量），从而在保持预训练知识的同时高效学习隐式推理能力。
3.6 推理过程
推理时，模型仅执行学生路径，无需教师路径和PMPS模块，因此不引入额外的计算开销。给定问题 (x)，模型依次执行：（1）问题编码得到初始隐向量；（2）迭代 (N_l) 步隐式推理；（3）自回归生成最终答案。整个推理过程仅需 (N_l) 个额外的前向传播步骤（每步处理一个隐向量token），相比于显式思维链需要生成并处理数十至数百个中间token，推理效率显著提升。