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

Abstract: This study investigates the existence of positional biases in Transformer-based language 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 language models, extending it to the domain of embedding learning. We examine positional biases at multiple stages of the training pipeline for an encoder-decoder neural retrieval model, namely language model 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 the beginning of the input content, with fine-tuning further aggravating this effect.

## introduction recent advancements have allowed transformerbased models to handle increasingly larger context lengths, resulting in the availability of language models (lms) that can accommodate input lengths reaching tens of thousands of tokens @xcite . however, studies assessing how well this context is captured by causal lms @xcite have shown that models are biased to information contained at the beginning or end of the input, losing information in the middle. instead of further analysing text generation, we extend this type of study to text representation learning, which has been a fundamental task for dense retrieval @xcite , and is also gaining attention in the context of retrieval-augmented generation @xcite and recommendation systems @xcite . specifically, we focus on web document retrieval, examining how well a single embedding represents a complete web document, while assessing the emergence of eventual position biases. we start by continuously pre-training and finetuning an encoder-decoder model similar to t5base @xcite but with a context length of 2048 tokens, following standard techniques to achieve a model that is representative of the state-of-the-art among the low-parameter scale. we leverage the ms-marco (v1) document collection @xcite , as this dataset is commonly used in retrieval evaluation benchmarks @xcite , and it is one of the major sources of training data for the fine-tuning of neural retrieval models @xcite . we found the existence of a dwell in the beginning effect, i.e. a positional bias displayed by the model where earlier parts of the input are dominant in the embedding. we track this behavior by evaluating the model on position-aware tasks during multiple stages of its training. from our experiments, we conclude that these positional biases start emerging during unsupervised contrastive pre-training, and that the heavy reliance on ms-marco data for fine-tuning will exacerbate this behavior. our models and code are available in a public github repository foot_0 . ## related work bi-encoders are now the state of the art approach to dense retrieval @xcite . current standard training setups leverage the usage of contrastive loss functions and methods such as ance @xcite to sample hard negative examples. other techniques that are often employed include in-domain pre-training @xcite and retrieval-aligned pre-training @xcite @xcite @xcite , which allow for a better fine-tuning starting point, consequently achieving stronger retrieval results. for long document retrieval, early methods dealt with the increased input length through heuristic aggregation strategies, which rely on segmenting the document into passages that are scored independently, with max-pooling being particularly effective @xcite . instead of aggregating scores, studies like parade @xcite considered the aggregation of passage-level representations. other authors @xcite used transformer architectures with sparse attention patterns @xcite to model the long inputs more efficiently, showing that, on ms-marco, the gains that arise from using such models are limited when compared to simple aggregation strategies. currently, llara @xcite achieves stateof-the-art performance in the ms-marco document retrieval task, by continually pre-training llama-7b @xcite with a retrievalaligned task. models like llara leverage context windows of up to 4096 tokens, relying on flashattention @xcite for fast and exact full attention computation, together with some variation of rotary position embeddings (rope) @xcite or attention with linear biases (alibi) @xcite . this enables stronger modeling of longer sequences, without the need of additional training, while resorting to full-attention computations. ## methodology this section details the training of a t5-base retriever with 2048 input length (t5-2k), adapting the t5 architecture to follow recent advancements in long-context language modeling, and following a state-of-the-art dense retrieval training pipeline. ## conclusions and future work this study investigated a dwell in the beginning effect on transformer-based models for document retrieval. through experiments with a t5 model and repllama, we observed that the embeddings tend to favor information located at the beginning of the input, leading to decreased performance when relevant information is elsewhere in the document. we investigate each step in the training pipeline, namely language model pre-training, contrastive pre-training, and contrastive fine-tuning, showing that biases emerge in the contrastive pretraining step, and that they persist throughout the fine-tuning process. our findings emphasize the importance of considering the quality of embeddings for long inputs, particularly in contexts where effectively capturing the entire sequence is essential for the downstream task. moreover, our results can further justify previous research which showed limited gains on long-sequence modeling for ms-marco, when compared to aggregation approaches @xcite . as for future work, we note that while our experiments focused on tied encoders, a similar study can be conducted using untied weights, given the size mismatch between queries and documents. furthermore, addressing the identified biases may involve devising more robust pre-training tasks, or curating better-distributed datasets, all while considering evaluation on appropriate retrieval benchmarks that require long-context modeling @xcite .