超越Token预测：CoCoMix如何革新Transformer预训练

AI资讯
1

Meta提出CoCoMix，结合连续概念和token预测，提升Transformer预训练效率和性能，实现弱监督到强监督的改进。

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

原文作者：机器之心

冷月清谈：

Meta的研究人员提出了一种名为CoCoMix（连续概念混合）的新型预训练框架，该框架结合了离散的下一个token预测和连续概念学习。CoCoMix利用预训练的稀疏自编码器（SAE）提取语义概念，并根据概念对模型输出的影响程度进行选择。然后，模型通过预测这些选定的概念进行训练，并将预测的概念压缩成单个连续概念，将其与token嵌入交错混合到隐藏状态中，从而直接影响下一个token的预测。

实验结果表明，CoCoMix在各种规模的模型上都提高了性能和采样效率。例如，在13.8亿参数的模型上，CoCoMix在减少21.5%训练token的情况下实现了与标准的下一个token预测相当的性能。此外，CoCoMix在弱监督到强监督的场景中表现出色，即使使用较小模型提取的概念来训练较大模型也能取得显著改进。
CoCoMix的核心优势在于其可解释性和可操作性。通过分析概念预测结果，可以理解模型关注的概念。此外，通过放大预测概念的大小，可以控制模型的输出生成。

怜星夜思：

1、CoCoMix 中使用的稀疏自编码器（SAE）是如何预训练的？预训练的数据集和方法对 CoCoMix 的最终效果有什么影响？
2、CoCoMix 提到了弱监督到强监督的场景，具体来说，这种场景有哪些实际应用？
3、除了语言模型预训练，CoCoMix 的这种连续概念混合的思想还可以应用到哪些其他领域？

原文内容

机器之心报道
编辑：蛋酱、陈陈

近年来，大型语言模型（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 在弱监督到强监督的场景中表现出显著改进，其中从小模型中提取的概念甚至可以用作监督更大模型训练的标签。


  • 论文标题：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 表示混合。


实验

实验部分，研究者主要通过以下几个问题对 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 模型中提取的概念中获益，显示了有效的弱到强监督。


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


与 KD 基线的比较

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

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

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


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


CoCoMix 的可解释性和可操纵性

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

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

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


CoCoMix 有效性分析

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

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

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

更多研究细节，可参考原论文。

© THE END 
转载请联系本公众号获得授权
投稿或寻求报道：[email protected]

2

弱监督到强监督还可以用在一些数据隐私比较敏感的场景，比如用公开数据集训练的小模型去指导基于私有数据集的大模型训练，这样可以避免直接使用私有数据。

3

弱监督到强监督在实际应用中很有潜力。比如，我们可以用少量标注数据训练一个小模型，然后用这个小模型提取的概念去指导大模型的训练，这样可以节省大量的标注成本。

4

在推荐系统中，可以用用户行为数据训练一个模型提取用户的兴趣概念，然后用 CoCoMix 将这些概念混合到推荐模型中，可能会提升推荐效果。

5

这个论文里好像没有具体说SAE的预训练细节，可能需要去看他们引用的相关论文。不过我感觉用不同的数据集预训练SAE，应该可以让CoCoMix适应不同的下游任务。

6

预训练数据肯定很重要啊，高质量的数据才能训练出好的SAE，进而提升CoCoMix的效果。估计他们用了很大的数据集，说不定也用了些数据增强之类的技巧。

7

我觉得可以应用到图像识别领域，比如用一些预训练的模型提取图像的特征，然后用 CoCoMix 的方法将这些特征混合到图像分类模型中。

8

CoCoMix 的思想还可以应用到强化学习领域，比如用一些专家经验提取一些状态-动作的概念，然后用 CoCoMix 将这些概念混合到强化学习agent的策略网络中。

9

我觉得在一些数据标注成本很高的领域，比如医学图像分析，CoCoMix 的这种弱监督学习方法应该很有用。

10

SAE的预训练应该是比较关键的一环，如果SAE本身学到的概念不够好，可能会影响CoCoMix的效果。我猜想预训练数据集的质量和大小很重要，另外SAE的架构和训练方法也应该会有影响。