在自然语言处理 (NLP) 快速发展的当下，我们正面临一个“双模型问题”。如果你正在构建一个检索增强生成 (RAG) 系统，通常需要两种截然不同的架构: 一个检索器 (通常是像 BERT 这样的双向编码器) 来处理嵌入和搜索，以及一个生成器 (像 GPT 或 Llama 这样的仅解码器 LLM) 来合成答案。

这种结构上的分离导致了低效。它不仅使部署成本加倍，还阻碍了任务之间的知识共享。如果单个大型语言模型 (LLM) 能够同时处理高质量的文本生成和高质量的句子表示，那会怎样？

这就是 UniMAE 的承诺，这篇名为 “Decoder-Only LLMs can be Masked Auto-Encoders” (仅解码器 LLM 可以成为掩码自编码器) 的论文提出了一种新颖的无监督训练方法。在这篇文章中，我们将详细拆解 UniMAE 如何强制标准的仅解码器 LLM 学习强大的句子嵌入，同时不丧失其生成流畅文本的能力。

问题所在: 为什么 LLM 难以处理嵌入

要理解 UniMAE，我们首先需要了解编码器和解码器之间的架构差异。

编码器模型 (例如 BERT) 使用双向注意力机制。句子中的每个 Token 都能同时“看到”所有其他 Token。它们通常使用一个特殊的 Token (如 [CLS]) 来聚合整个句子的含义，这使得它们非常适合创建用于搜索和聚类的向量嵌入。

仅解码器模型 (例如 Llama, GPT) 使用单向 (因果) 注意力机制。一个 Token 只能关注它之前的 Token。它们被训练用来预测序列中的下一个 Token。因此，它们缺乏一种专门的机制来表示整个句子。如果你简单地对 LLM 的所有 Token 嵌入取平均值，其质量通常远不如专用的编码器。

最近修复这个问题的尝试包括“魔改”LLM——要么平均其输出，要么在微调期间强制使用双向注意力。然而，这些方法通常会削弱模型的主要优势: 文本生成能力。

解决方案: UniMAE

研究人员提出了 UniMAE (单向掩码自编码器) 。其核心理念巧妙而直接: 强制模型将句子的所有语义信息压缩到最终的 [EOS] (序列结束) 标记中。

他们是如何做到的？通过训练模型仅使用该 [EOS] 嵌入来重建原始句子。

图 1: UniMAE 训练概览。左侧显示掩码自回归过程，右侧显示使用 [EOS] 嵌入恢复输入的微型解码器重建过程。

如上图 1 所示，该框架包含两个并行的训练目标:

  1. 掩码自回归 (MAR): LLM 生成文本，但有一个转折——部分输入被掩盖了。
  2. 掩码重建 (MRC): 一个微型、独立的解码器尝试使用 LLM 的 [EOS] Token 中的表示来重建句子。

让我们深入探讨这两个组件。

1. 掩码自回归 (MAR)

标准的 LLM 是通过自回归 (AR) 学习进行训练的: 给定一系列单词，预测下一个单词。UniMAE 采用了一种稍难的版本，称为掩码自回归 (MAR) 。

在 MAR 中，输入句子 \(X\) 通过随机掩盖一定比例的 Token 被破坏 (表示为 \(\tilde{X}\)) 。LLM 仍然必须预测正确的序列。

MAR 的损失函数计算如下:

() L _ { M A R } = \\sum _ { t } C E ( x _ { t } \\mid \\tilde { x } _ { i < t } ; \\Phi _ { l l m } ) , ()

在这里，模型尝试基于被扰动的先前 Token \(\tilde{x}_{i

为什么要这样做? 如果我们将干净的句子输入给模型，它可能只是简单地记住序列。通过引入噪声 (掩码) ，我们强制模型依赖语义上下文，而不是短程的相关性。这确保了模型实际上是在“思考”内容，这对于构建稳健的表示至关重要。

2. 掩码重建 (MRC)

这是 UniMAE 创新的核心。目标是让 [EOS] Token 的嵌入 (表示为 \(h_{eos}\)) 成为整个句子含义的容器。

首先，LLM 处理被掩盖的输入 \(\tilde{X}\)，以在序列末尾产生潜在表示:

() h _ { e o s } \\gets \\Phi _ { l l m } \\left( \\tilde { X } \\right) ()

由于解码器模型中的因果注意力掩码，[EOS] Token 是唯一关注了所有先前 Token 的位置。然而，如果没有激励机制，[EOS] Token 没有理由去“记住”句子的开头——它只关心接下来的内容 (通常什么都没有) 。

为了解决这个问题，研究人员初始化了一个微型解码器 (Tiny Decoder) 。 这是一个仅在训练期间使用的小型临时神经网络。它的工作是获取 \(h_{eos}\) 向量并重新创建原始句子。

这个微型解码器的输入由两部分 \(H_1\) 和 \(H_2\) 组成，它们将 [EOS] 嵌入与可训练的位置嵌入 (\(p\)) 相结合:

() \\begin{array} { c } { { H _ { 1 } [ h _ { e o s } + p _ { 0 } , \\ldots , h _ { e o s } + p _ { N } ] , } } \\ { { { } } } \\ { { H _ { 2 } [ h _ { e o s } , e _ { x _ { 1 } } + p _ { 1 } , \\ldots , e _ { x _ { N } } + p _ { N } ] , } } \\end{array} ()

微型解码器使用注意力机制来重建数据。但是，这里有个陷阱: 作者对注意力矩阵应用了一个掩码 (\(M\))。这迫使重建过程严重依赖 \(h_{eos}\)，而不是通过查看相邻的 Token 嵌入来“作弊”。

注意力计算公式化为:

() \\begin{array} { r l } & { Q = H _ { 1 } W ^ { Q } , K = H _ { 2 } W ^ { K } , V = H _ { 2 } W ^ { V } } \\ & { \\quad M _ { i j } = \\left{ \\begin{array} { l l } { 0 , } & { \\mathrm { i f ~ } i \\neq j \\land B _ { i j } = 1 } \\ { - \\infty , } & { \\mathrm { e l s e } } \\end{array} \\right. , } \\ & { \\quad \\quad A = \\mathrm { s o f t m a x } \\left( \\frac { Q ^ { T } K } { \\sqrt { d } } + M \\right) V . } \\end{array} ()

这个重建步骤的目标函数是最小化重建的 Token 与原始输入 Token 之间的差异:

() \\mathcal { L } _ { M R C } = \\sum _ { t } C E ( x _ { t } \\mid A , H _ { 1 } , H _ { 2 } , \\Phi _ { d e c } ) . ()

通过最小化这个损失，梯度回流到主 LLM 中，专门更新权重以确保 \(h_{eos}\) 捕获输入文本的高保真摘要。

联合优化

最终的训练方案结合了这两个目标。模型被训练成同时是一个优秀的生成器 (MAR) 和一个优秀的编码器 (MRC) 。

() \\mathcal { L } _ { U n i M A E } = \\alpha \\mathcal { L } _ { M A R } + \\mathcal { L } _ { M R C } , ()

这里，\(\alpha\) 是一个平衡两个任务的权重参数。有趣的是，微型解码器在训练后会被丢弃。在推理过程中，你只需将文本输入 LLM，提取 [EOS] 向量，并将其用作你的句子嵌入。

实验与结果

研究人员使用 大规模文本嵌入基准 (MTEB) 对 UniMAE 进行了评估，这是一套包含检索、聚类、分类等 56 个数据集的严格测试套件。他们在不同大小的 Llama-3 模型 (1B、3B 和 8B 参数) 上测试了该方法。

1. 嵌入性能

结果令人印象深刻。UniMAE 显著优于标准基线，包括“Mean” (平均池化) 和“Echo” (重复句子) 。

表 1: 56 个 MTEB 数据集上的结果，最佳结果加粗显示，次佳结果以此划线显示。

如表 1 所示，UniMAE 在无监督方法中取得了最先进的结果。

  • 显著性: 在 LLaMA-3.2-1B 模型上，UniMAE 从 39.32 的平均分 (使用平均池化) 跃升至 52.81 。
  • 扩展性: 这种性能提升在 1B、3B 和 8B 模型中均有效。
  • 对比: 它甚至优于 MNTP , 这是一种试图使 LLM 变为双向的流行方法，但该方法计算昂贵且从根本上改变了架构。

2. 向量空间可视化

为了证明嵌入具有语义意义，作者使用 t-SNE (一种用于可视化高维数据的技术) 可视化了向量空间。

图 2: BiorxivClusteringS2S 测试集上前 6 个类别的句子嵌入 t-SNE 可视化。

图 2 将基础模型 (a) 与 UniMAE 训练的模型 (b) 进行了比较。

  • 左侧 (基础模型) : 代表不同领域科学论文的点是分散且混合的。模型难以区分“神经科学”和“生物信息学”。
  • 右侧 (UniMAE) : 聚类清晰且紧凑。模型已经学会了有效地将语义相似的文本分组在一起。

3. 保留生成能力

这是关键测试。许多方法在改进嵌入的同时，破坏了 LLM 撰写连贯文本的能力 (灾难性遗忘) 。

图 3: (a) MTEB-15 上不同 MAR 掩码比率下的结果。(b) 语言模型任务上的性能。

看一看 图 3(b) (右侧的柱状图) :

  • Base (蓝色) : LLM 在生成任务上的原始性能。
  • MNTP (绿色) : 一种竞争方法。注意性能的大幅下降 (例如，在 LLaMA-8B 上从 72 降至 41) 。该模型已无法有效地生成文本。
  • UniMAE (橙色) : 性能与 Base 模型几乎相同。

这证实了 UniMAE 成功地将 LLM 转变为混合模型: 既是高性能编码器，也是高性能生成器。

4. 掩码的重要性

图 3(a) (折线图) 探讨了自回归部分的“掩码比率”。结果表明, 40% 到 60% 之间的掩码比率是最佳的。

  • 太低 (0%): 模型学不到稳健的表示，因为任务太简单 (简单的死记硬背) 。
  • 太高 (>60%): 信号丢失；模型无法重构含义。

结论与未来影响

UniMAE 代表了统一 NLP 架构的重要一步。通过将 [EOS] Token 视为必须重建输入的潜在变量，作者使仅解码器 LLM 能够执行表示任务 (如检索和聚类) ，其竞争力可与专用编码器相媲美。

主要收获:

  1. 架构无关: 它适用于标准的仅解码器 LLM，无需永久性的结构更改。
  2. 高效: 它仅需约 100 个训练步数即可达到 SOTA 无监督结果。
  3. 通用性: 生成的模型可以驱动 RAG 系统 (同时处理搜索嵌入和答案生成) 或作为领域适应的基础模型。

作者指出了一个局限性: 用于训练的“微型解码器”会被丢弃，这在技术上代表了训练阶段参数的浪费。然而，考虑到推理时的收益——使用单个模型处理所有任务——这个成本似乎微不足道。

UniMAE 预示着这样一个未来: 我们不再区分“嵌入模型”和“生成模型”，而是迈向真正通用的语言模型。