做过大模型蒸馏的朋友,大概率都遇到过这样的怪事:教师模型明明很强,生成结果也很漂亮,可学生模型把这些输出当成训练数据学完之后,效果反而不如直接用原始数据训练。有人说是温度参数没设对,有人怪学生模型太小,我在踩了好几轮坑之后才意识到,更隐蔽的瓶颈往往出在数据层——你喂给学生的数据,未必是学生当下“吃得下”的数据。
某高校研究团队最近提出的RSR(Ranking-Sampling-Reweighting)框架,就是专门解决这个问题的。它不改蒸馏损失,不改模型结构,只在数据准备和采样阶段做文章:给每条教师样本打一个“可学习性”分数,再分层筛选,最后在训练时动态重加权。这套思路看着简单,却能把“强师带出弱徒”的怪现象扳回来不少。这篇文章我会把 RSR 的原理、三个关键步骤、我复现时的完整流程,以及实际踩过的坑一次性讲清楚,适合正在做小模型蒸馏但总觉得数据不对劲的同学参考。
1. 为什么强师也会带出弱徒
1.1 教师“会做”不等于学生“能学”
蒸馏的本质,是让学生模型模仿教师模型的行为分布。教师模型是一个已经见过海量数据的“专家”,它输出的每个 token 都蕴含了它内部学到的复杂规律。问题在于,专家的表达水平太高时,学徒是看不懂的。
举个最直观的例子:让一个刚学完四则运算的小模型去模仿一个直接用微积分解题的大模型,小模型每个 token 都能抄下来,但它既理解不了为什么要这样变形,也接不住这种难度梯度的损失。训练的时候,损失函数会不断把压力推向学生模型,逼它去拟合它根本拟合不了的目标分布。后果就是梯度震荡、loss 漂高,最后学生模型连原本会的基础能力都丢了。
我见过不少团队在蒸馏时只盯着教师模型的“质量”看,以为教师越强、输出越干净,蒸馏效果就越好。实际上,教师输出里的每一句话、每一个推理步骤,对学生来说都可能是噪声,因为学生当前没有足够容量去把它们内化。所以选择教师输出,不能只看教师本身,还要看“教师输出与学生当前能力的匹配度”。
1.2 数据里的三个隐性陷阱
除了难度不匹配,蒸馏数据里还藏着三个很容易被忽略的陷阱。
第一个叫难度失配。教师模型认为某个问题很简单,只给了一个很简略的答案;但学生模型基础薄弱,它需要的是带有中间推导过程的、节奏更慢的答案。反过来,教师觉得“这题有点意思”的样本,学生可能完全跟不上。简单说,教师眼中的“难易分布”和学生眼中的“难易分布”不是同一个东西。
第二个叫分布偏移。教师生成的文本,和学生在推理阶段自己生成的文本,分布特征并不一致。训练时学生整天看教师那种“完美的、不卡壳”的文本,到了推理时自己生成了哪怕一个稍微不同的词,后面整个分布都开始跑偏。这有点像模拟器里练开车开得太顺,上了真实道路遇到一点不完美就熄火。教师分布和学生自己生成分布之间的差距,如果不处理,学生就是在用错误坐标系学习。
第三个叫噪声放大。教师模型虽然整体很强,但并不是每条输出都可靠。尤其是在长文本生成、复杂推理、低资源场景下,教师也会产出事实错误、逻辑断裂的内容。学生模型没有足够的判别能力,不会拒绝这些错误,它会像海绵吸水一样,把教师的幻觉一并吸收。而错误信息一旦进了学生模型,后续再想清洗就难了。
1.3 黄金数据应该长什么样
既然全盘照收不行,那理想的蒸馏数据应该是什么样?我把它概括成三个特征:学生学得动、信息量够、教师信得过。
“学得动”指的是样本难度落在学生模型的最近发展区里,接近“跳一跳够得着”的状态。“信息量够”意味着样本不能太简单,如果学生已经能完美预测每一个 token,那这条样本对训练就没有增量信息。“教师信得过”则指教师在这个样本上的输出没有幻觉、逻辑是自洽的。
这三个特征单独拎出来都不难理解,但难在它们会互相牵制。一条样本可能信息量很足,但难度太高;也可能教师很自信,但学生完全接不住。RSR 做的事情,就是把这个多目标的选择问题,拆解成可计算的打分、可控制的分层和可调整的加权,让“黄金数据”不再是一句抽象口号。
2. RSR 方法的核心拆解
2.1 三步走:Rank、Select、Reweight
RSR 的全称是 Ranking-Sampling-Reweighting,用一条流水线完成数据筛选。
第一步是Rank(排序打分)。对教师模型生成的所有候选样本,分别计算一个综合评分。评分至少包含三个子项:学生模型在当前参数下对该样本的学习难度、教师模型输出本身的可靠性、以及这条样本可能带来的信息增益。每个子项都会被归一化,再按权重相加,最终得到一个可比较的标量分数。
第二步是Select(分层采样)。拿到所有样本的分数之后,不直接截取 TopK,而是先把样本按照分数区间分成若干个桶,再从每个桶里按比例抽取。这样做的原因很现实:如果只选最高分的样本,大概率会筛出一堆“学生已经会了”的简单样本,模型会越学越平庸;如果只选最低分的样本,学生又会被远超能力的难题带崩。分层采样能强制保留一部分有梯度的中等难度样本。
第三步是Reweight(动态重加权)。在正式训练过程中,每条样本对 Loss 的贡献不是固定不变的。随着学生模型逐步变强,原本“太难”的样本会慢慢变成“刚好合适”,我们再给它增加权重。RSR 会周期性重新计算样本的可学习性分数,并用指数移动平均去平滑权重,避免训练中途权重剧烈跳动。
2.2 可学习性分数是怎么算的
综合评分里最关键也最容易做偏的,是“可学习性分数”。我实现的时候用了这样一个组合思路:
先算一个基础难度指标,用教师输出作为目标序列,计算当前学生模型在教师输出上的交叉熵损失。这个损失值越大,说明学生离拟合这条样本还差得远;损失值越小,说明学生已经很接近了。但注意,这里不能简单地认为“损失越小越好”。如果损失趋近于 0,意味着样本已经没有教学价值,训练时它对梯度的贡献接近于 0。
所以可学习性分数会设计成一个倒 U 形函数:损失太低或太高的样本都会得到低分,只有落在某个中间区间里的样本才能拿高分。这个区间的边界可以由训练初期一个小的探测集计算出来,比如取全体样本损失分布的 20% 到 80% 分位作为“合适区间”。
其次是教师可靠性分数。我这里用了一个不额外增加推理成本的方法:对同一条 prompt,让教师模型生成多个候选输出,然后计算这些输出之间的平均相似度。如果教师对同一个问题给出的多个答案高度一致,说明它在稳定输出,可靠性较高;如果多个答案差异巨大,说明教师自己也没底,这条样本的权重就要调低。
最后是信息增益分数。一个简化的估计方式是计算教师输出和学生模型自己对该 prompt 的最高置信度输出之间的 KL 散度。KL 散度越大,说明教师输出里含有学生当前分布里没有的信息,潜在增益越高。但这一步要和可学习性分数做平衡,否则会因为 KL 太大而选到教师答案里完全偏离学生能力的部分。
综合评分可以写成:
Score_i = α * Learnability_i + β * Reliability_i + γ * Informativeness_i
其中 α、β、γ 是三个超参数,一般初始可以设成 0.5、0.3、0.2,再根据具体任务微调。
2.3 为什么是分层采样,而不是简单 TopK
很多人会问:已经有了分数,直接按分数从高到低取前 N 条不就行了吗?我在初版实现里就是这么干的,结果效果很一般。
问题在于,单一分数压缩了多样性。TopK 选出来的样本,多半集中在一条“看似完美”的窄带里,它们的学生损失、教师可靠性都很相似,导致训练数据的难度分布极其单一。学生模型反复接触同一类难度、同一类句式,很快就会过拟合到这条窄带上。
分层采样是在 TopK 基础上加了一层“配额制”。我一般把分数从低到高切成 5 个桶,分别代表“极难、偏难、适中、偏易、极易”。然后按经验比例抽:适中桶抽 40%,偏难和偏易各抽 20%,极难和极易各抽 10%。这个比例不是死的,如果发现学生容易崩溃,就调低极难桶比例;如果发现学生输出越来越平淡,就调高偏难桶比例。
这样做的好处是既保证了整体数据以“黄金区间”为主,又保留了两翼的极端样本。极端样本中的极难样本虽然权重低,但它们能作为“上限刺激”防止学生过早收敛;极易样本则充当“热身”,帮助学生稳定训练初期梯度。
3. 实操流程与关键实现
3.1 跑通一个最小实验需要什么
复现 RSR 的硬件门槛不算高。我用的是一张 24GB 显存的商用显卡,教师模型和学生模型都是开源权重加载的,学生模型是 1B 级别的小模型,教师模型是 7B 级别的大模型。如果你手头资源更紧张,也可以用一个 API 来生成教师输出,只要把输出和对应的 prompt 保存下来,后续打分就能离线进行。
数据方面,你需要一份覆盖目标场景的 prompt 集合。数量不需要特别夸张,我做的对话压缩任务里用了 2000 条 prompt,教师模型对每条 prompt 生成 8 个候选输出,总共就得到 16000 条候选样本。这个量级足够看出不同筛选策略之间的差异了。
需要格外注意的是,教师模型和学生模型最好使用同一个 tokenizer,或者至少是同一个子词词表。否则计算学生模型在教师输出上的损失时,会出现大量 token 被映射成 UNK 的情况,可学习性分数就不能用了。这一点被我列为最容易踩的坑之一。
3.2 打分阶段的代码与计算细节
打分阶段需要做两件主要的事:第一,用教师模型生成候选样本并记录置信度;第二,用当前学生模型计算这些样本的损失和 KL 散度。下面是我精简后的伪代码,可以帮你理清流程。
# 假设已有 train_prompts 列表,每个元素是一个 prompt 字符串 scores = [] for prompt in train_prompts: # 教师模型生成候选输出 teacher_outputs = teacher.generate( prompt, num_return_sequences=8, max_new_tokens=256, return_dict_in_generate=True, output_scores=True ) # 计算教师可靠性:候选输出之间的平均 self-BLEU reliability = average_pairwise_similarity(teacher_outputs.sequences) for output_seq in teacher_outputs.sequences: output_text = tokenizer.decode(output_seq) # 学生模型在当前教师输出上的 loss student_loss = compute_cross_entropy(student_model, prompt, output_text) # 学生自己对该 prompt 生成的结果,用于算 KL 散度 student_own_output = student_model.generate(prompt, max_new_tokens=256) info_gain = kl_divergence(student_own_output, output_text) learnability = inverted_u_score(student_loss, loss_low=1.2, loss_high=4.0) score = 0.5 * learnability + 0.3 * reliability + 0.2 * info_gain scores.append({ "prompt": prompt, "teacher_output": output_text, "student_loss": student_loss, "reliability": reliability, "score": score })这里的inverted_u_score我建议用高斯型函数实现:当损失落在理想区间时给 1,偏离后按指数衰减。用硬截断容易在桶边界产生突变,影响后面重加权的稳定性。
打分阶段最耗时的是教师模型生成。为了省时间,我通常把候选生成和打分分两步跑,先生成所有候选文本并缓存到磁盘,再用学生模型离线打分。打分的计算量其实不大,因为学生模型只有 1B 级别,跑一遍前向也就几分钟。
3.3 分层采样和动态重加权的实现
拿到 scores 之后,先按分数从低到高排序,再切成 5 个等宽的桶。所谓等宽,不一定每个桶里样本数一样,而更建议按分数分位点来切,保证分数分布不同区间的代表性。然后按前面说的比例从每个桶里抽样。抽样时要设置随机种子,方便重复实验。
动态重加权让我改了好几个版本才稳定。最简单的版本是训练之前算好每个样本权重,以后不再变化,但这样不能体现“学生变强后样本难度下降”的动态过程。我最后采用的是每 500 步重新计算一次样本权重,方法是:
new_weight_i = clamp(learnability_score_i at current_step, 0.1, 2.0)
再用指数移动平均和旧权重混合:
weight_i = 0.8 * weight_i + 0.2 * new_weight_i
这样能避免第 500 步时由于单次计算噪声导致权重突跳。训练 Loss 中,每条样本的交叉熵会乘上这个权重。权重只作用于蒸馏数据对应的那部分 Loss,其他预训练或原始任务 Loss 不要动,否则整个模型的优化方向会被带偏。
3.4 训练阶段的配置建议
我给采样后的数据设定一个固定 batch size,通常 16 到 32 之间。学生模型的学习率建议比普通微调低一点,我用 1e-5 到 2e-5 比较稳。蒸馏任务不像预训练那样需要高学习率探索新区域,它更像在教师划定的范围里做精调,学习率太大会破坏学生原本学到的通用能力。
另外,如果教师模型和学生模型的 tokenizer 相同,可以考虑在交叉熵之外加一个 KL 蒸馏 Loss,目标是让学生的输出分布去逼近教师的 logits。不过这个并不是 RSR 的必要条件,RSR 只负责决定“哪条数据被用来训练”,不限制你选哪种蒸馏损失。我实测的时候发现,单纯用标准交叉熵学教师输出文本,再配合 RSR 分层,已经能取得明显提升。加 KL 蒸馏后效果还能再上一层,但训练时间至少要翻一倍。
还有一个细节:把教师生成候选时的不 desirable token 过滤掉。尤其是那些包含大量重复、乱码或者过短无意义的输出。在打分之前先做一次基础清洗,能避免后面把噪声当成高可靠性样本。
4. 复现效果与对比
4.1 我的对比设置
我挑了一个小规模对话摘要任务做复现测试,将训练集切成三份,分别使用三种数据策略训练相同结构的学生模型。第一种是随机采样,不筛选,只要教师生成就拿来学;第二种是TopK 采样,按 RSR 的综合 Score 排名后取前 20%;第三种是RSR 全流程,也就是打分加分层采样加动态重加权。
为了公平,三种策略保证每个 epoch 能见到的样本总数一致。训练步数一致,学习率和 batch size 也都一致。最后在同一套测试集上评测生成质量,用了一个包含流畅度、相关性和事实一致性的综合打分,满分为 100。
下表是我跑了三组不同随机种子后取平均的结果。这里不贴具体模型名,只看相对趋势。
| 策略 | 综合得分 | 学生 Loss 波动 | 生成多样性 |
|---|---|---|---|
| 随机采样 | 62.3 | 高 | 高 |
| TopK 采样 | 65.1 | 低 | 低 |
| RSR 全流程 | 69.8 | 中 | 中高 |
4.2 关键结果和解读
随机采样的数据量最大,理论上覆盖最全,但综合得分最低。原因是里面混了大量学生学不动的难样本和教师幻觉样本,训练时 loss 经常忽高忽低,说明模型被互相冲突的梯度拉扯了。
TopK 采样把分数最高的样本挑出来,综合得分比随机高了不少,但生成多样性骤降。学生输出明显变得“千篇一律”,遇到稍微偏离训练分布的输入就开始套模板。这正好印证了单一分数截断的缺陷:它在压缩难度分布。
RSR 全流程得分最高,并且生成多样性保持在中等偏上水平。它的 Loss 波动介于随机和 TopK 之间,说明分层里的困难样本仍然会带来一定梯度冲击,但这种冲击被重加权控制在了可接受范围内。说白了,RSR 不是让训练变得毫无波澜,而是把波澜控制在了“能刺激进步”的强度上。
需要说明的是,这组差异在简单任务上并不明显。我拿一个造句任务试过,三种策略都很快收敛,RSR 的优势只有 1 到 2 分。所以 RSR 更适合教师和学生能力差距较大、任务本身有层次感的场景。如果任务太简单,数据选择带来的增益会被随机噪声淹没。
5. 常见问题与排查
5.1 问题速查表
复现 RSR 时我遇到了不少问题,下面把典型的几种和对应排查方向整理成了一张表。
| 现象 | 可能原因 | 排查/解决 |
|---|---|---|
| 训练开始后 Loss 飙升 | 极难样本比例过高,学生模型被超出能力的样本干扰 | 调低“极难”桶的采样比例,或调低综合评分里可学习性分数的权重 |
| 训练很稳但评测不涨 | 筛选出的样本太简单,信息量不足 | 提高偏难桶比例,增大信息增益分数权重 |
| 生成结果模板化、重复严重 | 分层采样比例失衡,简单样本控场 | 减少极易桶比例,提高偏难桶比例 |
| 教师输出之间相似度虚高 | 教师生成都在复读训练数据 | 换更多样化的 prompt,或在生成时调高 temperature |
| 重加权后 Loss 剧烈震荡 | 权重更新太激进,没有做平滑 | 检查 EMA 的动量因子,我一般设 0.8 到 0.9 |
5.2 参数调整心得
RSR 里我最先调的是综合评分公式里的 α、β、γ。直觉上 α 最重要,但我把 α 从 0.5 一路调到 0.8,发现效果反而变差。原因还是那个老问题:可学习性倒 U 形函数本身已经把“难度”考虑进去了,权重再加大,会让极难样本的综合评分被低估得太狠,分层采样时极难桶可能完全没样本。
我的习惯是先把 β 设为 0.3 不动,只调 α 和 γ,并且保证 α + γ = 0.7。这样只在一个自由度上调参,更容易观察趋势。另一个重要参数是采样桶数。桶数太少,分层就失去了多样性控制;桶数太多,每个桶里样本数太少,随机抽样不稳定。5 个桶是我用得最顺的配置,如果你的数据量很少,建议减到 3 个。
5.3 它和 IFD、LFS 的区别
细心的同学可能会问,RSR 和已有的 IFD(Instruction Following Difficulty)、LFS(Learning from Strong outputs)等数据选择方法到底有什么不同?
IFD 的核心也是用学生模型计算教师输出的损失,但它主要是在做难度过滤,并没有把教师可靠性、信息增益和动态重加权揉在一起。RSR 更像一个组合框架,IFD 类方法可以作为其中的“可学习性评分”子模块。LFS 偏重于从教师候选中选择能让训练损失降低最多的样本,思路和梯度匹配有点像,但 LFS 没有显式处理教师噪声。RSR 的教师可靠性部分正好补上了这块短板。
实际选型时,我建议不要纠结于谁取代谁。你完全可以把 IFD 的分数放进 RSR 的评分公式里,作为可学习性分数的一个变体。数据选择方法之间是互补的,不是互斥的。
6. 可扩展方向与我的个人体会
6.1 多教师蒸馏与在线课程化
RSR 目前默认是一位教师、一套学生参数。如果把教师拓展到多个模型,只需要分别计算每位教师输出在综合评分里的可靠性,再在分层采样时让不同教师的样本尽量均衡。教师之间的分歧度其实可以作为另一个过滤信号:分歧大的样本通常教师置信度偏低,权重可以下调。
另一个我很看好的方向是把 RSR 做成“在线课程”。目前我是每 500 步重新计算权重,但如果算力允许,每 1 到 2 个 batch 就做一次增量式权重更新,那学生模型就相当于有了一条实时的学习路径。难点在于频繁调用学生模型前向会拖慢训练速度,所以实际工程里需要在“实时性”和“吞吐量”之间做取舍。
6.2 我的个人体会
把 RSR 完整跑通之后,我最大的感受是:大模型蒸馏里,数据选择常常被当成一件“预处理”小事,可它往往比模型结构选择和损失函数设计更能影响最终效果。判断一条训练样本好不好,不能只看教师模型的输出质量,还要看它在学生模型当前能力下是否“可学”。这个道理说起来简单,真正执行时却很容易被“越大越好、越难越好”的直觉带偏。
我后来做蒸馏项目,都会先跑一遍 RSR 的三步流程,哪怕只是用它把数据分层看看分布,也会比盲目全量训练更有方向感。个人建议先把随机种子定死、把打分缓存好,再开始调参数。否则每改一次参数都要重新生成候选输出,时间和耐心都会被快速消耗。
如果手头算力紧张,也不必强求每一步都做到位。哪怕只实现“Rank + 分层采样”,放弃动态重加权,效果也会比纯 TopK 好不少。等到确认瓶颈确实在数据层时,再回来把 Reweight 这一步加上。这大概就是 RSR 最实用的地方:它给了你一个从数据侧思考蒸馏的抓手,而不是让你在模型训练循环里盲调。