TS-Align: A Teacher-Student Collaborative Framework for Scalable Iterative Finetuning of Large Language Models

📄 arXiv: 2405.20215v4 📥 PDF

作者: Chen Zhang, Chengguang Tang, Dading Chong, Ke Shi, Guohua Tang, Feng Jiang, Haizhou Li

分类: cs.CL

发布日期: 2024-05-30 (更新: 2024-09-29)

备注: EMNLP-2024 Findings

💡 一句话要点

TS-Align：一种可扩展的迭代微调框架，用于对齐大型语言模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型对齐 师生学习 迭代微调 强化学习 奖励模型 偏好学习 自动化数据挖掘

📋 核心要点

  1. 现有LLM对齐方法依赖大量人工标注偏好数据，成本高且难以扩展，尤其是在模型需要定期更新时。
  2. TS-Align框架利用教师模型和学生模型协作，自动挖掘模型输出中的成对反馈数据，用于策略模型的微调。
  3. 实验表明，TS-Align框架对齐后的策略模型在多个数据集上显著优于基线模型，且教师模型的排序能力成功迁移到学生模型。

📝 摘要（中文）

本文提出了一种名为“TS-Align”的框架，旨在解决大型语言模型（LLMs）迭代对齐过程中人工标注数据成本高昂且难以扩展的问题。该框架通过利用大规模教师模型和小规模学生模型之间的协作，自动挖掘LLM输出中的成对反馈数据，并以此微调策略模型。这种自动挖掘过程高效且可扩展。通过在提出的师生协作框架内进行基于策略生成的迭代重复微调，实验结果表明，最终对齐的策略模型在七个对话或指令跟随数据集上的平均胜率为69.7%，优于基础策略模型。此外，教师模型的排序能力有效地提炼到学生模型中，从而产生了一个小规模但有效的奖励模型，用于策略模型的对齐。

🔬 方法详解

问题定义：现有大型语言模型（LLMs）的对齐通常依赖于人工标注的偏好数据，特别是当模型需要定期更新时。然而，人工标注数据的收集成本高昂，且难以扩展，这限制了LLM迭代对齐的效率和可行性。因此，如何降低对齐过程中的数据标注成本，实现LLM的可扩展迭代对齐是一个关键问题。

核心思路：TS-Align的核心思路是利用一个大规模的教师模型和一个小规模的学生模型进行协作，自动挖掘LLM输出中的成对反馈数据。教师模型作为知识的提供者，负责评估和排序学生模型的输出，从而生成伪标签数据。学生模型则通过学习这些伪标签数据来提升自身的性能，并最终作为策略模型进行部署。这种师生协作的方式避免了人工标注的成本，并实现了数据的自动生成和模型的迭代优化。

技术框架：TS-Align框架主要包含以下几个阶段：1) 数据生成阶段：学生模型生成一系列输出，教师模型对这些输出进行评估和排序，生成成对的偏好数据。2) 策略微调阶段：利用生成的成对偏好数据，对学生模型进行策略微调，使其更好地符合人类的偏好。3) 迭代优化阶段：重复数据生成和策略微调的过程，不断提升学生模型的性能。整个框架采用on-policy生成的方式，保证了数据分布的一致性。

关键创新：TS-Align最重要的创新点在于利用师生协作的方式，实现了LLM对齐数据的自动生成，从而避免了人工标注的成本。与传统的依赖人工标注的方法相比，TS-Align具有更高的可扩展性和效率。此外，通过将教师模型的排序能力提炼到学生模型中，可以得到一个有效的奖励模型，用于策略模型的对齐。

关键设计：在数据生成阶段，教师模型使用其强大的生成和评估能力，对学生模型的输出进行排序，生成成对的偏好数据。在策略微调阶段，可以使用常见的强化学习算法，如PPO（Proximal Policy Optimization），对学生模型进行微调。损失函数的设计需要考虑到偏好数据的特点，例如可以使用pairwise ranking loss。学生模型的规模可以根据实际需求进行调整，但通常需要小于教师模型，以降低计算成本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TS-Align框架对齐后的策略模型在七个对话或指令跟随数据集上的平均胜率为69.7%，显著优于基础策略模型。这表明TS-Align框架能够有效地提升LLM的性能，使其更好地符合人类的偏好。此外，实验还验证了教师模型的排序能力可以有效地提炼到学生模型中，从而产生一个有效的奖励模型。

🎯 应用场景

TS-Align框架可应用于各种需要持续对齐的大型语言模型，例如对话系统、智能助手、内容生成工具等。通过降低对齐成本，该框架可以加速LLM的迭代更新和优化，使其更好地服务于人类。此外，该框架还可以用于训练小规模但高性能的奖励模型，用于指导LLM的训练。

📄 摘要（原文）

Mainstream approaches to aligning large language models (LLMs) heavily rely on human preference data, particularly when models require periodic updates. The standard process for iterative alignment of LLMs involves collecting new human feedback for each update. However, the data collection process is costly and challenging to scale. To address this issue, we introduce the "TS-Align" framework, which fine-tunes a policy model using pairwise feedback data automatically mined from its outputs. This automatic mining process is efficiently accomplished through the collaboration between a large-scale teacher model and a small-scale student model. The policy fine-tuning process can be iteratively repeated using on-policy generations within our proposed teacher-student collaborative framework. Through extensive experiments, we demonstrate that our final aligned policy outperforms the base policy model with an average win rate of 69.7% across seven conversational or instruction-following datasets. Furthermore, we show that the ranking capability of the teacher is effectively distilled into the student through our pipeline, resulting in a small-scale yet effective reward model for policy model alignment.