Software 2.0

MoEE 混合专家嵌入模型

admin -

最近发现了一篇有趣的论文，题为“你的混合专家大模型 (Mixture-of-Experts LLM) 实际上是免费的嵌入模型”。[1] 最近的 LLM 架构趋势是解码器模型，由于其注意力方法，它不适合嵌入模型。然而，作者透露，混合专家 (MoE) LLM 可以作为嵌入模型，应用多种以嵌入为重点的任务，而无需进一步微调。在这篇博客中，首先，让我们回顾一下 MoE，我将介绍它的工作原理及其实际实现。

1、什么是混合专家 (MoE)？

混合专家 (MoE) 是一种具有多个子网络的架构，称为“专家”，每个子网络专门负责不同的任务或数据方面。 MoE 的优势之一是，它使 AI 模型能够以比相同或更大的模型更少的计算量进行预训练，同时保持或提高质量。因此，如果我们的预算有限，我们可以使用 MoE 实现比密集的、类似大小的传统模型更好的模型。最近的成功是，Mixtral 8 x 7B 在许多评估数据集上的表现优于 LLaMA 2 70B。

从现在开始，让我们研究一下 MoE 的架构。最近成功的 MoE 使用 Transformer 模型，因此我将重点介绍 Transformer 的流行 MoE 架构。MoE 主要有两个组件，如下所述。

本文 [3] 改编的 MoE 与 Transformer 架构之间的比较
  • MoE 层

MoE 在 Transformer 架构中用 MoE 层替换了前馈网络 (FFN) 层。每个 MoE 层都有一些专家（例如上图中的 4 个专家），每个专家由简单的 FFN 层组成。请注意，Transformer 中的其他组件（例如自注意力层）共享相同的权重。因此，MoE 的权重数量并不简单。例如，Mixtral 8 x 7B 的权重不是 8 x 7 = 56B，而是 47B，因为 MoE 层以外的其他层共享相同的权重。

  • 门控网络

门控网络或路由器是 MoE 中的关键组件。它接受输入标记并为每个标记选择最相关的专家。例如，在上图中，路由器的左侧选择第二个专家来处理单词“more”标记。同时，路由器确定第一个专家来处理单词“Parameters”标记。通常，门控网络选择与给定标记相关的前 k 个专家，并将标记发送给选定的专家；例如，Mixtral 8 x 7B 选择前 2 个专家。

我们如何选择前 k 个专家？我们使用 softmax 函数计算专家的重要性概率并保留前 k 个概率专家，如下所示。我提取了上图中的门控部分。

Top-k 专家选择流程图

门控网络有其权重。我们将 softmax 函数应用于输入单词 token 与门控网络权重的点积结果，然后得到专家与给定 token 相关的概率。根据概率，我们可以选择 top-k 相关专家。具有这种门控网络的 MoE 称为稀疏 MoE。

这些是理解 MoE 如何作为嵌入模型工作所需的基础知识。为了进一步理解，我建议阅读这篇博客 [2]。现在，让我们深入了解 MoE 实际上如何作为嵌入模型工作。

2、MoE 如何作为嵌入模型工作？

在深入讨论本节主题之前，让我们快速回顾一下嵌入（embedding）。最近，嵌入一直是深度学习模型中输入数据的内部表示，它具有语义和浓缩数据信息。我们通常将神经网络的最后一个隐藏状态提取为嵌入，如下所示。

嵌入模型的工作原理

我们通常使用基于编码器的模型来提取嵌入，因为与仅解码器的模型相比，它们可以通过双向注意力捕获语义。仅解码器的模型通常使用因果注意力来仅与前一个单词标记进行交互；因此，它们无法像编码器-解码器模型那样捕获丰富的语义，例如上下文信息。

人们普遍认为解码器模型不能用于嵌入提取。然而，作者发现 MoE 中的路由权重为解码器嵌入提供了补充信息。每一层中的路由权重反映了对输入标记的推理选择，因此它包含了隐藏状态的嵌入可能丢失的输入的语义信息。在数学公式中，我们可以将其描述为：

路由权重公式

g 是 softmax 函数，H 表示隐藏状态。我们将所有 MoE 层的路由权重连接起来，以避免丢失模型的推理选择。

为了充分利用路由路由权重和解码器嵌入，作者提出了一种称为 MoE 嵌入 (MoEE) 的方法来形成更全面的嵌入表示。MoEE 有两种类型。一种方法是基于连接的组合，如下所述。

MoEE(concat) 公式

这种方法很简单，我们只需连接路由权重和解码器嵌入。作者将此方法称为 MoEE(concat)。它可以保留每个路由权重捕获的不同信息，同时允许下游任务利用组合表示。

另一种方法是加权和集成。它对从路由权重和隐藏状态 (HS) 嵌入计算出的相似度得分进行加权和，表示为 MoEE (sum)。此方法用于比较两个句子的任务，例如语义文本相似度。

MoEE(sum) 公式

𝛂 是一个控制路由权重贡献的超参数。计算每对的相似度得分后，我们计算计算出的相似度得分与真实相似度之间的秩相关，例如 Spearman 秩相关。

对于实际使用，我认为 MoEE(concat) 易于使用。此外，作者利用 PromptEOL 技术 [4] 来增强 MoEE。此技术提示以下模板以约束 LLM 预测下一个标记中的语义信息。

PromptEOL 对嵌入任务的特定提示

现在，这是跨 MTEB 任务的性能表。

由 [1] 改编的性能表

使用 PromptEOL 的 MoEE 可以比监督和自监督方法更好地工作。请注意，此排行榜不是最新的，因此此结果不是 SOTA。这种方法的价值在于我们可以获得不错的嵌入任务结果，并且无需任何进一步的训练就可以使用它。

到目前为止，我们已经介绍了 MoEE 的工作原理。在下一节中，我们将使用 BERTopic 和聚类句子来实现 MoEE。

3、利用 BERTopic 实现 MoEE

在本节中，我们从预先训练的 MoE LLM 中提取嵌入，并使用 20 个新闻组数据集 [5] 将它们与 BERTopic 结合使用。供你参考，BERTopic 是一个超越传统统计主题建模的便捷主题建模库。它利用 Transformer 中的嵌入进行主题聚类，因此我认为它适合检查功能。首先，让我们准备一个环境。

3.1 环境设置

我使用了带有 Python 3.10 的 conda 环境。我在 Ubuntu 20.04 上使用 cuda 12.4、16 GB VRAM 进行了实验。你可能需要 32 GB RAM 来下载模型权重。

conda create -n moee python=3.10 -y
conda activate moee

接下来，我们需要通过 pip 安装以下库。

pip install transformers torch bitsandbytes bertopic accelerate

MoE 模型通常需要较高的 VRAM，因为我们需要提前将整个模型加载到 VRAM 中。因此，我们需要使用量化包 bitsandbytes 来节省 VRAM 内存。

我们需要克隆官方 GitHub 存储库。

git clone https://github.com/tianyi-lab/MoE-Embedding.git

所有准备工作都已完成。现在，让我们使用 MoEE 实现 BERTopic 的主题聚类。

3.2 利用 MoEE 和 BERTopic

现在，我们将使用 MoEE 作为 BERTopic 的嵌入模型并尝试主题聚类。原始存储库允许我们使用小型 MoE 模型，例如 Qwen-1.5-MoE-A2.7B 或 OLMoE-1B-7B。在这篇博客中，我将使用 OLMoE-1B-7B，它对于在 16 GB VRAM 上运行推理来说是负担得起的。首先，我们需要加载 OLMoE-1B-7B。

kwargs = {
        "base_model": 'allenai/OLMoE-1B-7B-0924',
        "normalized": False,
        "torch_dtype": torch.bfloat16,
        "mode": "embedding",
        "pooling_method": "mean",
        "attn_implementation": "sdpa",
        "attn": "bbcc",
    }

config = {
    'embed_method': 'prompteol',
    'emb_info': 'MoEE'
    }

embedding_model = MOEE(model_name_or_path='allenai/OLMoE-1B-7B-0924', **kwargs)

接下来，我们需要计算 20 个新闻组数据集的嵌入以通过 BERTopic。 稍后我会附上完整代码。

from sklearn.datasets import fetch_20newsgroups

docs = fetch_20newsgroups(subset='all', remove=('headers', 'footers', 'quotes'))['data']

dataset = MyDataset(docs)
dataloader = DataLoader(dataset=dataset, batch_size=8)
embeddings = None

for batch in tqdm(dataloader):
    with torch.no_grad():      
        embedding = embedding_model.encode(batch, **config)
        
        if embeddings is None:
            embeddings = embedding[0]
        else:
            embeddings = np.vstack((embeddings, embedding[0]))
    
    torch.cuda.empty_cache()

为了提前计算嵌入，我们使用torch.utils.data.DataLoader 作为迭代器，并对每个批处理文档进行编码。请注意，我们必须将嵌入作为 np.asarray 类型传递给 BERTopic。

当你想使用自己的 MoE 模型时，必须实现从每个 MoE 层获取路由权重。对于隐藏状态嵌入，我们可以利用 HuggingFace 转换器函数。我们只需要推理时传递 ut_hidden_​​states=True参数。

现在，我们可以运行主题建模了。

# Step 2 - Reduce dimensionality
umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine')

# Step 3 - Cluster reduced embeddings
hdbscan_model = HDBSCAN(min_cluster_size=15, metric='euclidean', cluster_selection_method='eom', prediction_data=True)

# Step 4 - Tokenize topics
vectorizer_model = CountVectorizer(stop_words="english")

# Step 5 - Create topic representation
ctfidf_model = ClassTfidfTransformer()

# Step 6 - (Optional) Fine-tune topic representations with 
# a `bertopic.representation` model
representation_model = KeyBERTInspired()

# All steps together
topic_model = BERTopic(
  embedding_model=embedding_model,          # Step 1 - Extract embeddings
  umap_model=umap_model,                    # Step 2 - Reduce dimensionality
  hdbscan_model=hdbscan_model,              # Step 3 - Cluster reduced embeddings
  vectorizer_model=vectorizer_model,        # Step 4 - Tokenize topics
  ctfidf_model=ctfidf_model,                # Step 5 - Extract topic words
  representation_model=representation_model # Step 6 - (Optional) Fine-tune topic representations
)

# topic modeling using BERTopic model
topics, probs = topic_model.fit_transform(docs, embeddings)

我们按照默认设置获得了 42 个主题；下面显示了一些示例。尽管我随机挑选了主题，但它可以很好地捕捉语义。

使用 MoEE 进行主题建模的随机结果

此外，这里是主题集群可视化。

作者的主题聚类可视化

请看主题聚类可视化中的红色圆圈。这个红色圆圈指的是主题 0，与计算机有关。更接近的主题也与机械词有关，例如图形、数字和打印机。

这种方法向我们展示了我们可以在不进行任何训练的情况下获得不错的嵌入。虽然仍有提升质量的空间，以相当于 SOTA 监督模型，但本文的发现是进一步改进无需训练的嵌入提取方法的良好一步。

这是我的完整代码。你需要将此文件放入 MoE-Embedding 目录的顶部。

原文链接：Unlocking Mixture-of-Experts (MoE) LLM : Your MoE model can be embedding model for free

汇智网翻译整理，转载请标明出处