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

KzWcx6mFqHAy · 2025-2-17 15:59:39

近年来，大型语言模型（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）

首先是两个核心结果：

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

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

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

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

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

如表 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 进行了详细分析，以验证每个提出的组件的效果。角度如下：

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

此处，研究者使用了 69M Transformer，并在 OpenWebText 数据集中的 20B token 上进行训练。
更多研究细节，可参考原论文。