理解 On-Policy Distillation:基于 Sparse Crosscoders 的机制可解释性视角

TL;DR 用稀疏交叉编码器与 swap readout 揭示 on-policy distillation 的真实机制：学生并未获得教师新特征，而是重新加权师生共享的已有特征，SFT 预热则提前部分完成该重加权。

问题

问题背景

LLM 推理能力后训练中，On-policy Distillation (OPD) 是主流技术，通常用强教师模型在 student rollout 上蒸馏，提升推理精度。但该过程如何改变 student 内部表征，业界缺乏认知。

现有方法局限

• 现有 interpretability 工具如 sparse crosscoders 可以同时编码多个模型到共享特征字典，识别模型特有特征。
• 局限：这类编码器将所有模型映射到同一组特征激活，无法区分某个 student checkpoint 自身对特征的使用强度变化；它只能告诉“哪些特征是某模型特定的”，不能回答“训练前后 student 对已有特征的依赖权重如何变化”。
• 因此无法解释 OPD 到底蒸馏了什么：是新增特征？复用教师特征？还是仅调节共享特征？

为什么这个问题难/重要

• 技术挑战：需要跨 checkpoint 追踪同一特征的使用频率，同时区分 student 与 teacher 的表征贡献；普通 crosscoder 丢失了模型内部的时序信息。
• 业界关注度：OPD 在 reasoning LLM 后训练中被广泛采用，但缺少机制理解导致训练不稳定、调参依赖经验、无法预判 warm-up 的必要性。

行业类比

类似 推荐系统中的知识蒸馏：只观察 student 输出概率提升，却不清楚 embedding 空间中哪些维度被重加权，导致后续 A/B 测试和退化分析缺乏依据。

核心洞察

• swap readout 方法解决了稀疏 crosscoder 在跨模型比较时无法区分特征使用变化的缺陷。标准 crosscoder 将所有模型编码到同一组特征激活上，导致无法判断激活差异来自特征本身变化还是使用方式变化。swap readout 通过分别读取每个学生 checkpoint 自身的特征激活，独立衡量每个特征的使用变化，即使这些 checkpoint 在训练 crosscoder 时未见。这一方法学创新为后续分析提供了定量基础，也适用于其他多模型或训练阶段对比的机制解释任务。
• OPD 的核心机制是重新加权学生已共享的特征，而不是从教师传递新特征或创造新特征。作者通过特征字典和激活频率的定量分析发现，超过 98% 的学生频繁使用特征在 OPD 后 firing rate 变化在 20% 以内，并且没有观察到新特征产生。这与“知识蒸馏转移教师知识”的常见直觉相悖，揭示学生内部表征空间基本不变，只是调整了现有特征的重要性。对工程实践的启示是：蒸馏优化可以集中于特征激活模式调整，而非试图改变模型权重结构，或利用特征干预直接提升模型性能。
• SFT warm-up 的作用不是添加特征，而是通过预先和补充调整共享特征权重来增强 OPD 效果。warm-up 阶段已经提升了部分 OPD 后续会提升的特征，并且改变了 OPD 单独不会触及的特征类别（如对话格式、推理风格、数学符号），这些改变在 OPD 后持续。干预实验表明，直接将这种 reweighting 施加到未 warm-up 的学生特征上，可以接近 warm-up 后的准确率。这说明 warm-up 的价值在于特征使用层面的初始化，为优化蒸馏流程中的热身阶段提供了新视角，可以针对性地设计特征调整策略。

方法

方法概述

输入：学生模型在 OPD 前、OPD 后两个 checkpoint 以及教师模型，在相同推理轨迹上的残差流激活，特别关注 decision tokens 上的表示。

稀疏跨编码器：训练一个 Sparse Crosscoder，使三个模型共享一个特征字典。每个模型有其独立的编码器，将自身残差流激活映射到共享字典系数；解码器则从字典重构激活。训练目标同时优化重构误差与 L1 稀疏惩罚。相比单模型 SAE，crosscoder 能将不同模型的语义轴对齐，但标准输出是所有模型合并后的统一激活，无法回答“学生模型自己如何使用特征”。

Swap Readout：为此提出 swap readout，对每个 checkpoint 单独读取其特征激活。具体做法是用该 checkpoint 对应的编码器将激活映射到共享字典，得到只属于该 checkpoint 的特征系数，不与其他模型混合。这样能测量训练前后学生使用每个特征的变化（如 firing rate 相对变化），并且可扩展到 crosscoder 训练时未见过的中间 checkpoint。

输出：每个特征的使用变化分数、被重加权的特征集合、以及针对 decision tokens 的最大变化特征。后续分析包括：检查 OPD 是否创造新特征或传递教师专属特征；分解 SFT warm-up 对共享特征的重加权方向；通过特征级干预验证因果性。

与同类方法差异：标准 crosscoder 输出统一特征激活，无法区分单个模型内部特征使用的演变；swap readout 为每个 checkpoint 单独编码，从而首次以特征级粒度追踪蒸馏训练的动态。

实验

实验设计

使用 稀疏跨编码器 学习一个特征字典,共享于学生初始模型、OPD 后模型与教师模型之间。提出 swap readout 方法,分别读取各 checkpoint 的特征激活,量化训练对学生特征使用的影响,即使该 checkpoint 未参与跨编码器训练。实验覆盖三种 OPD 配置,并分析常见的 SFT warm-up 阶段。

• 对比对象:直接 OPD(无 warm-up)、warm-up + OPD、纯教师特征传递假设。
• 关键步骤:特征归因、决策 token 定位、特征级干预验证。

关键发现

• OPD 并未创建新特征,也未将教师自有特征传递给学生;超过 98% 学生频繁使用的特征其 firing rate 变化在 20% 以内。
• 决策 token 处特征变化最大。
• SFT warm-up 同样不添加新特征,而是以两种方式重加权共享特征:提前执行 OPD 的部分重加权;改变 OPD 单独无法触及的特征(对话格式、推理风格、数学符号),且该变化在 OPD 后持续。
• 特征级干预:在不改变权重的条件下,将直接蒸馏学生的特征重加权为 warm-up 学生的状态,准确率接近 warm-up 学生;打乱特征则无效果。

与基线对比的解读

常规 crosscoder 分析只能标记模型特有特征,无法区分模型对相同特征的使用程度变化;swap readout 弥补了这一缺陷,支持跨 checkpoint 的特征使用追踪。相比无 warm-up 的直接蒸馏, warm-up 通过预先重加权和补充额外特征方向,使 OPD 更高效,且干预实验证实了 warm-up 引入的特征重加权是关键贡献,而非额外特征或教师特征迁移。

行业影响

落地场景

论文揭示 On-policy Distillation (OPD) 主要重新加权学生已有的特征，而非注入新知识。这一发现可直接指导 LLM 推理能力蒸馏 的实践：在电商客服、代码助手、金融分析等场景中，团队常需将大模型能力压缩到小模型以降低延迟和成本。利用 sparse crosscoder 和 swap readout 分析特征使用变化，能提前判断哪些特征应被强化或抑制，避免盲目增加训练数据。

商业价值

• 降本：SFT warm-up 已能完成部分特征重加权，后续 OPD 迭代可缩短，减少 GPU 时耗；且发现无需改变权重、仅通过特征级干预（类似 激活操控）即可接近 warm-up 学生效果，为无梯度蒸馏提供可能，进一步降低训练开销。
• 增效：通过提前干预关键特征（如推理风格、数学符号），可快速适配不同业务领域，避免从头蒸馏。
• 体验提升：更可控的特征调整有助于保留通用能力，减少灾难性遗忘。

接口与集成

可集成进现有 LLM 训练栈：在 OPD 前使用 crosscoder 分析学生-教师共享特征字典，识别需重加权的特征集合；随后两种路径：

1. 将特征权重调整作为先验，指导 warm-up 数据构造（如针对对话格式、推理风格的样本）。
2. 直接对推理时的特征激活施加缩放因子，类似于 SAE 特征 steering 工具，已存在于 TransformerLens 等库中。

具体 use case：某云服务商提供模型蒸馏 API，客户上传教师模型后，平台自动运行 crosscoder 分析，生成 warm-up 数据并报告需干预的特征列表，客户可选择仅通过特征缩放适配其垂直场景（如医疗问诊格式），无需重新训练权重。

局限

• **稀疏交叉编码器的近似性**：论文使用 sparse crosscoders 学习共享特征字典，但稀疏字典可能无法完整捕获模型表征，尤其是低密度或复合特征。模型内部变化可能以非线性方式体现，线性读取和稀疏约束可能导致部分重加权效应被低估或误判。此外，crosscoder 训练涉及教师与学生激活的联合重构，其解不唯一（文中提到 N 矩阵的不确定性），这可能影响特征归因的稳定性。
• **实验设置与泛化边界**：研究仅在三种 OPD 设置（特定模型与推理任务）上验证，未覆盖不同模型规模、架构或非推理领域（如对话、代码生成）。SFT warm-up 的结论可能高度依赖教师 rollout 质量和数据分布。因此，当前结论的普适性有限，不能直接推广到所有蒸馏场景，这也是论文尚未明确指出的潜在边界。
• **因果推断的局限**：尽管 swap readout 和特征级干预展示了 warm-up 重加权与性能提升的相关性，但干预实验（将 warm-up 特征重加权施加到直接蒸馏学生）只能说明特征使用模式与行为相关，不能证明重加权是性能提升的充分或必要原因。同时，对教师特征在学生中的具体功能含义缺乏深层语义验证，因果链仍需更严格的反事实或消融设计来支撑。