《Pooling and Attention: What Are Effective Designs for LLM-Based Embedding Models?》
链接:https://arxiv.org/pdf/2409.02727
github:https://github.com/yixuantt/PoolingAndAttn
在搜广推场景,双塔模型是业界标配,尤其是搜索引擎、推荐系统和语义文本相似性场景,普遍需要获得user和item / query和doc 好的embedding表示,便于大规模高效计算,主要是通过对比表示学习实现。但LLM的能力,如何在这种场景下用上?本文也许能带给你一些启发。
作者是 Yixuan Tang 和 Yi Yang,来自香港科技大学。论文探讨了在大型语言模型(LLM)基础上构建文本嵌入模型时,有效的池化(Pooling)和注意力(Attention)机制设计。
摘要:
- 论文指出,尽管基于LLM的嵌入模型在公共嵌入基准测试中取得了最新性能,但关于LLM嵌入模型的有效设计仍存在疑问。
- 这些模型通常在不同的数据集上训练,使用不同的基础LLM或训练设置,且公共嵌入基准的评估往往没有报告统计显著性,使得难以确定哪些设计真正有助于最终性能。
- 研究通过在相同的训练数据和基础模型上训练一系列LLM嵌入模型,但池化和注意力策略不同,来大规模实验。
- 结果显示,没有一种通用的解决方案:双向注意力和额外的可训练池化层在文本相似性和信息检索任务中表现更好,但在聚类和分类任务中并没有显著超越更简单的设计,如EOS-last token池化和默认因果注意力。
- 论文提出了一种新的池化策略,多层可训练池化(Multi-Layers Trainable Pooling),它使用交叉注意力网络转换所有隐藏层的输出,而不是仅使用最后一层。
- 这种方法在文本相似性和检索任务中被证明比现有池化方法有统计上的优越性。
1. 引言:
- 文本嵌入是捕捉文本语义信息的高维表示,对信息检索和语义文本相似性等任务至关重要。
- 论文讨论了编码器模型(如BERT)和LLM作为嵌入模型的使用,并指出LLM在嵌入模型基准测试中的领先地位。
2. 常见的池化和注意力策略:
- 论文回顾了现有LLM嵌入模型中常用的池化和注意力策略,包括EOS-last token池化、均值池化和可训练池化层。
- 注意力策略部分讨论了因果注意力和双向注意力。
3. 多层可训练池化:
具体来说,多层可训练池化策略的步骤可以概括为:
这种策略的优势在于能够从LLM的不同层中捕获和整合信息,这可能包含关于输入文本的不同方面的语义信息。通过这种方式,模型能够生成更加丰富和信息量更大的文本嵌入,从而在某些任务(如文本相似性和信息检索)中提高性能。然而,这也带来了额外的计算复杂性和训练成本,需要在实际应用中进行权衡。
- 论文提出了一种新的池化策略,利用LLM中所有层的隐藏状态,并使用可训练网络进行转换。
-
多层可训练池化(Multi-Layers Trainable Pooling)策略是论文中提出的一种新的池化方法,旨在从大型语言模型(LLM)的所有隐藏层中捕获更丰富的语义信息。这种方法与传统的仅使用最后一层或特定层的池化策略不同,它通过以下几个关键步骤来实现:
-
输入层:首先,选择LLM的所有隐藏层的输出作为输入。对于使用因果注意力的LLM,使用所有层的EOS(End-of-Sequence)标记的隐藏状态作为输入。对于使用双向注意力的LLM,则考虑所有token的平均隐藏状态。
-
层权重矩阵:引入一个可训练的层权重矩阵,以反映不同层对于最终嵌入的重要性。通过将层权重矩阵与LLM的输出相结合,可以调整每个层的贡献度。
-
交叉注意力网络:通过交叉注意力网络进一步转换和提炼隐藏状态。这个网络计算固定、可训练的查询(Q)与由输入派生的键(K)和值(V)之间的注意力。这种方法允许模型从多层隐藏状态中筛选出最相关的信息,并将其编码到语义潜在空间中。
-
多层感知器(MLP):交叉注意力网络的输出再通过一个多层感知器(MLP)来产生最终的嵌入表示。
- 将LLM的所有层的隐藏状态作为输入。
- 通过可训练的层权重矩阵对不同层的重要性进行加权。
- 使用交叉注意力网络来处理加权后的层表示,以提炼和编码相关信息。
- 通过MLP生成最终的嵌入输出。
- 通过实验验证了不同层隐藏状态编码了不同信息,并提出中间层可能对下游任务有用。
4. 池化和注意力实验:
- 论文通过实验评估了不同池化和注意力策略的有效性,使用相同的数据集和训练协议来确保公平比较。
5. 实证分析:
- 分析了不同池化和注意力策略在多个任务上的有效性,并发现没有一种策略在所有任务上都是最优的。
6. 稳健性检查:
- 使用Qwen2-0.5B作为基础模型,进一步验证了分析的稳健性。
7. 相关工作:
- 论文讨论了基于编码器的嵌入模型和基于LLM的嵌入模型的相关研究。
8. 结论:
- 论文总结了研究发现,指出双向注意力和额外的可训练池化层在某些任务上表现更好,但并非所有任务都适用。同时,提出了一种新的多层可训练池化方法,以捕获更广泛的语义信息。
- 欢迎微信扫码关注公众号:nlp之路,关注发送LLM,免费领取LLM电子书