Pooling And Attention: What Are Effective Designs For LLM-Based Embedding Models?

📄 arXiv: 2409.02727v2 📥 PDF

作者: Yixuan Tang, Yi Yang

分类: cs.CL, cs.IR

发布日期: 2024-09-04 (更新: 2024-09-05)

备注: https://github.com/yixuantt/PoolingAndAttn

💡 一句话要点

针对LLM嵌入模型，研究Pooling和Attention的有效设计方案，并提出多层可训练Pooling方法。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: LLM嵌入模型 池化策略 注意力机制 文本相似度 信息检索

📋 核心要点

  1. 现有LLM嵌入模型训练数据、基模型和训练设置各异，缺乏统一的有效设计方案。
  2. 论文核心在于探索不同Pooling和Attention策略对LLM嵌入模型性能的影响，并提出多层可训练Pooling方法。
  3. 实验表明，双向Attention和可训练Pooling在文本相似性和检索任务中表现优异，多层可训练Pooling更胜一筹。

📝 摘要（中文）

大型语言模型（LLM）在生成任务中的显著进步，推动了对基于LLM的嵌入模型的研究。这些模型采用不同的池化和注意力策略，在公共嵌入基准测试中取得了最先进的性能。然而，关于LLM嵌入模型的有效设计仍然存在疑问。这些模型通常在不同的数据集上训练，使用不同的LLM基础模型或训练设置。此外，公共嵌入基准测试的评估通常未能报告统计显著性，难以确定哪些设计真正有助于最终性能。这使得从业者难以找到LLM嵌入模型的最佳训练方案。本研究通过使用相同的训练数据和基础模型，但采用不同的池化和注意力策略，训练了一系列基于LLM的嵌入模型，从而进行了一项大规模实验。结果表明，没有万能的解决方案：虽然双向注意力和额外的可训练池化层在文本相似性和信息检索任务中表现优异，但在聚类和分类任务中，它们并没有显著超过像EOS-last token池化和默认因果注意力这样的更简单的设计。此外，我们提出了一种新的池化策略，即多层可训练池化，它使用交叉注意力网络转换所有隐藏层的输出，而不仅仅是最后一层。与现有的池化方法相比，该方法在文本相似性和检索任务中具有统计上的优越性。总的来说，本文阐明了LLM嵌入模型的有效训练策略。

🔬 方法详解

问题定义：现有基于LLM的嵌入模型在训练时使用了不同的数据集、基础模型和训练设置，这导致难以确定哪些设计选择真正对最终性能有贡献。此外，现有方法在评估时缺乏统计显著性分析，使得难以区分不同设计之间的优劣。因此，需要系统地研究不同的Pooling和Attention策略对LLM嵌入模型性能的影响，并找到更有效的训练方法。

核心思路：论文的核心思路是通过控制变量法，使用相同的训练数据和基础模型，只改变Pooling和Attention策略，来评估不同策略对LLM嵌入模型性能的影响。此外，论文还提出了一种新的Pooling策略，即多层可训练Pooling，旨在利用所有隐藏层的输出信息，从而提高嵌入模型的性能。

技术框架：整体框架包括：1) 使用预训练的LLM作为基础模型；2) 对LLM的输出进行Pooling操作，得到固定长度的嵌入向量；3) 使用对比学习目标函数训练嵌入模型；4) 在不同的下游任务上评估嵌入模型的性能。主要模块包括LLM编码器、Pooling层和对比学习损失函数。

关键创新：论文的关键创新在于提出了多层可训练Pooling策略。与传统的只使用最后一层隐藏层输出的Pooling方法不同，多层可训练Pooling利用所有隐藏层的输出，并通过一个交叉注意力网络来学习不同隐藏层输出的权重，从而更有效地利用了LLM的信息。

关键设计：多层可训练Pooling的关键设计包括：1) 使用交叉注意力网络来学习不同隐藏层输出的权重；2) 将所有隐藏层的输出连接起来，作为交叉注意力网络的输入；3) 使用可训练的查询向量来学习不同隐藏层输出的重要性。损失函数采用对比学习损失，例如InfoNCE。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，双向注意力和额外的可训练池化层在文本相似性和信息检索任务中表现优异。更重要的是，论文提出的多层可训练池化方法在文本相似性和检索任务中，相比于现有的池化方法，取得了统计上显著的性能提升。具体提升幅度未知，但该方法证明了利用所有隐藏层信息进行Pooling的有效性。

🎯 应用场景

该研究成果可应用于各种需要文本嵌入的场景，例如文本相似度计算、信息检索、文本聚类和文本分类。通过选择合适的Pooling和Attention策略，可以提高LLM嵌入模型在这些任务上的性能，从而提升相关应用的效率和准确性。该研究还有助于开发更有效的LLM嵌入模型训练方法，推动自然语言处理领域的发展。

📄 摘要（原文）

The significant advancements of Large Language Models (LLMs) in generative tasks have led to a growing body of work exploring LLM-based embedding models. While these models, employing different pooling and attention strategies, have achieved state-of-the-art performance on public embedding benchmarks, questions still arise about what constitutes an effective design for LLM-based embedding models. However, these models are often trained on different datasets, using different LLM base models or training settings. Moreover, evaluations on public embedding benchmarks often fail to report statistical significance, making it difficult to determine which designs truly contribute to final performance. This complicates the process for practitioners seeking optimal training recipes for LLM-based embedding models. In this study, we conduct a large-scale experiment by training a series of LLM-based embedding models using the same training data and base model but differing in their pooling and attention strategies. The results show that there is no one-size-fits-all solution: while bidirectional attention and an additional trainable pooling layer outperform in text similarity and information retrieval tasks, they do not significantly surpass simpler designs like EOS-last token pooling and default causal attention in clustering and classification tasks. Furthermore, we propose a new pooling strategy, Multi-Layers Trainable Pooling, which transforms the outputs of all hidden layers, rather than just the last layer, using a cross-attention network. This method proves to be statistically superior in text similarity and retrieval tasks compared to existing pooling methods. Overall, this paper sheds light on effective training strategies for LLM-based embedding models.