LoRi:面向隐式推理的低秩蒸馏框架
摘要
LoRi提出了一种面向隐式思维链推理的低秩蒸馏框架,该框架在共享低秩子空间中对齐教师和学生轨迹,从而提升数学推理基准上的性能。
arXiv:2606.05315v1 公告类型:新
摘要:隐式思维链(iCoT)方法旨在将推理过程内化于大型语言模型中,但其表现往往逊于显式思维链提示。我们通过实验发现,隐状态推理轨迹呈现出低秩结构。基于这一观察,我们提出了一种低秩蒸馏框架,通过一阶和二阶统计量在共享低秩张量子空间中对齐教师和学生轨迹,从而实现推理迁移。该公式能够捕捉推理的全局结构,同时支持紧凑的潜在推理过程。我们在多个模型家族(包括LLaMA和Qwen)上,针对不同规模的数学推理基准评估了该方法。我们的方法持续提升了性能,尤其在具有挑战性的多步任务上,接近显式思维链的准确率,并优于先前的隐式思维链蒸馏方法。
查看缓存全文
缓存时间: 2026/06/05 08:05
# LoRi: 用于隐式推理的低秩蒸馏(Low-Rank Distillation for Implicit Reasoning) 来源:https://arxiv.org/html/2606.05315 Ryan Solgi¹、田野天一¹、郑璋¹ ¹加州大学圣塔芭芭拉分校,美国 [email protected], [email protected] ###### 摘要 隐式思维链(iCoT)方法旨在将推理过程内化到大语言模型中,但其效果通常不如显式思维链提示。我们通过实验发现,隐藏状态的推理轨迹具有低秩结构。受此观察启发,我们提出了一种低秩蒸馏框架,该框架通过一阶和二阶统计量,在共享的低秩张量子空间中对齐教师和学生轨迹,从而实现推理能力的迁移。由此产生的公式能够捕捉推理的全局结构,同时支持紧凑的潜在推理过程。我们在多个模型家族(包括LLaMA和Qwen)上,针对不同规模和数学推理基准进行了评估。我们的方法始终如一地提升了性能,特别是在具有挑战性的多步任务上,其准确率接近显式思维链,并优于先前的iCoT蒸馏方法。代码可在 https://github.com/rmsolgi/lori 获取。 LoRi: 用于隐式推理的低秩蒸馏(Low-Rank Distillation for Implicit Reasoning) Ryan Solgi¹、田野天一¹、郑璋¹ ¹加州大学圣塔芭芭拉分校,美国 [email protected], [email protected] ## 1 引言 大语言模型(LLMs)在显式思维链(CoT)提示下展现出强大的推理能力(Wei et al., 2022 (https://arxiv.org/html/2606.05315#bib.bib18);Zelikman et al., 2022 (https://arxiv.org/html/2606.05315#bib.bib20))。然而,显式CoT计算开销大,可能助长对非鲁棒模式的依赖,并且对推理过程敏感(Li et al., 2024 (https://arxiv.org/html/2606.05315#bib.bib19);Lin et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib21);Wang et al., 2022 (https://arxiv.org/html/2606.05315#bib.bib22);Yao et al., 2023 (https://arxiv.org/html/2606.05315#bib.bib24))。因此,近期工作探索了更高效的推理范式,以减少对显式文本理由的依赖(He et al., 2026 (https://arxiv.org/html/2606.05315#bib.bib23);Li et al., 2025b (https://arxiv.org/html/2606.05315#bib.bib25))。 隐式思维链(iCoT)将推理过程编码在潜在表示中,而非显式文本中(Deng et al., 2024a (https://arxiv.org/html/2606.05315#bib.bib26);Hao et al., 2024 (https://arxiv.org/html/2606.05315#bib.bib27))。近期的iCoT蒸馏方法将显式理由迁移到潜在推理状态中(Shen et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib6);Wu et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib28);Kuzina et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib7);Wei et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib8)),但在困难数学推理任务上的表现仍逊于显式CoT。一个关键挑战在于,显式推理步骤与潜在动态之间缺乏明确的对应关系,使得“token到潜在”的迁移本质上是不适定的。现有方法依赖于局部监督(Shen et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib6))或采样的中间状态(Kuzina et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib7)),这可能无法完全捕捉全局推理轨迹。 见标题说明 (a) LoRi 的高层概述。 见标题说明 (b) GSM8K-Hard 上的代表性推理结果。 图 1:所提出的低秩 iCoT 蒸馏框架及其推理性能概览。(a) 从教师隐藏状态学习到的低秩因子用于将教师和学生表示投影到一个共享的低秩子空间中。然后通过损失函数匹配投影后的低秩表示。(b) 在GSM8K-Hard上,LoRi在不同模型规模下均提升了相较于先前隐式CoT基线的推理准确率。先前的工作表明,模型表示具有低维几何结构(Yu and Wu, 2023 (https://arxiv.org/html/2606.05315#bib.bib9);Golowich et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib10);Model et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib13);Park et al., 2024 (https://arxiv.org/html/2606.05315#bib.bib11))。虽然近期研究通过轨迹几何和可分性来分析推理过程(Sun et al., 2026 (https://arxiv.org/html/2606.05315#bib.bib12);Zhou et al., 2026 (https://arxiv.org/html/2606.05315#bib.bib17);Li et al., 2025a (https://arxiv.org/html/2606.05315#bib.bib16)),但它们并未直接检验跨token和跨层的隐藏状态的低秩结构。通过堆叠跨层和CoT token的隐藏状态,我们发现归一化累积奇异值随秩快速增长(附录A (https://arxiv.org/html/2606.05315#A1)),表明推理轨迹可以由低维子空间很好地近似。 受此观察启发,我们提出了一种低秩 iCoT 蒸馏框架,该框架通过低秩统计表示来对齐教师推理轨迹的全局几何结构。学生无需模仿显式推理步骤,而是学习教师推理动态背后的低维子空间(图 [1 (https://arxiv.org/html/2606.05315#S1.F1)(a))。这使得学生能够用较短的潜在轨迹捕捉主要的推理结构,从而实现长度不变且高效的蒸馏。 我们的主要贡献总结如下: - **低秩 iCoT 蒸馏。** 我们提出了一种 iCoT 蒸馏框架,通过使用一阶和二阶隐藏状态统计量,在共享的低秩子空间中对齐教师和学生轨迹,从而实现推理能力的迁移。 - **高效、长度不变的全局推理迁移。** 该公式将长 CoT 推理迁移到与序列长度无关的短潜在轨迹中,从而用少量潜在步骤实现高效推理,并且无需额外的中间监督。 - **跨模型和基准的一致提升。** 在不同模型和规模上,LoRi 始终优于先前的 iCoT 方法,准确率提升高达 ∼\sim12%,并在 GSM8K-Hard 上获得强增益(高达 ∼\sim10%),如图 [1](https://arxiv.org/html/2606.05315#S1.F1)(b) 所示。在更大规模下,LoRi 显著缩小了与显式 CoT 之间的差距。 ## 2 背景与相关工作 #### iCOT 蒸馏。 在推理中,模型根据输入 x\\boldsymbol{x},在中间推理过程 τ\\boldsymbol{\\tau} 的条件下预测最终答案 y\\boldsymbol{y}: p(y,τ∣x)=p(τ∣x)p(y∣x,τ),p(\\boldsymbol{y},\\boldsymbol{\\tau}\\mid\\boldsymbol{x}) = p(\\boldsymbol{\\tau}\\mid\\boldsymbol{x})\\,p(\\boldsymbol{y}\\mid\\boldsymbol{x},\\boldsymbol{\\tau}), 其中 τ\\boldsymbol{\\tau} 可能对应显式文本理由(显式 CoT)或隐式内部推理过程(iCoT)。遵循先前工作(Shen et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib6);Kuzina et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib7);Wei et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib8)),我们考虑一个师生蒸馏框架,将推理能力从 CoT 迁移到 iCoT。 - **教师模型。** 教师生成一个自然语言推理轨迹 r=(r1,⋯,rN)\\boldsymbol{r}=(r_{1},\\cdots,r_{N}),并定义联合条件分布 pT(y,r∣x)=pT(r∣x)pT(y∣x,r)p_{T}(\\boldsymbol{y},\\boldsymbol{r}\\mid\\boldsymbol{x}) = p_{T}(\\boldsymbol{r}\\mid\\boldsymbol{x})\\,p_{T}(\\boldsymbol{y}\\mid\\boldsymbol{x},\\boldsymbol{r})。 - **学生模型。** 学生构建一个潜在推理轨迹 {zt}t=1K\\{\\boldsymbol{z}_{t}\\}_{t=1}^{K},其中 K≪NK \\ll N,每个 zt\\boldsymbol{z}_{t} 代表一个隐藏推理状态。该轨迹递归生成: zt=fθ(x,z1:t−1),t=1,...,K,\\boldsymbol{z}_{t} = f_{\\boldsymbol{\\theta}}(\\boldsymbol{x},\\boldsymbol{z}_{1:t-1})\\,,\\qquad t=1,\\ldots,K, 其中 fθf_{\\boldsymbol{\\theta}} 是学生模型,z1\\boldsymbol{z}_{1} 从 x\\boldsymbol{x} 初始化。最终答案自回归生成: pS(y∣x,z1,...,zK)p_{S}(\\boldsymbol{y}\\mid\\boldsymbol{x},\\boldsymbol{z}_{1},\\ldots,\\boldsymbol{z}_{K})。 - **训练。** 学生通过如下形式的目标函数训练: L=Lreason+λLtask,\\mathcal{L} = \\mathcal{L}_{\\mathrm{reason}} + \\lambda \\mathcal{L}_{\\mathrm{task}}, 其中 Lreason\\mathcal{L}_{\\mathrm{reason}} 将推理行为从显式 CoT 教师迁移到潜在推理动态中,Ltask\\mathcal{L}_{\\mathrm{task}} 监督答案预测,λ\\lambda 平衡两个目标。 尽管 iCoT 提高了推理效率,但常常降低推理准确率。现有方法主要区别在于如何将推理从显式 CoT 迁移到 iCoT。逐步内化(Deng et al., 2024a (https://arxiv.org/html/2606.05315#bib.bib26))通过迭代微调逐渐去除 CoT token,而 COCONUT(Hao et al., 2024 (https://arxiv.org/html/2606.05315#bib.bib27))用隐藏状态动态替换文本推理。近期方法依赖蒸馏(Deng et al., 2024b (https://arxiv.org/html/2606.05315#bib.bib32)):CODI(Shen et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib6))在答案边界处蒸馏推理,PCCoT(Wu et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib28))为中间推理状态添加 token,KAVA(Kuzina et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib7))对齐 KV 缓存动态,SIM-CoT(Wei et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib8))使用步骤级监督来引导推理轨迹。 #### 低维结构。 流形假设指出,高维数据通常位于一个捕捉其内在结构的低维流形上(Bengio et al., 2013 (https://arxiv.org/html/2606.05315#bib.bib15);Fefferman et al., 2016 (https://arxiv.org/html/2606.05315#bib.bib14);Chen et al., 2022 (https://arxiv.org/html/2606.05315#bib.bib3))。近期工作表明,LLM 表示也具有类似的低维几何结构,包括激活的近似低秩结构(Yu and Wu, 2023 (https://arxiv.org/html/2606.05315#bib.bib9);Chen et al., 2024 (https://arxiv.org/html/2606.05315#bib.bib4);Liu et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib5))、线性语义方向(Park et al., 2024 (https://arxiv.org/html/2606.05315#bib.bib11))、表示流形结构(Model et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib13))以及输出 logits 中的低秩行为(Golowich et al., 2025 (https://arxiv.org/html/2606.05315#bib.bib10))。 #### 推理几何。 在推理任务中,Sun et al. (2026) (https://arxiv.org/html/2606.05315#bib.bib12) 表明推理轨迹会经过特定于步骤的子空间,这些子空间随深度增加而变得更加可分。Zhou et al. (2026) (https://arxiv.org/html/2606.05315#bib.bib17) 将推理描述为由逻辑结构塑造的平滑流动,而 Li et al. (2025a) (https://arxiv.org/html/2606.05315#bib.bib16) 提出了一个推理流形,其中正确轨迹集中在低维区域,而错误则由偏离引起。这些研究共同表明,高维推理轨迹可以被低维子空间或流形有效近似。 ## 3 LoRi 方法 我们将教师的推理过程视为隐藏状态空间中的一条轨迹,该轨迹位于低维子空间附近。从这个角度来看,蒸馏不必强制执行逐步骤的对应;相反,它应该将学生的潜在推理动态与教师轨迹的主导结构对齐。受此观点启发,我们提出了 LoRi(低秩 iCoT,Low-Rank iCoT),这是一个将推理能力从显式思维链迁移到紧凑隐式潜在过程的蒸馏框架。LoRi 结合了两个互补的目标:理由级对齐,它保留了教师轨迹的全局几何结构;以及锚点级对齐,它规范了从潜在推理到答案生成的过渡。结合基于预计算低秩表示的高效训练方案,LoRi 实现了从 CoT 教师到 iCoT 学生的可扩展、长度不变的蒸馏。 见标题说明 图 2:LoRi 中低秩理由级对齐的概述。 ### 3.1 低秩 iCoT 蒸馏 #### 蒸馏损失。 蒸馏的目标是将教师显式理由 r 中编码的推理结构迁移到学生的潜在推理动态 z 中,同时保留学生生成显式逐步解和最终答案的能力。我们用复合目标训练学生: L = L_LR + λ L_CE, 其中 L_CE 是监督学生显式输出序列的交叉熵,L_LR 在低秩子空间中将学生的潜在推理状态与教师的隐藏表示对齐,λ 平衡这两项。直观地说,L_LR 鼓励两个模型共享教师推理轨迹的主导低秩结构,从而提高学生的稳定性和泛化能力。 #### 低秩项。 我们将 L_LR 定义为两个互补成分的和: L_LR = L_rationale + β L_anchor, 其中 L_rationale 将学生的潜在推理动态与教师在整个推理轨迹上的隐藏状态对齐,而 L_anchor 将学生在答案预测位置的表示与教师对应的隐藏状态对齐,β 控制其相对权重。前者捕捉教师推理过程的全局低秩结构,后者提供局部信号,引导从内部推理到答案生成的过渡。 ### 3.2 理由级对齐(L_rationale) 为了迁移推理能力,我们在低秩子空间中将学生的隐藏状态与教师的推理轨迹对齐。令 H_T ∈ R^(N×L×H) 表示教师关于理由 token 的隐藏状态,其中 L 是层数,H 是隐藏维度。类似地,令 H_S ∈ R^(K×L×H) 表示学生在其隐式推理步骤中的隐藏状态。我们不直接匹配这些高维张量,而是将它们投影到一个共享的低秩子空间,并通过一个统计匹配目标来对齐它们的低维表示(图 2 (https://arxiv.org/html/2606.05315#S3.F2))。 #### 推理轨迹的 Tucker 表示。 根据 Tucker 分解 (De Lathauwer et al., 2000) (https://arxiv.org/html/2606.05315#bib.bib36),一个张量 X ∈ R^(N×L×H) 被分解为一个张量核 Q ∈ R^(r_N×r_L×r_H) 和正交因子矩阵 U_N ∈ R^(N×r_N)、U_L ∈ R^(L×r_L)、
相似文章
通过混合层蒸馏和关键信息的逐步注意力改进小模型的推理能力
本文提出一种新颖的思维链蒸馏框架,通过混合层模块的动态层对齐,将教师模型对关键信息的逐步注意力转移到学生模型中。该方法通过明确指导学生模型在推理过程中逐步聚焦关键信息,在数学和常识推理基准测试中实现了一致的性能提升。
LARK:基于可学习性的轨迹选择方法用于高效推理蒸馏
LARK提出了一种基于可学习性的推理轨迹选择方法,用于大语言模型蒸馏。该方法采用可学习性因子和χ²正则化选择策略,平衡效率与泛化能力,在多个模型和任务上持续优于基线方法。
通过近未来引导弥合在线蒸馏中的推理轨迹
本文指出了在线蒸馏大语言模型时token级监督的局限性,并提出TOPD方法,利用近未来轨迹信息更好地识别发散推理状态并将引导分布到多个token上,在AIME基准测试中取得了性能提升。
偏离时回溯:缓解大语言模型推理蒸馏中的双重暴露偏差
本文介绍了一种名为Motab的新型大语言模型推理蒸馏流水线,通过动态监控学生生成并在偏离时回溯到安全状态并借助教师干预,同时缓解离策略和在线策略暴露偏差,取得了约3%的平均性能提升。
OPRD:在策略表示蒸馏
OPRD提出了一种新的知识蒸馏方法,该方法在策略部署期间跨层对齐学生和教师的隐藏状态,消除了来自词空间KL估计的采样方差。实验表明,OPRD在数学推理基准(AIME 2024/2025、AIMO)上优于输出空间基线,同时速度快1.44倍,内存使用减少54%。