[Paper Review] SeDR: Segment Representation Learning for Long Documents Dense Retrieval

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

Topic Modeling4 citations

TL;DR

SeDR proposes segment representation learning for long document dense retrieval using a Segment-Interaction Transformer that enables document-aware, segment-sensitive representations with quadratic complexity only in segment length, not full document length. It outperforms existing methods on MS MARCO and TREC-DL, achieving state-of-the-art results through improved segment interaction and a novel Late-Cache Negative strategy to overcome GPU memory limitations in training.

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.

Motivation & Objective

• To address the limitations of dense retrieval (DR) on long documents due to Transformer’s quadratic complexity and finite embedding capacity.
• To overcome information loss from document truncation or splitting-and-pooling by preserving document-level context in segment representations.
• To mitigate GPU memory bottlenecks in training that limit negative sampling for long documents.
• To develop an efficient, scalable method that maintains high representation quality for long documents without increasing model parameters significantly.

Proposed method

• Proposes a Segment-Interaction Transformer that enables self-attention between segments of the same document, allowing segment representations to be document-aware while preserving O(n_s²) complexity per segment.
• Uses a modified Transformer architecture with segment-level attention to capture inter-segment relationships, improving upon independent encoding or global attention mechanisms.
• Introduces Late-Cache Negative, a strategy that stores and reuses encoded document representations from earlier training steps to provide additional negative samples beyond in-batch negatives.
• Employs a dynamic negative sampling mechanism where top-K hard negatives are selected based on similarity, with K=100 and cache size C=50 found optimal for balancing MRR and Recall.
• Constructs a segment-based dense retrieval index where each document is represented by multiple segment embeddings, enabling fine-grained relevance matching.
• Trains the model using contrastive learning with both in-batch and cached negatives to improve semantic discrimination.

Experimental results

Research questions

• RQ1 Can segment-level representation learning with inter-segment interaction improve long document retrieval beyond splitting-and-pooling or independent encoding?
• RQ2 Does a Segment-Interaction Transformer with localized attention achieve better performance and efficiency than global or sparse attention mechanisms in long document DR?
• RQ3 Can Late-Cache Negative effectively alleviate the negative sampling bottleneck caused by GPU memory limits during long document training?
• RQ4 How do different configurations of cache size and hardness in negative sampling affect retrieval performance on long documents?

Key findings

• SeDR achieves state-of-the-art performance on MS MARCO and TREC-DL, with MRR@100 of 0.632 and Recall@100 of 0.527 on TREC-DL 20.
• The Segment-Interaction Transformer outperforms other segment-interaction patterns, including Longformer and global attention, with fewer parameters and higher efficiency.
• Late-Cache Negative with C=50 and K=100 achieves the best balance between MRR@100 and Recall@100, demonstrating its effectiveness in mitigating GPU memory constraints.
• t-SNE visualization shows that SeDR avoids embedding collapse seen in global-attention models and maintains diverse, document-aware segment representations.
• The model reduces information loss from document splitting while maintaining computational efficiency, outperforming both max-pooling baselines and other segment-interaction methods.
• Empirical results confirm that segment representations with inter-segment interaction significantly improve retrieval performance compared to independent encoding.