TAKE框架：轨迹感知知识估计实现文本数据集高效蒸馏

背景

大规模文本语料库已成为现代 NLP 中一个不易察觉的瓶颈——不仅体现在存储开销上，更体现在训练、微调和持续学习所积累的计算成本上。为了缓解这一问题，数据集蒸馏（dataset distillation）旨在将原始数据集压缩到极小规模，同时尽可能保持下游任务的性能。然而，现有方法大多针对图像领域设计，难以直接迁移到文本数据。本文提出一种文本数据集蒸馏框架，能够将语料压缩至原始规模的 0.1%，同时保留下游任务的保真度。

核心内容

本文的核心贡献是 Trajectory-Aware Knowledge Estimation (TAKE)，一种基于轨迹感知的知识估计方法，用于文本数据集蒸馏。整体框架如下：

1. 影响函数量化贡献：作者从影响函数（influence functions）出发，这是评估每个样本对下游目标函数贡献的自然且原则性的工具。影响函数可以定量衡量一个训练样本对模型参数或最终损失的影响程度。
2. 轨迹感知的知识估计：传统的单点影响评估只考虑训练终点，而 TAKE 将影响函数沿训练轨迹进行卷积，将训练过程中每个样本的知识贡献整合为一个单一的 per-sample 知识分数（knowledge score）。该分数能够捕捉那些在训练过程中反复提供信息量的样本，而不仅仅是最终时刻有影响力的样本。
3. 离散最优传输指导原型选择：得到每个样本的知识分数后，将其作为样本权重，纳入离散最优传输（discrete Optimal Transport）目标中。该目标从合成生成的候选池（synthetically generated candidate pool）中选择原型样本，使得选出的原型集合能够最大程度地保留原始数据集的分布信息，同时权重反映了样本的知识重要性。

TAKE 在多个文本分类和自然语言理解（NLI）任务上进行了评估，采用极端压缩率（0.1%，或每类 20 个样本）。实验表明，在该压缩率下，TAKE 能够实现数据效率，且不牺牲下游任务保真度。该方法在理论上具有坚实的基础，并对 coreset 构建和数据中心 AI（data-centric AI）具有更广泛的意义。作者已开源代码。

关键要点

• 核心问题：大规模文本语料库的训练成本成为瓶颈，需要一种高效的文本数据集蒸馏方法。
• 方法创新：TAKE 通过沿训练轨迹卷积影响函数，获得每个样本的轨迹感知知识分数，克服了传统影响函数只考虑终点的局限性。
• 原型选择机制：将知识分数作为权重，融入离散最优传输目标，从合成候选池中选取代表性子集，保证分布覆盖与信息保留。
• 极端压缩能力：在压缩至原始规模 0.1%（每类仅 20 个样本）的条件下，下游任务准确率仍能与全量数据训练持平。
• 评估任务：文本分类和自然语言推理（NLI），涵盖多种 benchmark。
• 理论支撑：基于影响函数和最优传输，方法具有严谨的数学基础，不依赖经验启发式。
• 开源发布：代码已公开，便于复现与拓展。

意义与影响

TAKE 为文本数据集蒸馏提供了一种理论上自洽且实践有效的新范式。其意义体现在：

• 走向数据效率：在数据规模持续膨胀的背景下，TAKE 表明极低的数据量仍可保持模型性能，为低成本微调和持续学习铺平道路。
• 推广影响函数：将影响函数从静态评估扩展到动态轨迹感知，提升了其在实际训练过程中的表征能力。
• 连接 coreset 与蒸馏：将蒸馏问题形式化为带权重的原型选择，与 coreset 构建和最优传输紧密关联，推动数据选择方法的发展。
• 数据中心 AI：强调数据质量而非数量的重要性，为主动学习、数据清洗、可解释数据贡献评估提供新工具。

总体而言，TAKE 不仅解决了文本蒸馏的实际问题，还深化了对训练数据贡献理解的理论框架，其思想可扩展到其他模态和任务。