TAKE: 面向文本数据集蒸馏的轨迹感知知识估计
摘要
本文提出TAKE(轨迹感知知识估计),一个文本数据集蒸馏框架,利用影响函数和最优传输将数据集缩减至原始大小的0.1%,同时保持下游任务性能。
arXiv:2607.11898v1 公告类型:新
摘要:大规模文本语料库已成为现代自然语言处理中的隐形瓶颈,不仅体现在存储上,更体现在训练、微调和持续学习的累积成本中。本文提出一个文本数据集蒸馏框架,将语料库缩减至原始大小的0.1%,同时保持下游任务性能。我们通过影响函数(influence functions)的视角进行蒸馏,该函数量化每个样本对下游目标的贡献,这是一种自然且基于原理的选择依据。我们引入了轨迹感知知识估计(TAKE),它将训练轨迹上的基于知识的影响卷积为单个样本知识得分,从而捕获信息丰富的样本。这些得分作为离散最优传输目标中的样本权重,指导从合成候选池中选择原型。我们在极端压缩(0.1%或每类20个样本)下的文本分类和自然语言推理任务上评估TAKE的下游准确率,结果表明无需牺牲任务性能即可实现数据效率。该方法具有理论依据,对核心集构建和数据中心AI有更广泛的影响。我们在 https://github.com/votrinhan88/take 上发布源代码。
查看缓存全文
缓存时间: 2026/07/15 04:21
# 文本数据集蒸馏的轨迹感知知识估计
来源:https://arxiv.org/html/2607.11898 (2026年3月)
###### 摘要
大规模文本语料库已成为现代 NLP 中的隐性瓶颈,这不仅体现在存储上,更体现在训练、微调和持续学习的累积成本中。我们提出了一种文本数据集蒸馏框架,可将语料库减少至原始大小的 0.1%,同时保持下游任务保真度。我们通过影响函数的视角进行蒸馏,该函数量化每个样本对下游目标的贡献,这是一种自然且有原则的选择基础。我们引入了**轨迹感知知识估计**(TAKE),它将沿训练轨迹的基于知识的影响卷积为每个样本的知识得分,从而捕获信息丰富的样本。这些得分作为离散最优传输目标中的样本权重,指导从合成生成的候选池中选择原型。我们在文本分类和自然语言推理任务上,在极端压缩条件下(0.1% 或每类 20 个样本)评估了 TAKE 的下游准确率,结果表明数据效率可以在不牺牲任务保真度的情况下实现。该方法具有理论基础,对核心集构建和数据中心 AI 具有更广泛的意义。我们在 https://github.com/votrinhan88/take 发布了源代码。
关键词:文本数据集蒸馏⋅影响函数⋅最优传输。
## 1 引言
大型语言模型的进步依赖于大规模监督语料库,然而,在微调循环、超参数搜索和持续学习过程中,整理和训练此类数据的成本不断累积。真正的负担不在于单次运行,而在于整个开发生命周期中的累积支出。语料库规模(而非模型规模)往往是约束瓶颈。因此,减小语料库规模是一个系统性优先事项,直接关系到成本、碳足迹和数据治理。
数据集蒸馏(DD)由 [1 (https://arxiv.org/html/2607.11898#bib.bib1)] 首次提出,它用一个较小的代理数据集 D̃ ≪ D 替换大型数据集 D,使得在 D̃ 上训练的模型性能与在 D 上训练的模型相当。换句话说:能教同一门考试的最小编教材是什么?这使 DD 完全属于数据中心 AI 范畴,通过作用于数据本身来补充模型压缩。早期 DD 方法在图像空间中运行,通过元学习、梯度或轨迹匹配来优化合成像素 [1 (https://arxiv.org/html/2607.11898#bib.bib1),2 (https://arxiv.org/html/2607.11898#bib.bib2),3 (https://arxiv.org/html/2607.11898#bib.bib3)]。文本打破了这些假设:不可微分的令牌解码阻碍了基于梯度的合成,且嵌入空间与特定架构紧密耦合。因此,将 DD 适配到文本需要根本不同的策略。
现有的文本 DD 方法 [4 (https://arxiv.org/html/2607.11898#bib.bib4),5 (https://arxiv.org/html/2607.11898#bib.bib5),6 (https://arxiv.org/html/2607.11898#bib.bib6),7 (https://arxiv.org/html/2607.11898#bib.bib7),8 (https://arxiv.org/html/2607.11898#bib.bib8)] 至少在以下方面存在不足:在极端压缩下,均匀加权会将预算耗尽在无信息样本上;它们的目标缺乏任务对齐或全局优化;它们的输出是嵌入,既不可审计也不可迁移。我们在 §2.1 (https://arxiv.org/html/2607.11898#S2.SS1) 中调研了该领域,并表明没有先前的方法能同时解决上述所有问题。这些差距促使 TAKE 具备三个期望:(1) **基于知识的加权**——根据样本的任务对齐下游贡献对其进行加权,将预算导向信息最丰富的样本;(2) **数据集级蒸馏目标**——在整个训练分布上优化,而不是逐批次优化,确保全局覆盖;(3) **人类可读输出**——可直接检查、审计且跨架构可迁移的蒸馏语料库。
我们提出了**轨迹感知知识估计**(TAKE),这是首个满足所有三个期望的文本 DD 方法。我们通过影响函数 [9 (https://arxiv.org/html/2607.11898#bib.bib9)] 解决第一个期望,该函数量化每个样本对下游损失的反事实影响,并提供具有理论依据且任务对齐的加权方案(§3 (https://arxiv.org/html/2607.11898#S3))。然而,在单个检查点评分会引入**难样本偏差**(HSB):在收敛时,噪声样本主导影响分数,而干净样本和中等样本(对于鲁棒的蒸馏数据集最有价值)则受到抑制。TAKE 通过将沿完整训练轨迹的影响分数整合为每个样本的知识得分来纠正这一点,该得分捕获干净样本和中等样本,同时降低噪声样本的权重。然后,TAKE 微调一个 LLM 以生成人类可读的候选实例池。知识得分和合成候选共同输入一个离散最优传输(OT)目标,该目标全局选择蒸馏语料库。具体来说,我们贡献了:
- **理论**:我们在文本 DD 设置中形式化了难样本偏差(HSB),证明了单检查点影响分数偏向噪声样本,并推导了知识重加权分布匹配的形式化间隙界限。
- **方法**:TAKE 是首个同时解决样本加权和分布覆盖的文本 DD 方法——通过轨迹集成的知识得分和与 DD 目标对齐的离散 OT 目标。
- **实证**:在极端压缩条件下(0.1% 或每类 20 个实例),TAKE 在六个语言基准上达到或超过了先前的文本 DD 方法,产生了跨不同骨干家族泛化的人类可读蒸馏语料库。
本文其余部分组织如下:§2 (https://arxiv.org/html/2607.11898#S2) 回顾相关工作;§3 (https://arxiv.org/html/2607.11898#S3) 提出理论动机;§4 (https://arxiv.org/html/2607.11898#S4) 详细介绍 TAKE 及我们的 DD 流程;§5 (https://arxiv.org/html/2607.11898#S5) 报告实验和讨论;§6 (https://arxiv.org/html/2607.11898#S6) 反映局限性和伦理问题;§7 (https://arxiv.org/html/2607.11898#S7) 总结。
## 2 相关工作
### 2.1 文本数据集蒸馏
数据集蒸馏由 [1 (https://arxiv.org/html/2607.11898#bib.bib1)] 作为双层元学习问题引入:合成一小部分图像,使得在其上训练的模型能够匹配在完整数据集上训练的模型。后续方法通过梯度匹配 [2 (https://arxiv.org/html/2607.11898#bib.bib2)]、分布匹配 [10 (https://arxiv.org/html/2607.11898#bib.bib10)] 和轨迹匹配 [3 (https://arxiv.org/html/2607.11898#bib.bib3),11 (https://arxiv.org/html/2607.11898#bib.bib11)] 提高了可伸缩性和保真度。这些方法有一个共同假设:数据是连续且可微的,因此合成集可以通过反向传播直接优化。文本违反了可微性假设:令牌解码不可微,输出必须语义连贯,且嵌入空间与特定架构紧密耦合。因此,将 DD 适配到文本需要根本不同的策略——从合成嵌入的元学习到基于 LLM 的可读语料库生成,每一代都在用一项限制交换另一项限制。
SLDD 和 DDTC [4 (https://arxiv.org/html/2607.11898#bib.bib4),5 (https://arxiv.org/html/2607.11898#bib.bib5)] 直接适配了元学习公式 [1 (https://arxiv.org/html/2607.11898#bib.bib1)],优化合成文本嵌入,使得在其上训练的模型能够最小化任务损失。在这两者中,只有 SLDD 通过将蒸馏嵌入映射到每个令牌的最近词袋部分恢复了可解释性。DDAL [6 (https://arxiv.org/html/2607.11898#bib.bib6)] 优化了 BERT 模型的注意力标签,以匹配类别级数据集统计量,实现了部分任务对齐,但产生了既锁定架构又不可解释的输出。
第二代方法通过重新思考整个流程来解决可解释性问题。DiLM [7 (https://arxiv.org/html/2607.11898#bib.bib7)] 通过 LLM 生成可读候选文本,并通过梯度匹配(在训练语料库上)从中进行选择,随后进行 k 中心选择;但在离散检查点匹配梯度只是任务损失的代理,导致任务对齐仅为部分。DaLLME [8 (https://arxiv.org/html/2607.11898#bib.bib8)] 学习从蒸馏嵌入到合成文本的逆映射,恢复了可解释性并尝试通过 k 质心聚类实现全局覆盖;但由于没有任务损失目标,任务对齐完全缺失。表 1 (https://arxiv.org/html/2607.11898#S2.T1) 展示了这一进展:每一代都解决了一些差距,同时留下了其他差距。
表 1:文本 DD 方法概览。
| 差距 | SLDD | DDTC | DDAL | DiLM | DaLLME | TAKE |
|------|------|------|------|------|--------|------|
| 任务对齐 | ✓ | ✓ | △ | △ | − | ✓ |
| 全局优化 | △ | △ | △ | △ | △ | ✓ |
| 可解释性 | △ | − | − | ✓ | ✓ | ✓ |
| 重要性加权 | − | − | − | − | − | ✓ |
−:缺失,△:部分,✓:已解决
一个更微妙的差距贯穿所有方法:每个现有方法都应用**均匀加权**,将每个训练样本视为同等信息量。在适中的蒸馏预算下,这也许可以辩护,但极端压缩正是蒸馏最重要的场景,此时该假设失败。样本在学习中的影响差异很大(干净 vs. 噪声,简单 vs. 困难),但之前没有方法考虑这种结构。没有重要性加权,蒸馏预算会被耗尽在无信息样本上,使选择偏向于对鲁棒性最重要的边界情况。
### 2.2 影响函数与难样本偏差
影响函数(IF)[12 (https://arxiv.org/html/2607.11898#bib.bib12),9 (https://arxiv.org/html/2607.11898#bib.bib9)] 通过 I(z)=∇⊤H_{θ}^{-1}∇ 量化训练点的留一法效应,但依赖于单次检查点,且需要昂贵的二阶计算。基于轨迹的方法 [13 (https://arxiv.org/html/2607.11898#bib.bib13),14 (https://arxiv.org/html/2607.11898#bib.bib14),15 (https://arxiv.org/html/2607.11898#bib.bib15)] 将归因扩展到多个检查点,但它们是**测试条件的**:设计用来归因固定的测试预测。自影响近似通过将每个训练点作为其自身查询,将它们适配到语料库级评分。然而,自影响仅在最终参数处评估一次,引入了系统性偏差,破坏了语料库级评分。我们称之为**难样本偏差**(HSB):在收敛时,干净样本的梯度几乎为零,因此噪声样本或困难样本主导影响分数。这种失败模式在相邻工作中已有充分记录:[16 (https://arxiv.org/html/2607.11898#bib.bib16)] 表明网络先拟合干净模式,后拟合噪声模式;数据集制图 [17 (https://arxiv.org/html/2607.11898#bib.bib17)] 证实单检查点统计量混淆了难度与噪声;Co-teaching [18 (https://arxiv.org/html/2607.11898#bib.bib18)] 记录了基于梯度的课程中噪声样本的系统性过度选择。在标准训练中,HSB 通常可以通过正则化缓解;然而在 DD 的极端低数据场景中,它会放大:有偏选择抑制了对模型鲁棒性至关重要的**信息丰富**样本。据我们所知,之前没有文本 DD 方法识别或纠正了这一点。
### 2.3 用于分布匹配的最优传输
分布匹配已被提出作为可处理的 DD 代理 [19 (https://arxiv.org/html/2607.11898#bib.bib19)]:如果蒸馏集和训练集沿优化轨迹诱导出相等的期望梯度,那么训练的模型会收敛到相似的参数。这是一个充分条件,但不是等价条件(形式化间隙界限出现在定理 1 (https://arxiv.org/html/2607.11898#Thmtheorem1) 中)。常见的散度在本文相关的极端预算下存在实际缺陷。MMD [20 (https://arxiv.org/html/2607.11898#bib.bib20)] 缺乏样本级分配结构,需要辅助评分启发式,这重新引入了分布匹配本意要消除的全局优化间隙。Kullback–Leibler 散度和 Jensen–Shannon 散度需要密度估计,这在低数据场景下不可靠 [21 (https://arxiv.org/html/2607.11898#bib.bib21)]。Monge-Kantorovich 意义上的 OT 避免了这两个缺陷:它天生操作于离散测度,尊重嵌入空间的度量几何,并产生一个显式的传输计划,直接将训练样本分配给原型——无需辅助启发式。通过 Sinkhorn-Knopp [22 (https://arxiv.org/html/2607.11898#bib.bib22)] 进行熵正则化使其在语料库规模上可处理。OT 已在 NLP 和机器学习中得到广泛应用 [23 (https://arxiv.org/html/2607.11898#bib.bib23),24 (https://arxiv.org/html/2607.11898#bib.bib24),25 (https://arxiv.org/html/2607.11898#bib.bib25),26 (https://arxiv.org/html/2607.11898#bib.bib26)],但从未用于文本 DD,也从未与任何 DD 设置中的知识重加权源分布相结合。TAKE 填补了这两个空白:κ\_n 定义了一个非均匀源测度,而离散熵 OT 计划无需任何辅助选择步骤即可选择覆盖知识加权训练分布的原型。
## 3 理论动机
### 3.1 用于数据集蒸馏的重加权分布匹配
数据集蒸馏(DD)目标寻求 |D̃| = K ≪ N 的 D̃,以最小化原始数据上的损失:
min_{D̃} E_{θ0~p(θ0)} [ E_{z~D} [ ℓ(z; F(θ0; D̃, η, T)) ] ] (1)
其中 F(θ0; D̃, η, T) 表示从初始化 θ0 开始在 D̃ 上使用学习率 η 运行 T 步梯度下降得到的模型。这个双层问题对于离散文本来说是难处理的。我们通过梯度场等价论证 [19 (https://arxiv.org/html/2607.11898#bib.bib19)] 将其松弛为分布匹配:如果 D̃ 在整个训练过程中诱导出与 D 相同的期望梯度场,那么两次运行会收敛到相似的参数。由于 ∇_θ E_{z~P} [ℓ(z;θ)] = E_{z~P} [∇_θ ℓ(z;θ)],在任务相关的特征空间中匹配分布可以控制梯度场,这使得分布匹配成为等式 1 (https://arxiv.org/html/2607.11898#S3.E1) 的可处理代理。这产生了分布匹配(DM)目标:
min_{|D̃|=K} d(P_{D̃}, P_D) (2)
其中 d 是任务相关特征空间中的散度。等式 2 (https://arxiv.org/html/2607.11898#S3.E2) 是一个代理,而不是等价;形式化间隙界限出现相似文章
面向低资源开源Text-to-SQL模型的知识蒸馏
本文提出了一种知识感知的Text-to-SQL框架,利用知识蒸馏在低资源环境下通过构建任务特定的知识库并生成合成训练数据来提升性能。在七个基准上的实验表明,该方法带来了显著的改进,尤其是对于开源模型。
轨迹即师:通过能量导航蒸馏实现少步离散流匹配
本文介绍了轨迹塑造离散流匹配(TS-DFM),该方法以引导式导航取代盲目随机跳跃,显著提升了文本生成效率并降低了计算成本。与传统多步基线相比,该方法在保持推理成本不变的同时,实现了更低的困惑度和更快的速度。
用于自动化AI导师评估的知识蒸馏
本文介绍了FATE,一个80亿参数的语言模型,利用知识蒸馏评估AI导师,在BEA 2025共享任务的四个教学维度基准上实现了高达22.63个百分点的提升。
KDFlow:面向大语言模型的用户友好且高效的知识蒸馏框架
KDFlow是一种新颖的大语言模型知识蒸馏框架,采用解耦架构,使用SGLang进行教师推理,FSDP2进行学生训练,相比现有框架实现了1.44倍至6.36倍的加速。
同策略蒸馏(5分钟阅读)
本文引入同策略蒸馏,通过在教师提供的token级KL正则化下,在学生自身轨迹上训练学生模型,解决训练-推理分布不匹配问题,统一了前向KL、反向KL和JSD损失,其中反向KL更适用于较小的学生模型。