[论文解读] Learning by Distilling Context

Charlie Snell, Dan Klein+1 more|arXiv (Cornell University)|Sep 30, 2022

Topic Modeling被引用 7

一句话总结

本文提出上下文蒸馏（context distillation）方法，通过在自动生成的补全结果上进行微调，使语言模型能够内化上下文标记（如指令、草稿、示例）带来的性能提升。该方法在 SPIDER Text-to-SQL 任务上比梯度下降高出 9% 的准确率，并实现了将逐步推理能力迁移到下游任务的能力。

ABSTRACT

Language models significantly benefit from context tokens, such as prompts or scratchpads. They perform better when prompted with informative instructions, and they acquire new reasoning capabilities by generating a scratch-pad before predicting the final answers. However, they do not extit{internalize} these performance gains, which disappear when the context tokens are gone. Our work proposes to apply context distillation so that a language model can improve itself by internalizing these gains. Concretely, given a synthetic unlabeled input for the target task, we condition the model on ``[instructions] + [task-input]'' to predict ``[scratch-pad] + [final answer]''; then we fine-tune the same model to predict its own ``[final answer]'' conditioned on the ``[task-input]'', without seeing the ``[instructions]'' or using the ``[scratch-pad]''. We show that context distillation is a general method to train language models, and it can effectively internalize 3 types of training signals. First, it can internalize abstract task instructions and explanations, so we can iteratively update the model parameters with new instructions and overwrite old ones. Second, it can internalize step-by-step reasoning for complex tasks (e.g., 8-digit addition), and such a newly acquired capability proves to be useful for other downstream tasks. Finally, it can internalize concrete training examples, and it outperforms directly learning with gradient descent by 9\% on the SPIDER Text-to-SQL dataset; furthermore, combining context distillation operations can internalize more training examples than the context window size allows.

研究动机与目标

• 解决语言模型在移除上下文标记（如指令、草稿）后性能下降的问题。
• 使语言模型能够内化复杂能力（如逐步推理和特定任务指令），而无需依赖外部上下文。
• 通过迭代蒸馏克服上下文窗口限制，将超过上下文窗口容量的更多示例内化。
• 开发一种通用训练方法，在提升模型性能的同时保留其原有能力（如遵循指令）。

提出的方法

• 使用教师模型根据包含指令和任务输入的详细提示生成草稿和最终答案。
• 从教师模型的生成结果中提取最终答案，并微调同一模型，使其仅通过简化的学生提示直接预测该答案。
• 对三类信号应用蒸馏：抽象任务指令、具体示例（如 SPIDER 中的示例）以及逐步推理（如 8 位数加法）。
• 通过混合使用蒸馏数据和少样本学习示例，防止微调过程中的灾难性遗忘。
• 通过迭代应用蒸馏，使模型内化超过上下文窗口容量的更多示例，并观察到性能持续提升。
• 通过在相关下游任务（如涉及加法的字谜题）上测试蒸馏后的推理能力，评估其可迁移性。

实验结果

研究问题

• RQ1 语言模型能否通过上下文蒸馏内化抽象任务指令和解释，并支持迭代更新？
• RQ2 如 8 位数加法等逐步推理能力能否被有效内化并迁移到其他推理任务中？
• RQ3 在上下文窗口受限的情况下，上下文蒸馏能否在学习具体示例时优于直接的梯度下降微调？
• RQ4 结合多个蒸馏步骤是否能使模型内化超过上下文窗口大小限制的训练样本？
• RQ5 内化后的推理能力能否泛化到需要类似认知操作的新任务上？

主要发现

• 在 SPIDER Text-to-SQL 数据集上，对于 8-shot 微调，上下文蒸馏使性能比直接梯度下降微调高出 9%。
• 在对 8 位数加法进行逐步推理蒸馏后，模型的直接加法准确率从 1% 提升至 17%，表明实现了有效内化。
• 蒸馏后的模型在指令遵循任务上保持了强劲性能（RougeL: 57 → 58），表明未发生原始能力的灾难性遗忘。
• 蒸馏后的加法能力成功迁移到相关推理任务，在涉及加法的合成字谜题中性能从 17% 提升至 30%。
• 迭代式上下文蒸馏使模型能够内化超过上下文窗口容量的训练样本，且性能随蒸馏步骤增加而持续提升。
• 该方法在不同类型训练信号上均表现出泛化能力：包括抽象指令、具体示例和结构化推理，显示出广泛的适用性。