Dwell in the Beginning: How Language Models Embed Long Documents for Dense Retrieval (2404.04163v2)

Abstract: This study investigates the existence of positional biases in Transformer-based models for text representation learning, particularly in the context of web document retrieval. We build on previous research that demonstrated loss of information in the middle of input sequences for causal LLMs, extending it to the domain of representation learning. We examine positional biases at various stages of training for an encoder-decoder model, including LLM pre-training, contrastive pre-training, and contrastive fine-tuning. Experiments with the MS-MARCO document collection reveal that after contrastive pre-training the model already generates embeddings that better capture early contents of the input, with fine-tuning further aggravating this effect.

• The paper reveals that modified T5 models exhibit a 'dwell in the beginning' bias, favoring early content in dense retrieval tasks.
• It employs a training regimen combining unsupervised and supervised contrastive learning to optimize long-document representations.
• Experimental results demonstrate competitive retrieval performance on MS-MARCO, underscoring the need to mitigate positional biases in transformer models.

Positional Biases in Transformer-Based Models for Dense Retrieval

Introduction

The efficiency of transformer-based models in handling and interpreting large sequences of text proves a continually interesting area of paper, particularly in tasks like web document retrieval. A recent examination into these models reveals noticeable positional biases—the tendency for these models to disproportionately favor information at the beginning of inputs during text representation learning. Such biases present notable implications for the effectiveness of document retrieval systems, especially those relying on dense retrieval methods. Through a series of experiments utilizing the MS-MARCO document collection and a variant of the T5 model adapted for elongated input sequences, insights were gained into the nature and implications of these biases.

Model and Training Overview

At the core of the paper is a modified T5 model, dubbed 2K-T5, crafted to accommodate an input length of 2048 tokens. This adaptation includes replacing traditional positional embeddings with Rotary Position Embeddings (RoPE) for better context handling over extended text lengths. The model undergoes a structured training regimen:

• LLM Pre-training: Utilizes the MS-MARCO corpus to adapt the model to longer contexts and document-specific content distributions.
• Unsupervised Contrastive Pre-training: Aims to align the model better with its final retrieval task by training it to distinguish between closely related text samples.
• Supervised Contrastive Fine-tuning: Directly trains the model on the retrieval task using the MS-MARCO dataset, progressively refining its ability to extract and prioritize useful text representation.

Experimental Observations and Results

Retrieval Performance

Initial benchmarks place the 2K-T5 model competitively with state-of-the-art protocols in the domain of dense retrieval, striking a balance between Mean Reciprocal Rank (MRR) and Recall measurements. Despite its leaner architecture, the model performs comparably with larger, more complex systems, indicating efficient learning and generalization capabilities over extensive text inputs.

Positional Biases Analysis

Critical to the paper was the assessment of positional biases, colloquially termed the "dwell in the beginning" effect. Modifications to the retrieval documents, where relevant passages were systematically repositioned, elucidated the model's preferential treatment towards information located at the start of the input. Notably, this bias was amplified with further fine-tuning on MS-MARCO data, suggesting that the distribution of relevant information in training datasets could encourage or exacerbate such biases.

The paper further dissected the training pipeline to pinpoint when biases emerge, discovering a clear inclination towards initial input segments post contrastive pre-training. Interestingly, LLM pre-training did not visibly introduce these biases, suggesting that the methods and objectives of subsequent training phases significantly influence the model's attention to various sections of the input.

Implications and Future Directions

The presence of positional biases in transformer-based models for text representation has tangible implications, especially in applications requiring nuanced comprehension across the full span of lengthy documents. The "dwell in the beginning" phenomenon could, for instance, undermine the effectiveness of retrieval-augmented generation tasks where exhaustive document understanding is paramount.

Future explorations might focus on refining pre-training and fine-tuning methodologies to mitigate these biases, perhaps by incorporating diversified data distributions or by introducing training objectives that penalize positional predilections. Moreover, extending the analysis to other languages and domains could further elucidate the nuances of model behavior in dense retrieval settings.

Concluding Remarks

The paper's insights into positional biases within transformer-based models for dense retrieval highlight a critical area of focus for the continued development and refinement of AI-driven information retrieval systems. As the field progresses, understanding and addressing these biases will be vital in harnessing the full potential of LLMs in comprehensive, context-aware applications.