[论文解读] SeDR: Segment Representation Learning for Long Documents Dense Retrieval

Junying Chen, Qingcai Chen+2 more|arXiv (Cornell University)|Nov 20, 2022

Topic Modeling被引用 4

一句话总结

SeDR 通过一种仅在段落长度上具有二次复杂度（而非整个文档长度）的段落交互 Transformer，提出了一种用于长文档密集检索的段落表示学习方法，实现了文档感知、段落敏感的表示。该方法在 MS MARCO 和 TREC-DL 上优于现有方法，通过改进段落间交互并引入一种新颖的 Late-Cache Negative 策略，克服了训练过程中 GPU 显存限制的问题，实现了最先进性能。

ABSTRACT

Recently, Dense Retrieval (DR) has become a promising solution to document retrieval, where document representations are used to perform effective and efficient semantic search. However, DR remains challenging on long documents, due to the quadratic complexity of its Transformer-based encoder and the finite capacity of a low-dimension embedding. Current DR models use suboptimal strategies such as truncating or splitting-and-pooling to long documents leading to poor utilization of whole document information. In this work, to tackle this problem, we propose Segment representation learning for long documents Dense Retrieval (SeDR). In SeDR, Segment-Interaction Transformer is proposed to encode long documents into document-aware and segment-sensitive representations, while it holds the complexity of splitting-and-pooling and outperforms other segment-interaction patterns on DR. Since GPU memory requirements for long document encoding causes insufficient negatives for DR training, Late-Cache Negative is further proposed to provide additional cache negatives for optimizing representation learning. Experiments on MS MARCO and TREC-DL datasets show that SeDR achieves superior performance among DR models, and confirm the effectiveness of SeDR on long document retrieval.

研究动机与目标

• 为解决 Transformer 的二次复杂度和有限嵌入容量导致的长文档密集检索（DR）性能瓶颈。
• 通过在段落表示中保留文档级上下文，缓解因文档截断或分段池化导致的信息损失。
• 缓解训练过程中因 GPU 显存限制而造成的负样本采样瓶颈。
• 开发一种高效可扩展的方法，在不显著增加模型参数量的前提下，保持长文档的高质量表示。

提出的方法

• 提出一种段落交互 Transformer，支持同一文档内段落之间的自注意力机制，使段落表示具备文档感知能力，同时每段保持 O(n_s²) 的复杂度。
• 采用改进的 Transformer 架构，引入段落级别注意力，以捕捉段落间关系，优于独立编码或全局注意力机制。
• 提出 Late-Cache Negative 策略，通过存储并重用早期训练步骤中编码的文档表示，提供超出批次内负样本的额外负样本。
• 采用动态负样本采样机制，基于相似度选择 top-K 难负样本，K=100 且缓存大小 C=50 时在 MRR 和 Recall 之间达到最佳平衡。
• 构建基于段落的密集检索索引，每个文档由多个段落嵌入表示，实现细粒度的相关性匹配。
• 使用对比学习进行模型训练，结合批次内负样本与缓存负样本来提升语义区分能力。

实验结果

研究问题

• RQ1 通过段落级表示学习并引入段落间交互，是否能超越分段池化或独立编码方式，在长文档检索中取得更好效果？
• RQ2 采用局部注意力机制的段落交互 Transformer，在长文档密集检索中是否相比全局注意力或稀疏注意力机制展现出更优的性能与效率？
• RQ3 Late-Cache Negative 是否能有效缓解因 GPU 显存限制导致的负样本采样瓶颈？
• RQ4 负样本采样中缓存大小与负样本难度的不同配置如何影响长文档检索性能？

主要发现

• SeDR 在 MS MARCO 和 TREC-DL 上均达到最先进性能，TREC-DL 20 的 MRR@100 达到 0.632，Recall@100 达到 0.527。
• 段落交互 Transformer 在参数更少的情况下优于 Longformer 和全局注意力机制，展现出更高的效率与性能。
• 当 C=50 且 K=100 时，Late-Cache Negative 在 MRR@100 与 Recall@100 之间实现最佳平衡，有效缓解了 GPU 显存约束。
• t-SNE 可视化显示，SeDR 避免了全局注意力模型中常见的嵌入坍塌问题，保持了多样化且具备文档感知能力的段落表示。
• 该模型在减少文档分段导致的信息损失的同时保持了计算效率，优于最大池化基线方法及其他段落交互方法。
• 实证结果证实，引入段落间交互的段落表示显著提升了检索性能，优于独立编码方式。