CoCoMix：Meta提出新型Transformer预训练框架，性能优于知识蒸馏

AI资讯
1

Meta提出的CoCoMix预训练框架，结合token预测与连续概念，提升LLM性能与效率，尤其在弱监督场景表现突出。

原文标题：比知识蒸馏好用，田渊栋等提出连续概念混合，再度革新Transformer预训练框架

原文作者：数据派THU

冷月清谈：

Meta等机构的研究者提出了连续概念混合（CoCoMix）预训练框架，旨在通过结合离散的下一个 token 预测与连续概念，提升大型语言模型的性能和效率。CoCoMix 通过预训练的稀疏自编码器学习连续概念，并将其混合到模型的隐藏状态中，以此提高模型的推理能力和对概念的理解。实验结果表明，CoCoMix 在多个语言建模基准和不同规模的预训练模型上均表现出优越性，尤其是在采样效率和弱监督到强监督的场景中。与传统的下一个 token 预测和知识蒸馏相比，CoCoMix 能够以更少的训练 token 达到相当的性能，并能有效利用小模型中提取的概念来指导更大模型的训练。此外，CoCoMix 还具有良好的可解释性和可操纵性，能够帮助研究者分析模型关注的概念，并通过调整概念的大小来控制模型的输出生成。该研究为大型语言模型的预训练提供了一种新的有效方法。

怜星夜思：

1、CoCoMix通过混合连续概念来优化LLM预训练，那么这种方法相较于传统的token预测，在处理长文本和复杂推理任务时，有哪些具体的优势？
2、文章提到CoCoMix在弱监督到强监督的场景中表现出显著改进，那么在实际应用中，如何有效地利用小模型提取的概念来指导大模型的训练，避免出现负迁移或知识冲突的情况？
3、CoCoMix是如何实现模型的可解释性和可操纵性的？这种可解释性和可操纵性对于模型的debug和改进有什么帮助？

原文内容

来源：人工智能前沿讲习
本文约2500字，建议阅读5分钟
本文介介绍了Transformer预训练框架。


来自 Meta 等机构的研究者提出了一种新颖且高效的预训练框架：连续概念混合（Continuous Concept Mixing, CoCoMix），其将离散的下一个 token 预测与连续概念相结合。

近年来，大型语言模型（LLMs）的进展彻底改变了自然语言处理领域，并因此成为各种现实应用中的核心技术，例如代码助手、搜索引擎和个人 AI 助手。

这些突破的核心在于对「下一个 token 预测」的范式。

然而，自然语言 token 代表的意思通常是表层的（例如 the 或 a 这样的功能性词汇），需要模型进行大量训练才能获得高级推理和对概念的理解能力，同时也限制了它们处理长期任务（如规划）的能力。

为了解决这一问题，最近的研究探索了超越 token 层面信号的方法。例如有研究表明稀疏自编码器（Sparse Autoencoders, SAEs）能够通过捕捉高级语义概念，有效地分离出大型语言模型（LLMs）中有意义的潜在特征。

本文，来自 Meta 等机构的研究者提出了一种新颖且高效的预训练框架：连续概念混合（Continuous Concept Mixing, CoCoMix），其将离散的下一个 token 预测与连续概念相结合。

CoCoMix 用来预测从预训练的稀疏自编码器中学习到的连续概念，并通过与 token 隐藏表示交错的方式将其混合到模型的隐藏状态中。

具体来说，本文使用经过预训练的 SAE 提取语义概念，并根据归因（attribution）分数选择最具影响力的概念，这些分数量化了每个概念对模型输出的影响。然后，模型通过交叉熵损失训练，从其隐藏状态中预测这些选定的概念。一旦预测出多个概念，就将它们压缩为单个连续概念，并通过与 token 嵌入交错的方式混合（或插入）到隐藏状态中，从而直接贡献于下一个 token 的预测。

本文通过在多个语言建模基准和不同规模的预训练模型（从百万级到十亿级参数规模）上进行了评估，从而证明了 CoCoMix 的有效性。

结果表明，CoCoMix 采样效率更高，优于标准的下一个 token 预测、知识蒸馏以及插入停顿 token。本文发现，在端到端的框架中结合概念学习和交错技术对于性能提升至关重要。

CoCoMix 性能有多好？举例来说，当将 CoCoMix 应用于一个 1.38B 规模的模型时，CoCoMix 在减少 21.5% 训练 token 的情况下，实现了与下一个 token 预测相当的性能。

此外，CoCoMix 在弱监督到强监督的场景中表现出显著改进，其中从小模型中提取的概念甚至可以用作监督更大模型训练的标签。


图片
图片1080×703 112 KB


  • 论文标题：LLM Pretraining with Continuous Concepts
  • 论文地址：https://arxiv.org/pdf/2502.08524
  • 项目地址：https://github.com/facebookresearch/RAM/tree/main/projects/cocomix

此前，Meta 提出了大型概念模型（LCM），同样也是通过概念而非 token 进行学习和推理，被许多声音认为是大模型范式变革的新起点。现在看来，CoCoMix 出现，让 Meta 在创新算法逐步取代「连续预测下一个 token」这条道路上又迈出了重要的一步。

CoCoMix 介绍


CoCoMix 是一个使用连续概念扩展下一个 token 预测的框架。

其核心训练流程包括：一个概念选择框架（参见图1左），以及两个用于学习和利用连续概念的训练步骤（step）（参见图1右）。

  • 首先，本文使用归因分数选择重要概念，该分数衡量了每个概念对输出的影响。
  • 然后，本文提出通过交叉熵损失从模型的隐藏状态中预测选定的概念，使模型能够隐式地学习哪些概念应被编码为隐藏表示。
  • 最后，本文利用预测的概念创建一个连续概念，并将其交错插入到隐藏状态中，使模型能够显式地学习如何使用连续概念以及 token 隐藏状态。直观上，模型选择性地学习哪些概念对下一个 token 预测有用，以及如何将这些概念与 token 表示混合。



图片
图片1080×603 276 KB

实验


实验部分，研究者主要通过以下几个问题对 CoCoMix 进行了实证评估：

  • CoCoMix 能否提高 LLM 预训练中下一个 token 预测的性能？(图 2 和图 3）
  • 与其他知识提炼方法相比，CoCoMix 在弱到强监督设置中是否有所改进？(表 1 和图 4）
  • CoCoMix 是否引入了模型的可解释性和可操纵性？(图 5）
  • CoCoMix 的每个建议组件对性能有何贡献？(图 6）

首先是两个核心结果：

  1. 在相对大规模的预训练设置中与 NTP 的比较；
  2. 与 KD 基线的比较，尤其是在从小模型中提取的概念用于指导大模型的弱到强监督场景中。

大规模使用 CoCoMix 改进 NTP


如图 3 所示，CoCoMix 在各种规模的模型上都能持续显著提高下游任务的整体性能。结果还表明，较大的模型（如 386M 和 1.38B）可以从使用从较小的 124M 模型中提取的概念中获益，显示了有效的弱到强监督。


图片
图片1080×528 140 KB

如图 2 所示，在十亿级规模的模型上，CoCoMix 与 NTP 相比持续提高了性能。例如，CoCoMix 的性能与 NTP 相近，但使用的 token 却减少了 21.5%，显示了很高的采样效率。最后，值得注意的是，使用 CoCoMix 所获得的性能增益随着训练步骤的增加而增加，显示出很强的泛化性能。

图片
图片1080×440 150 KB

与 KD 基线的比较


研究者还将 CoCoMix 与 KD 基线在多种情况下进行了比较，包括：

  1. 较强的教师模型教授较小的学生模型；
  2. 弱到强的监督，即较弱的教师教授较大的学生模型；
  3. 分布转移，即学生在不同于教师预训练分布的语料库上接受训练。

如表 1 所示，在这些模型配置中，CoCoMix 都比 KD 有所改进。特别是，CoCoMix 在从弱到强的监督设置中表现出了显著的性能提升，例如在 386M 中平均复杂度提高了 2.8，而 KD 并没有表现出很大的改进。这是因为较弱的教师可能会引入噪声或次优知识，尤其是当学生的能力超过教师时。

图片
图片1080×450 172 KB

在图 4 中也可以观察到这一趋势，当学生的能力超过教师时（尤其是在分布转移场景中），使用 KD 训练的模型在训练中途会落后于标准训练。相比之下，CoCoMix 选择性地利用了有用的概念，从而实现了持续的性能提升。

图片
图片1080×460 158 KB

CoCoMix 的可解释性和可操纵性


CoCoMix 的另一个核心优势在于其可解释性和模型操纵。具体来说，由于模型经过训练后可以预测其隐藏状态中的概念，因此研究者可以根据概念预测结果来分析其重点关注的概念。此外，通过放大预测概念 z_t 的大小，可以控制模型的输出生成。

为了验证可操纵性是否达到预期效果，研究者在预训练模型的 SAE 潜在空间 c 中对同一概念的激活进行操纵，并确认输出是否表现出相应的概念。这里使用的是用 CoCoMix 训练的 386M 参数模型，其中预训练模型为 GPT-2。

如图 5 所示，当与「网站地址」相关的概念被放大时，两个模型都开始生成实际的网站地址。这表明本文模型成功地学习了 GPT-2 对齐概念。

图片
图片1080×443 182 KB

CoCoMix 有效性分析


随后，研究者对 CoCoMix 进行了详细分析，以验证每个提出的组件的效果。角度如下：

  1. 归因分数对选择概念的有效性；
  2. 概念预测与直接隐藏状态预测（即用连续损失预测隐藏状态，而不是用 SAE 将隐藏状态离散化）之间的比较；
  3. 压缩权重的稀疏性；
  4. 通过分析概念预测和混合的贡献进行成分分析；
  5. 通过比较将概念向量添加到原始隐藏状态和混合（将概念向量与 token 隐藏表示交错），对概念调节进行设计选择；
  6. CoCoMix 与 Pause token（即添加可学习 token）之间的比较。

此处，研究者使用了 69M Transformer，并在 OpenWebText 数据集中的 20B token 上进行训练。

编辑：王菁



关于我们

数据派THU作为数据科学类公众号，背靠清华大学大数据研究中心，分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识，努力建设数据人才聚集平台、打造中国大数据最强集团军。



图片
图片1080×564 45.8 KB

新浪微博：@数据派THU

微信视频号：数据派THU

今日头条：数据派THU

2

从理论上讲，CoCoMix的优势在于它能更好地捕捉token之间的语义关系。传统的token预测更多关注的是token的共现概率，而CoCoMix关注的是token背后的概念。对于长文本和复杂推理，模型需要理解token之间的语义关系，而不仅仅是知道它们经常一起出现。因此，CoCoMix更有可能提高模型在这方面的能力。

3

避免负迁移，可以考虑引入注意力机制，让大模型自主选择哪些概念是重要的，哪些概念是不重要的。这样，即使小模型提取的概念中存在一些噪声，大模型也可以通过注意力机制将其过滤掉。其次，可以采用对抗训练的方式，让大模型学习区分哪些概念是来自小模型，哪些概念是来自自身，从而提高模型的鲁棒性。