模型预训练模式“变天”？Meta推出预训练框架，训练token减少21.5%
模型预训练模式“变天”？Meta推出预训练框架，训练token减少21.5%
Meta公司近日推出了一种名为“连续概念混合（CoCoMix）”的预训练框架，该框架通过预测从预训练稀疏自编码器中学习到的连续概念，并将这些概念混入模型的隐藏状态中，实现了更高的样本效率。实验结果显示，使用CoCoMix训练的模型在多个语言建模基准测试中表现出色，特别是在十亿级模型上持续提高了“预测下一个token”的性能。
“预测下一个token”（NTP，Next Token Prediction），最早由美国数学家克劳德·艾尔伍德·香农（Claude Elwood Shannon）于 1948 年在《通信的数学理论》一书中提出。
它指的是通过观察一系列给定的 token，预测序列中的下一个 token。如今，“预测下一个 token”已经成为自然语言处理的核心机制。
OpenAI 前首席科学家、联合创始人伊利亚·苏茨克维（Ilya Sutskever）曾多次表示，只要能够预测下一个 token，就能帮助人类达到通用人工智能。
近日，通过结合“预测下一个 token”概念与连续概念，Meta 提出一款名为“连续概念混合（CoCoMix，Continuous Concept Mixing）”的预训练框架。（注：连续概念，指的是连续思维链技术，这是一种旨在提升模型推理能力的方法。）
图 | CoCoMix 的相关论文的第一作者吉洪·塔克（Jihoon Tack）（来源：https://jihoontack.github.io/）
CoCoMix 能够预测那些从预训练的稀疏自编码器中学习到的连续概念，并能通过与 token 隐藏表示交错的方式，将这些概念混入模型的隐藏状态中。实验结果显示，CoCoMix 的样本效率更高，并且始终优于“预测下一个 token”、知识蒸馏和插入暂停标记（inserting pause tokens）。（注：插入暂停标记，是一种训练大型语言模型的方法，旨在模拟人类在思考过程中的停顿，从而提升模型的生成文本质量和逻辑性。）
实验中，研究人员通过在多个语言建模基准测试上进行评估，以及在规模从百万级参数到十亿级参数不等的预训练模型上进行评估，借此证明了 CoCoMix 的有效性。
例如，当将 CoCoMix 用于 1.38B 大小的模型时，CoCoMix 实现了与“预测下一个 token”相当的性能，但是训练 token 能够减少 21.5%。
此外，在从弱监督到强监督的场景中，CoCoMix 从小型模型中提取的概念，甚至可以作为真实标签来监督较大模型的训练。
并且，CoCoMix 能够通过直接检查和修改预测概念来增强可解释性和可操纵性，从而能为指导模型的内部推理过程带来帮助。
图 | CoCoMix 的概况（来源：arXiv）
CoCoMix：在十亿级模型上持续提高“预测下一个 token”的性能
研究人员表示，在打造 CoCoMix 的时候他们使用预先训练的稀疏自编码器来提取语义概念，并根据归因分数选择最具影响力的概念，以此来量化每个概念对于模型输出的影响。
然后，研究人员开始进行模型训练，并使用交叉熵损失从其隐藏状态来预测被选定的概念。一旦预测出多个概念，研究人员就将它们压缩成一个连续的概念，并通过与 token 嵌入交错混合到隐藏状态之中，从而直接为“预测下一个 token”带来助力。
这样做的另一个好处，可以针对预测出的概念进行分析，从而实现内容的可控生成以及提高模型的可解释性。
为了针对 CoCoMix 进行实证评估，研究人员提出了以下研究问题：首先，CoCoMix 能否提高大模型预训练中“预测下一个 token”的性能？其次，与其他知识提取方法相比，CoCoMix 在从弱到强的监督设置上是否有所改善？其次，CoCoMix 是否引入了模型可解释性和可操纵性？最后，CoCoMix 的每个拟议组件对性能有何贡献？（注：知识提取，是指从非结构化或半结构化的数据源中自动提取结构化信息的过程。拟议组件，指的是正在讨论或计划中的组件，这些组件可能尚未完全实现或标准化，但被认为是终版模型的重要组成部分。）
针对此，在训练设置上，研究人员使用了在 124M 大小的 GPT-2 上进行训练的预训练开源稀疏自编码器。与此同时，他们使用三个不同数量的活动参数来训练 CoCoMix，包括 69M、386M 和 1.38B，上下文长度则为 1024。
由于连续概念的交错，CoCoMix 使用的浮点运算比“暂停 token”少，但是比“预测下一个 token”要多。（注：暂停 token，是一种特殊的标记，用于在模型的预训练和微调过程中引入可学习的“暂停”机制，旨在提高 Transformer 模型在处理多样化任务时的性能。）
研究人员表示，其中一个实验使用了 200B 个 token 外，其余所有实验均使用 20B 个训练 token 进行。在评估 CoCoMix 时，对于所使用的基准模型或基准数据集，他们使用了标准的“预测下一个 token”预训练程序，以及使用了在预训练中常见的知识蒸馏。
实验中，他们排除了那些需要训练多个模型的知识蒸馏基准模型。对于知识蒸馏，研究人员在平衡 KL 散度与“预测下一个 token”损失的同时，针对教师模型和学生模型输出之间的 KL 散度进行最小化。（注：KL 散度，是一种衡量两个概率分布之间差异的指标。）
在评估设置上，研究人员使用验证困惑度的预训练数据集和 7 个下游任务，以便针对模型的常识推理能力和阅读理解能力进行基准测试。与此同时，他们将 OpenWebMath 作为一个预训练数据集，以便证明即使在一个不同的语料库上训练模型，那些从一个预训练语言模型中学到的概念仍然可以用于 CoCoMix。
另外，研究人员还展示了以下两个核心结果：第一，他们展示了在相对大规模的预训练设置中与“预测下一个 token”的比较；第二，他们展示了与知识蒸馏基线的比较。特别是在弱到强的监督场景中，研究人员证明从小模型中提取的概念可被用于指导更大的模型。
与此同时，研究人员在 200B token 上训练“预测下一个 token”和 CoCoMix。如下图所示，CoCoMix 始终能够显著提高各种尺寸模型的下游任务性能。
（来源：arXiv）
实验结果表明，从较小的 124M 模型中提取的概念，能让 386M 和 1.38B 等较大模型受益，从而显示出从弱到强的监督有效性。此外，如下图所示，CoCoMix 在十亿级模型上持续提高了“预测下一个 token”的性能。同时，使用 CoCoMix 所带来的性能增益，会随着训练步骤的增加而增加，也就是说它拥有很强的泛化性能。
（来源：arXiv）
研究人员还比较了 CoCoMix 和知识蒸馏基线在多个场景中的表现。这些场景包括：让更强的教师模型教导更小的学生模型的场景；从弱监督到强监督的场景，即让较弱的教师教导较强的学生模型；分布偏移场景，即让学生模型在与教师模型的预训练分布在不同的语料上进行训练。如下表所示，CoCoMix 在所有模型配置中均比知识蒸馏有所改进。
（来源：arXiv）
特别是，在弱监督到强监督的设置中，CoCoMix 表现出显著的性能提升。例如，在 3.86 亿参数规模的模型中，将平均困惑度降低了 2.8，而知识蒸馏方法并没有显示出较大改进。这是因为较弱的教师模型可能会引入嘈杂知识或次优知识，尤其是当学生模型的能力超过教师模型时。
在下图中也可以观察到这一趋势：使用知识蒸馏训练的模型，在训练中途落后于标准训练，这是因为学生模型的速度超过了教师模型。相比之下，CoCoMix 能够有选择性地利用有用的概念，从而获得一致的性能提升。
（来源：arXiv）
CoCoMix 的另一个核心优势是其可解释性和模型导向性。具体来说，当模型被训练用来预测隐藏状态下的概念时，可以根据概念预测来分析它关注的概念。
为了验证这种可操纵性是否可以按照预期工作，研究人员在预训练模型的稀疏自编码器潜在空间中引导相同概念的激活，并确认输出是否表现出相应的概念。
研究人员使用用 CoCoMix 训练的 386M 参数模型，其中预训练模型是 GPT-2。如下图所示，当与“网站地址”相关的概念被放大时，这两个模型都开始生成实际的网站地址。这表明研究人员的模型已经成功地学习了 GPT-2 对齐的概念。
（来源：arXiv）
与此同时，研究人员还对 CoCoMix 进行了详细分析，以验证每个拟议组件的效果。期间，他们使用了 69M 模型，针对从 OpenWebText 数据集中采样的 20B token 进行训练。
在归因得分有效性的分析上，他们先是分析归因得分是否有效地提取了重要概念。为了证明这一点，研究人员使用概念提取的激活值 ct 来训练 CoCoMix。之所以这样做，是因为激活值能够很好地反映这一概念的重要性。如下图所示，使用
归因得分显著提高了性能，与基于激活值的选择相比，样本效率提高了 17.5%。
（来源：arXiv）
在组件的分析上，研究人员分析了本次方法之中每个组成部分的贡献：首先，分析概念预测方程 2 的贡献；其次，分析概念插入方程的贡献。下图中的结果表明，这两个组件对于性能改进都至关重要。
具体来说，仅仅使用预测损失的概念就可以适度减少困惑。同时，将概念插入与预测相结合，可以提高损失函数的有效性，从而能够实现进一步的性能提升。这显示插入让模型能够有效地利用预训练的大模型潜在推理。
（来源：arXiv）
而对于未来的探索方向，研究人员认为通过其他选择标准来提升 CoCoMix 的性能，或者通过去除不良概念以减少偏差，将会是非常有价值的探索方向。
参考资料：
https://arxiv.org/pdf/2502.08524