揭秘MoE大型語言模型：你的MoE模型可以免費(fèi)變身為嵌入模型
专家混合（MoE）示意图来自NVIDIA的技术博客
我最近发现了一篇有趣的论文，题目是“您的专家混合型大型语言模型实际上是免费使用的嵌入模型”[1]。最近的大型语言模型（LLM）架构趋势是解码器模型，因为它们的注意力机制，这些模型不适合用作嵌入模型。然而，作者们揭示了专家混合（MoE）大型语言模型实际上可以作为嵌入模型来应用各种嵌入任务，而无需额外的微调。在这篇博客中，首先让我们回顾一下MoE，并简要介绍其工作原理和实际应用。
目录:- 什么是混合专家模型（MoE）？
- 混合专家模型（MoE）作为嵌入模型的工作原理是什么？
- 实际应用：如何利用MoE与BERTopic结合
混合专家（MoE）是一种架构，包含多个名为“专家”的子网络。每个“专家”专注于不同的任务或数据的某个特定方面。MoE的一个优势是它能让AI模型在保持或提升质量的同时，用更少的计算量进行预训练比相同或更大规模的模型。因此，如果我们预算有限，我们可以使用MoE来获得比密集型、类似大小的常规模型更好的模型。近期的一个成功案例是，Mixtral 8 x 7B 在许多评估数据集上的表现优于LLaMA 2 70B。
从现在起，我们来了解一下MoE的架构。最近成功的MoE大多基于Transformer模型，因此我将重点介绍流行的基于Transformer的MoE架构。MoE主要有两个主要组成部分。
MoE与Transformer架构的对比（参考论文[3]提出的）
- 教育部层级分类
MoE 将 transformer 架构中的 FFN 层替换为 MoE 层。每个 MoE 层配备几个专家（如上图所示，每个 MoE 层通常配备 4 个专家），每个专家由一个简单的 FFN 层组成。值得注意的是，transformer 中的其他组件，如自注意力层，共享相同的权重。因此，MoE 的权重数量并不是直接等于其层的数量那么简单。例如，Mixtral 8 x 7B 的实际权重数量不是 8 x 7 = 56B，而是 47B，因为除了 MoE 层之外的其他层也共享相同的权重。
- 门限网络
一个门控网络或路由器是MoE中的一个关键组件。它接收输入令牌并为每个令牌选择最相关的专家。例如，在上面的示意图中，例如，图的左侧部分选择第二个专家来处理“more”这个词令牌。同时，路由器选择第一个专家来处理“Parameters”这个词令牌。通常，门控网络会选择最相关的前k个专家并将令牌发送给选定的专家；例如，Mixtral 8x7B选择前两个专家。
我们如何选择排名前k位的专家？我们使用softmax函数来计算专家的重要性的概率，并保留概率最大的前k位专家。我从上述图示中提取了筛选部分。
作者绘制的顶尖专家挑选流程图
gating 网络有自己的权重。我们通过 softmax 函数处理输入词元与 gating 网络权重的点积结果，从而得到专家与给定词元相关的概率。根据这个概率，我们可以选出 top-k 个最相关的专家。这种 MoE 也因此被称为“稀疏 MoE”。
这些是理解MoE作为嵌入模型所需的基本知识。我推荐读读这篇博客 [2]。现在，让我们来深入了解一下MoE是如何工作的，就像一个嵌入模型一样。
2. MoE 是如何运作作为嵌入模型的？ 嵌入的快速回顾在进入本节的主题之前，我们先快速回顾一下关于嵌入的知识。最近，嵌入在深度学习模型中被用作输入数据的内部表示，它包含了语义和数据信息的浓缩。我们通常会提取神经网络的最后一个隐藏状态作为嵌入。
作者画的模型嵌入的工作原理图解
我们通常使用编码器模型来提取嵌入，因为它们可以通过双向注意力捕捉语义，而仅解码器的模型则做不到这一点。仅解码器的模型通常使用因果性注意力仅与前一个词交互，因此它们无法像编码器-解码器模型那样捕捉丰富的语义，例如上下文。
MoE 作为嵌入模型是如何工作的？这是一种普遍的看法，解码器模型无法用来提取嵌入。然而，作者发现MoE中的路由权重为解码器嵌入提供了补充的有用信息。每一层的路由权重反映了对输入令牌的推理过程，因此它包含了输入的语义信息，而隐藏状态的嵌入可能无法完全保留这些信息。用数学公式，我们可以这样描述它：
路由权重计算公式
g 为 softmax 函数，H 表示隐藏状态。我们将所有 MoE 层的路由权重拼接在一起，以避免丢失模型的推理路径选择。
为了充分利用路由权重和解码器嵌入的潜力，作者提出了一种称为MoE嵌入（MoEE）的方法，以形成更全面的嵌入表示方式。MoEE主要有两种类型。其中一种是基于拼接的组合方法，下面将对此进行详细说明。
MoEE concat 公式
这种方法很简单，我们只需将路径权重与解码器的嵌入拼接起来。作者将这种方法称为MoEE（concat），即拼接方法，它可以在保留每个路由权重捕获的各自独有的信息的同时，允许下游任务利用组合后的表示。
另一种方法是加权求和整合。它对由路由权重和隐状态（HS）嵌入计算得出的相似度得分进行加权求和，简称为MoEE（求和）。这种方法常用于比较两个句子的任务，如语义相似性，例如语义文本相似性。
MoEE和sum公式
𝛼 是一个超参数，用于控制路由权重的贡献程度。计算每对的相似性得分后，我们计算出的相似性得分与实际相似性之间的等级相关系数，例如 Spearman 等级相关系数。
从实用角度来看，我认为MoEE（concat）使用起来很方便。此外，他们还利用PromptEOL技术[4]来增强MoEE。这项技术通过提示模板来约束大模型预测下一个令牌的语义信息。
PromptEOL的特定嵌入任务用提示
现在，这里是在不同MTEB任务中的表。
性能表（参考自[1]）
带有PromptEOL功能的MoEE可以表现更佳，优于监督和自监督方法。需要注意的是，该排行榜并不是最新的，因此该结果并不处于最先进（SOTA）的位置。该方法的价值在于，我们能够为嵌入任务获得相当不错的结果，并且可以在无需额外训练的情况下使用。
到目前为止，我们已经介绍了MoEE是如何工作的。接下来，我们将使用BERTopic来实现MoEE的功能，并对句子进行聚类。
3. 实际操作：结合MoEE与BERTopic在本节中，我们从预训练的MoE LLM中提取嵌入，并利用20新闻组数据集[5]与BERTopic进行处理。首先，值得一提的是，BERTopic是一个超越传统统计主题建模的便捷主题建模库。它利用Transformer模型的嵌入来进行主题聚类，我认为它非常适合用来检验其能力。首先，我们来准备好所需的环境。
环境配置我使用了一个装有Python 3.10的conda环境。在配备了CUDA 12.4和16GB VRAM的Ubuntu 20.04系统上进行了实验。你可能需要32GB的RAM来下载模型的权重。
conda create -n moee python=3.10 -y
conda activate moee（创建一个名为moee的环境，并激活它。）
接下来我们需要通过使用 pip 安装下面这些库。
pip install transformers torch bitsandbytes bertopic accelerate安装transformers, torch, bitsandbytes, bertopic和accelerate库
MoE模型一般需要大量的显存，因为需要先将整个模型加载到显存中。因此，我们需要使用bitsandbytes，这是一套用于量化的工具，从而节省显存内存。
我们需要做的是克隆官方的GitHub仓库（repository）。
git clone https://github.com/tianyi-lab/MoE-Embedding.git克隆此仓库到本地
所有准备工作已完成。现在，用MoEE实现BERTopic主题聚类。
结合MoEE和BERTopic来利用现在，我们将使用MoE作为BERTopic的嵌入模型并尝试主题分类。原仓库允许我们使用小型MoE模型，例如Qwen-1.5-MoE-A2.7B或OLMoE-1B-7B等。在本博客中，我将使用OLMoE-1B-7B，它适合在16GB VRAM上运行推理。首先，我们需要加载这个模型。
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。（我会稍后附上完整代码。）
从sklearn.datasets导入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的变换器函数。进行推理时，我们只需在参数中添加output_hidden_states=True。
现在可以开始做主题建模了。
# 步骤2 - 减少维度
umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine')
# 步骤3 - 对降维后的嵌入进行聚类处理
hdbscan_model = HDBSCAN(min_cluster_size=15, metric='euclidean', cluster_selection_method='eom', prediction_data=True)
# 步骤4 - 对主题进行分词
vectorizer_model = CountVectorizer(stop_words="english")
# 步骤5 - 创建主题表示
ctfidf_model = ClassTfidfTransformer()
# 步骤6 - （可选）使用 `bertopic.representation` 模型优化主题表示
representation_model = KeyBERTInspired()
# 所有步骤一起
topic_model = BERTopic(
embedding_model=embedding_model, # 步骤1 - 提取嵌入向量
umap_model=umap_model, # 步骤2 - 减少维度
hdbscan_model=hdbscan_model, # 步骤3 - 对降维后的嵌入进行聚类处理
vectorizer_model=vectorizer_model, # 步骤4 - 对主题进行分词
ctfidf_model=ctfidf_model, # 步骤5 - 提取主题词
representation_model=representation_model # 步骤6 - （可选）使用 `bertopic.representation` 模型优化主题表示
)
# 使用 BERTopic 模型进行主题建模
topics, probs = topic_model.fit_transform(docs, embeddings)我们默认设置下得到了42个主题；下面是几个示例。尽管我随机选择了这些主题，它们也能够很好地反映语义。
使用 MoEE，进行一些主题分析的随机结果
此外，这里有一个主题集群的可视化。
作者的议题集群可视化
请看主题集群可视化中的红色圆圈，它代表与计算机相关的主题0。与此相近的主题也与计算机相关的专业术语相关，例如图形、数码和打印机。
这种方法表明，我们可以在不进行训练的情况下获得不错的嵌入。尽管还有提升空间以达到当前最优的监督模型水平，本文的研究发现也为无监督的嵌入提取方法的进一步改进提供了良好的起点。
这是我的整段代码。你需要把这个文件放在MoE-Embedding目录的最顶层。
MMoE 的要点
参考[1] 李子悦，周天一，您的专家混合模型实际上是一个免费的嵌入模型 (2024), arXiv
[2] Omar S. 等，Mixture of Experts (2023)，Hugging Face
[3] William Fedus, Barrett Zoph 等，Switch Transformers：通过简单高效的稀疏性扩展到万亿参数模型论文（2021 年），arXiv 平台
[4] 江汀等人，使用大规模语言模型进行句子嵌入扩展 （2023年），arXiv
[5] 20 讨论组
这个故事发表在Generative AI上。你可以在LinkedIn上联系我们，并关注Zeniteq以获取最新的AI故事动态。
订阅我们的Newsletter和观看我们的YouTube频道，获取最新的人工智能新闻和动态。让我们一起创造人工智能的未来吧！
共同學(xué)習(xí)，寫下你的評論
評論加載中...
作者其他優(yōu)質(zhì)文章
