Choosy Babies Need One Coach: Inducing Mode-Seeking Behavior in BabyLlama with Reverse KL Divergence

📄 arXiv: 2410.22081v1 📥 PDF

作者: Shaozhen Shi, Yevgen Matusevych, Malvina Nissim

分类: cs.CL

发布日期: 2024-10-29

💡 一句话要点

利用反向KL散度诱导BabyLlama的Mode-Seeking行为，单教师优于多教师。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 蒸馏训练 反向KL散度 Mode-Seeking BabyLM 单教师模型 知识迁移 语言模型优化

📋 核心要点

  1. 现有BabyLM模型的蒸馏训练通常采用多教师集成，但可能导致学生模型学习到平均化的知识。
  2. 论文提出使用反向KL散度作为目标函数，鼓励学生模型学习教师模型的特定模式，实现更专注的学习。
  3. 实验结果表明，在反向KL散度下，单教师模型在多数任务上表现优于或等同于多教师模型，且优化技术进一步提升性能。

📝 摘要（中文）

本研究介绍了我们在第二届BabyLM挑战赛Strict-Small赛道的提交方案。我们采用教师-学生蒸馏框架，以BabyLLaMa模型为基础。为了使学生的学习过程更加集中，我们将目标函数替换为反向Kullback-Leibler散度，该散度已知会导致计算学习器中出现mode-seeking（而非mode-averaging）行为。我们进一步试验了使用单个教师（而不是两个教师的集成），并实施了额外的优化策略来改进蒸馏过程。实验表明，在反向KL散度下，单教师模型在大多数任务中通常优于或匹配多教师模型。此外，结合先进的优化技术进一步提高了模型性能，证明了我们提出的方法的有效性和鲁棒性。这些发现支持了我们的观点，即“挑剔的婴儿需要一位教练”。

🔬 方法详解

问题定义：论文旨在提升BabyLLaMa模型在BabyLM挑战赛中的性能。现有方法，特别是基于多教师集成的蒸馏方法，可能导致学生模型学习到教师模型的平均行为，无法有效捕捉教师模型的特定优势。这种mode-averaging行为限制了学生模型的性能。

核心思路：论文的核心思路是利用反向Kullback-Leibler (KL) 散度作为蒸馏的目标函数。与传统的KL散度不同，反向KL散度倾向于使学生模型的分布覆盖教师模型分布的特定模式（mode-seeking），而不是平均所有模式。这使得学生模型能够更专注于学习教师模型的优势。

技术框架：整体框架是教师-学生蒸馏。首先，使用一个或多个教师模型（BabyLLaMa）进行训练。然后，使用反向KL散度作为损失函数，训练学生模型（BabyLLaMa）模仿教师模型的行为。此外，论文还探索了不同的优化策略来进一步提升蒸馏效果。

关键创新：关键创新在于使用反向KL散度作为蒸馏目标函数，从而诱导学生模型产生mode-seeking行为。这与传统的蒸馏方法中使用的KL散度不同，后者倾向于产生mode-averaging行为。此外，论文还探索了单教师蒸馏策略，发现其在反向KL散度下可以取得比多教师更好的效果。

关键设计：关键设计包括：1) 使用反向KL散度作为损失函数，具体形式未知（论文未给出公式）；2) 探索单教师与多教师的配置；3) 采用额外的优化策略，具体细节未知（论文未详细说明）。BabyLLaMa模型本身的网络结构和参数设置沿用了原始论文。

📊 实验亮点

实验结果表明，在反向KL散度下，单教师模型在大多数任务中表现优于或匹配多教师模型。此外，结合先进的优化技术进一步提高了模型性能，证明了该方法的有效性和鲁棒性。具体的性能提升数据和对比基线未知（论文摘要未提供具体数值）。

🎯 应用场景

该研究成果可应用于提升小型语言模型的性能，尤其是在数据有限或计算资源受限的场景下。通过反向KL散度蒸馏，可以使小型模型更有效地学习大型模型的知识，从而在各种自然语言处理任务中获得更好的表现。该方法也可能推广到其他机器学习领域，例如图像识别和语音识别。

📄 摘要（原文）

This study presents our submission to the Strict-Small Track of the 2nd BabyLM Challenge. We use a teacher-student distillation setup with the BabyLLaMa model (Timiryasov and Tastet, 2023) as a backbone. To make the student's learning process more focused, we replace the objective function with a reverse Kullback-Leibler divergence, known to cause mode-seeking (rather than mode-averaging) behaviour in computational learners. We further experiment with having a single teacher (instead of an ensemble of two teachers) and implement additional optimization strategies to improve the distillation process. Our experiments show that under reverse KL divergence, a single-teacher model often outperforms or matches multiple-teacher models across most tasks. Additionally, incorporating advanced optimization techniques further enhances model performance, demonstrating the effectiveness and robustness of our proposed approach. These findings support our idea that "choosy babies need one coach".