SimReg：通过嵌入相似性正则化在预训练阶段实现更高性能

# SimReg：通过嵌入相似度正则化在预训练中实现更高性能

来源：https://arxiv.org/html/2605.08809
Yan Sun1, Guoxia Wang1, Jinle Zeng1, JiaBin Yang1, Shuai Li1, Li Shen3, Dacheng Tao4, DianHai Yu1, Haifeng Wang1
1 百度公司 2 中山大学 3 南洋理工大学
{sunyan25, wangguoxia, zengjinle, yangjiabin01, yudianhai, wanghaifeng}@baidu.com
[email protected], [email protected], [email protected]

###### 摘要

使用下一词预测（next-token prediction）预训练大型语言模型（LLMs）带来了显著的进步，然而此类模型中词嵌入（token embeddings）的上下文依赖性导致了较高的类内方差和类间相似度，从而阻碍了表示学习的效率。虽然基于相似度的正则化在监督微调和分类任务中已显示出益处，但其在大规模 LLM 预训练中的应用和有效性仍鲜有探索。在本工作中，我们提出了 **SimReg**，这是一种嵌入相似度正则化损失，它显式地鼓励序列中具有相同真实标签的词表示更加相似，同时通过对比损失强制其与不同标签的词保持分离。我们的分析表明，这种机制通过扩大多分类边际（multi-classification margins）从而带来收益，进而实现更高效的分类。在稠密和混合专家（Mixture-of-Experts, MoE）架构上的广泛实验表明，SimReg  consistently 将训练收敛加速超过 30%，并在标准基准测试中将平均零样本（zero-shot）下游性能提升了超过 1%。进一步的消融研究和分析为超参数调整和损失有效性提供了实用的见解。

## 1 引言

LLMs 已成为现代人工智能的基石，并在自然语言理解（Radford et al., 2019）、推理（Wei et al., 2022）和多模态交互（Lin et al., 2025b）等广泛领域中展现出卓越的能力。尽管 LLMs 正沿不同方向进步，但它们从根本上共享一个一致的基本原则，即下一词预测。LLMs 的核心机制是从先前上下文的嵌入中预测下一个词的类别分布，这也可以看作是在前序上下文的组合表示上定义的一个分类问题。借助巨大的模型参数和庞大的训练数据，它表现出非凡的泛化能力，在多样化的研究领域中引入新的解决方案，并进一步推动了各类应用的采用（Topsakal and Akinci, 2023），同时也面临着日益增长效率挑战（Shen et al., 2024）。无论是针对数据的（Fan et al., 2025; Deng et al., 2026）还是针对权重的（Li et al., 2024a; Sun et al., 2025; Lin et al., 2025a）方法都吸引了大量的研究兴趣。

与传统的分类不同，语言模型预测并不依赖于严格绑定标签的稳定对象。例如，在图像分类中，猫的图片始终与其标签相关联，导致同一类内的嵌入具有高度一致性。相比之下，在语言任务中，用于预测某个词的表示由多样化的上下文特征组成，其中许多特征与标签本身无关。因此，预测相同词的嵌入可能会有显著差异。例如，“The cat jumps over walls”（猫跳过墙）和“A child paints near walls”（孩子在墙附近画画）中“walls”的表示源自完全不同的上下文，这使得分类过程更具挑战性。

最近针对语言模型微调的一致性学习进展为解决这一挑战提供了潜在的解决方案（Huang et al., 2021; Gunel et al., 2021; Yin et al., 2023）。然而，这一研究方向尚未扩展到预训练阶段，也未在大规模预训练实践中得到广泛采用。后训练（Post-training）通常以小学习率和有限的数据集进行，这使得显著改变学习参数的几何结构变得困难。这些见解促使我们将这种方法进一步扩展到大规模预训练中。

> **图 1**：（左）SimReg 损失的工作流程。（右）我们比较了通过“仅交叉熵”和“交叉熵+SimReg”训练的 LLaMA-7B 模型样本中词嵌入的余弦相似度。仅使用“交叉熵”无法在词特征之间施加足够的可分性，所有词对的余弦值均超过 0.5。引入 SimReg 后，特征可分性普遍增强（平均余弦值至少降低 0.1），从而为分类训练提供更强的支持。更多结果见附录 A.4。

在本工作中，我们发现仅使用交叉熵的大规模预训练未能对词嵌入施加强一致性。为此，我们添加了一个一致性正则化项 **SimReg**，以增强大模型在预训练期间的表示能力。对于序列中的每个词，所有词被划分为正组（positive groups）和负组（negative groups）。目标函数惩罚组间相似度，将嵌入拉向同类样本并推离不同类样本。为了确保每个词都有有效的对比对，SimReg 在每个正组中引入自样本相似度，并进一步计算基于组级别而非样本级别的平均损失，这平衡了不同词的贡献，使其在长期预训练运行中保持高度稳定性。我们还提供了全面的理论理解，以解释它如何有助于改善交叉熵损失。我们在稠密模型和 MoE 模型上进行了广泛评估，包括 LLaMA-350M、1.3B、3B、7B（Touvron et al., 2023）以及 Mixtral-8×1B（Jiang et al., 2024）。SimReg 损失在预训练中始终能将收敛速度加快超过 30%。当使用超过 52B 词元进行训练时，它能在下游通用任务的平均性能上带来超过 1% 的提升。我们研究了 SimReg 的超参数敏感性，发现其具有广泛的适用性。我们总结主要贡献如下：

-   • 我们探讨了在大规模预训练任务中采用一致性正则化的优势，并提出了一系列改进措施以解决现有方法的训练不稳定问题，从而在整个长期预训练过程中实现稳定的性能提升。
-   • 我们提供了关于 SimReg 损失对交叉熵损失益处的详细理论分析，以及它如何改善多分类边际。
-   • 我们进行了广泛的实验以验证其对预训练任务的显著改进，实现了超过 30% 的平均训练加速，并在下游任务上获得了超过 1% 的收益，为社区提供了详细的经验见解。

## 2 相关工作

**对比学习。** 机器学习中特征相似度约束的系统性探索可以追溯到其在计算机视觉（CV）任务和对比学习中的早期发展（Oord et al., 2018; Khosla et al., 2020）。它们通过构建虚拟数据对并加入额外的监督损失信号来增强基线分类模型的训练，这有助于模型提取更具判别性的特征。它通常用于对抗输入级的噪声扰动，从而提高泛化能力（Geng et al., 2021; Shi et al., 2022; Huang and Gong, 2022; Zhou et al., 2024; Wang et al., 2024）。通常，数据对由原始样本及其扰动副本构成，模型被训练以最小化它们的表示相似度。随后，监督对比学习被扩展以纳入类别信息。通过利用可用标签构建类别一致的数据对，模型不仅被训练将同类样本拉近，还将不同类样本推开（Wang and Liu, 2021; Wen and Li, 2021; Ye et al., 2022; Denize et al., 2023）。最近的研究表明，对比学习还可以在不同领域产生的任务和数据的跨域特征提取中实现更高的效率（Verma et al., 2021; Wang et al., 2022; Xie et al., 2022; Azuma et al., 2023）。在多模态大模型训练中，这种学习范式常被用于对齐跨领域的知识映射，并捕捉不同模态下相同知识的表示能力（Yuan et al., 2021; Mai et al., 2022; Liu et al., 2024b; Sun et al., 2024）。总之，对比学习为机器学习社区提供了一个高效且通用的表示学习范式。

**LLMs 中的嵌入一致性。** 特征相似性的研究也被视为组合泛化（Lake, 2019; Wiedemer et al., 2023）和嵌入一致性正则化（Yin et al., 2023）。Gao et al. (2021) 学习句子嵌入并实现了更高的泛化效率。随后，这一概念被广泛扩展到词元级别（Gao et al., 2023; Wang and Yu, 2023）、词级别（Kenter and De Rijke, 2015; Antoniak and Mimno, 2018）和上下文级别（Laskar et al., 2020）。大多数这些任务主要关注小规模或微调设置。作为现代语言模型的基石，下一词预测范式已广泛应用于各种下游任务（Li et al., 2024b; Chen et al., 2024）。最近的研究进一步调查了词嵌入的相似性和分散性，这突出了嵌入的可分性成为一个关键方向（de Andrade et al., 2023; Tao et al., 2024; Hu et al., 2024）。

## 3 问题设定与方法

在本节中，我们介绍如何将 SimReg 纳入 LLMs 的预训练，并解释它为何有助于提高性能。在此之前，我们形式化 LLMs 的整体预训练设置，并介绍后续分析中使用的符号。

**通用预训练。** 在介绍训练框架之前，我们首先定义本文中的符号。我们将 LLM 预训练的进展视为通过在一般分布 $\mathcal{D}$ 下最小化交叉熵损失 $\ell$ 来学习最优权重 $\mathbf{w}$。我们将模型分解为两个级联函数 $f_P \circ f_E$，其中 $f_P$（logits 生成模块）由 $\mathbf{w}_P$ 参数化，$f_E$（嵌入生成模块）由 $\mathbf{w}_E$ 参数化，整体参数表示为 $\mathbf{w}=[\mathbf{w}_P, \mathbf{w}_E]$。基于此分解，语言模型的通用预训练目标可以形式化为：

$$
\min_{\mathbf{w}} \mathbb{E}_{(\mathbf{x}_i, y_i) \sim \mathcal{D}} \left[ \ell(f_P \circ f_E(\mathbf{x}_i), y_i) \right], \quad (1)
$$

其中 $(\mathbf{x}_i, y_i)$ 是从分布 $\mathcal{D}$ 中采样的（数据，标签）对。这里，$f_E$ 和 $f_P$ 的选择完全灵活，意味着 SimReg 损失原则上可以应用于网络中任何有效的词嵌入。我们在第 5.3 节的后续实验中进一步探索了该组件的最佳放置位置。

交叉熵损失作为语言建模中的基本训练目标。它衡量预测的词分布与真实 one-hot 分布之间的差异，从而指导模型最大化正确下一个词的概率。模型通常采用大规模特征提取器以获得可分离的表示。记词嵌入为 $\mathbf{e}_i = f_E(\mathbf{x}_i)$，对应的 logits 为 $\mathbf{z}_i = f_P(\mathbf{e}_i)$，样本级交叉熵损失的群体风险为：

$$
L_{ce} = \frac{1}{n} \sum_i \left( -z_{i, y_i} + \log \left( \sum_j \exp(z_{i, j}) \right) \right). \quad (2)
$$

> **图 2**：（a）我们分析了 C4 数据集中 10 亿训练样本的词 ID 分布，发现只有约 2% 的词以极高频率出现，导致分类数据中出现明显的长尾效应。（b）我们观察到，嵌入的对比相似度损失在达到基本阈值后不再继续下降，此后特征相似度不再进一步优化。仅仅增加模型规模并不能改善这一性能。通常，更大的可分性可以增强不同样本之间的区分度，从而产生更鲁棒和更具判别性的表示。虽然公式 (2) 对样本进行了平均，但语言任务的独特特性引入了一项挑战：词（词元）的分布高度不平衡，导致频繁出...