Reducing Representational Interference in Retrieval Embeddings with Trainable Prefix Tokens

AI summary

Overview

Research area: Information retrieval (cs.IR), specifically dense retrieval embedding models and the internal representation behavior of Transformer encoders/decoders.

Technical level: Intermediate. The method is simple to describe, but the motivation draws on attention-sink literature and the evaluation assumes familiarity with contrastive retrieval training, BEIR-style benchmarks, and pooling strategies.

Scope: The paper proposes prepending a small set of trainable [DUMP] prefix tokens to queries and passages during embedding-model fine-tuning, and shows across six models that this improves retrieval and long-context nDCG@10 without changing the architecture, pooling strategy, or loss.

What This Paper Is About

Embedding models squeeze a variable-length passage into one fixed-size vector, and this compression can mix genuinely semantic content with leftover token-level artifacts — for example, the disproportionate attention that "attention sink" positions attract despite carrying little meaning. The authors test whether giving the model a few dedicated, learnable prefix slots ([DUMP] tokens) lets it route that uninformative signal away from the pooled representation. The goal is a plug-and-play improvement to dense retrievers that requires no new architecture or auxiliary objective.

Key Contributions

1. A minimal intervention: K trainable [DUMP] tokens are added to the tokenizer vocabulary and prepended to both queries and passages. The model architecture, pooling strategy (e.g., [CLS] or last-token), and contrastive retrieval loss are all left unchanged, and the tokens receive no separate supervision — they are learned purely through the InfoNCE retrieval objective.

2. Broad empirical validation across architectures: The method is evaluated on six embedding models — encoder models gte-en-mlm-base, NeoBERT-base, EuroBERT-210m, and bge-m3-retromae, plus decoder models ettin-decoder-150m and MobileLLM-R1-140M — on a hard-negative BEIR subset, a Wikipedia subset from MTEB, and LongEmbed. The authors report average relative improvements of 2.46% on standard retrieval benchmarks and 2.41% under long-context retrieval.

3. A tunable-design study: Performance is compared for K ∈ {1, 2, 4} [DUMP] tokens (with K = 4 used in the main experiments), showing that more prefix tokens help but not monotonically.

4. Mechanistic evidence: Token-level analysis on gte-en-mlm-base over LongEmbed NQ shows that [DUMP] tokens change activation norms for low-information tokens and redirect attention away from the [CLS] token toward the learned prefixes.

Main Findings

• Consistent average gains across all six models (Table 1, nDCG@10 averages): gte-en-mlm-base 0.5922 → 0.5966; NeoBERT-base 0.5941 → 0.5967; EuroBERT-210m 0.5786 → 0.5884; bge-m3-retromae 0.6030 → 0.6066; ettin-decoder-150m 0.4377 → 0.4775; MobileLLM-R1-140M 0.5860 → 0.6093. The largest average jump is the decoder model ettin-decoder-150m.

• Gains are not uniform per dataset. Several individual cells regress or tie: gte-en-mlm-base on NQ (0.3828 → 0.3793), NeoBERT-base on NQ (0.3944 → 0.3830) and HotpotQA (0.5469 → 0.5440), and bge-m3-retromae on HotpotQA (0.5454 → 0.5430) and on NQ (0.4062 → 0.4062, unchanged). The claim is about averages, not per-dataset wins.

• Long-context gains on LongEmbed (Table 2): gte-en-mlm-base 0.4797 → 0.5024; NeoBERT-base 0.3391 → 0.3452; EuroBERT-210m 0.4313 → 0.4363; bge-m3-retromae 0.5548 → 0.5679; ettin-decoder-150m 0.3330 → 0.3498; MobileLLM-R1-140M 0.5580 → 0.5593.

• Token count matters but is not monotonic (Figure 1). For gte-mlm-base, two [DUMP] tokens give the best average improvement, and four tokens slightly hurt NQ while improving most other datasets. For MobileLLM-R1-140M, four tokens give the largest average gain. Using [DUMP] tokens at all is generally better than using none, but the best count varies by model and dataset.

• Low-information tokens carry unusually high activation norms in the baseline (Figure 2a). With [DUMP] training, the mean L2 norm of low-information tokens (identified via NLTK English stopwords) is reduced, and the learned prefixes D1–D4 settle into a stable activation range close to the [CLS] token.

• Attention is rerouted (Figure 2b). Last-layer self-attention from low-information tokens, averaged over heads, shifts away from [CLS] and toward the learned [DUMP] prefixes, suggesting the prefixes act as alternative sinks that reduce those tokens' direct influence on the pooled sequence representation.

Methodology in Plain English

The authors take an existing embedding model, add a handful of new special tokens called [DUMP] (using four indexed tokens, [DUMP1] through [DUMP4], in the main runs), and glue them to the front of every query and every passage before encoding. Because self-attention lets every position look at every other position, these prefix slots participate in building the contextual representation. Their embeddings sit in the model's normal embedding matrix and are updated by backpropagation from the same contrastive loss everyone else uses — the loss pushes each query's embedding closer to its positive passage and further from negatives, with cosine similarity divided by a temperature and an InfoNCE objective.

Nothing else changes: pooling still reads from its default target (the [CLS] token for encoder models, the last token for decoder models), and there are no extra losses. The same prefixed encoding is used at inference time.

Training uses 110k instances — 100k from RLHN, chosen for its relabeled hard negatives (which reduce false-negative noise), plus 10k from MLDR to add long-document examples. Models are fine-tuned for one epoch with maximum sequence length 8192, learning rate 1×10⁻⁴, a contrastive temperature scale of 20.0, bfloat16 precision, gradient caching, a global batch size of 256, and one negative passage per query. Evaluation reports nDCG@10 within the MTEB framework, using hard-negative versions of Natural Questions, HotpotQA, Quora, DBPedia, and MS MARCO, the English Wikipedia subset from MTEB, and LongEmbed for long-context retrieval.

For the analysis, the authors take gte-en-mlm-base trained on LongEmbed NQ and split tokens into three groups: the [CLS] sequence representation token, the learned prefixes D1–D4, and low-information tokens. They then measure mean L2 norm of final-layer hidden states per group, and last-layer attention from low-information tokens to [CLS] versus D1–D4 (averaged over heads, scaled by 10⁵ for readability).

Why This Matters

Impact on research: The paper connects two previously separate threads — attention-sink/outlier-token analysis in Transformers and the compression bottleneck in dense retrieval — and shows the connection has practical payoff. It also offers a rare inside-the-model explanation for a retrieval gain: the improvement is tied to observable changes in activation norms and attention routing, not just a higher benchmark number. The result is notable because the intervention barely qualifies as an architectural change, so it can be bolted onto existing retrieval training pipelines as a baseline or a low-cost add-on.

Real-world applications:

• Enterprise and legal document search, where long passages must be compressed into one vector and local evidence is easily diluted.
• Retrieval-augmented generation pipelines, where cleaner passage embeddings translate into better-grounded generated answers.
• Long-context question answering and multi-document summarization that depend on retrieving from lengthy articles.
• Multilingual and domain-specific search deployments, since the method applies to any contrastive embedding fine-tuning run without redesign.

Industry relevance: Adding a few tokens to the vocabulary and the input sequence is cheaper than re-architecting a retriever, changing pooling, or adding auxiliary losses, and it preserves compatibility with existing inference stacks and vector indexes. However, the gains are small in absolute nDCG@10 terms (fractions of a point for the encoder models) and not positive on every dataset, so the paper positions this as a marginal, low-cost improvement rather than a replacement for stronger embedding models.

Future Directions

• Choosing K principledly. The paper shows the best number of [DUMP] tokens varies by model and dataset and calls it a tunable hyperparameter, but no rule or search procedure for picking it is established.
• Extending beyond stopwords. Low-information tokens are identified using NLTK English stopwords on one model; whether the same routing effect holds for other languages, domains, or for special-token and positional artifacts more broadly is left open.
• Combining with existing fixes. The paper does not report interactions with alternative pooling strategies, other contrastive objectives, or hard-negative mining schemes, so whether the effect stacks with those is unknown.
• Accounting for the cost side. The presented results do not quantify inference-time overhead from the extra prefix positions or how gains scale with training data beyond the 110k-instance setup described.

Target Audience

Practitioners building or fine-tuning dense retrieval and embedding models who want a low-effort, architecture-preserving improvement; researchers studying attention sinks, outlier tokens, and representational interference in Transformers; and engineers integrating retrieval into RAG or long-document search systems who care about long-context compression quality. Readers should be comfortable with contrastive retrieval training and nDCG@10 evaluation to get the most from the results tables.

Authors’ abstract

Embedding models compress variable-length text into fixed-size vector representations, but this compression can entangle semantic content with residual token-level information, such as artifacts arising from attention sink behavior. We propose a simple method that prepends trainable [DUMP] tokens to the input during embedding model training. The method keeps the model architecture, pooling strategy, and training loss unchanged, while optimizing the added tokens through the standard contrastive retrieval objective. By introducing trainable prefix positions that interact with regular tokens during encoding, [DUMP] tokens help reduce the influence of residual or uninformative token-level signals on the final retrieval representation. Experiments on standard retrieval and long-context benchmarks show that our method consistently improves over finetuned baselines across encoder and decoder embedding models. Further analysis shows that [DUMP] tokens alter token-level activation and attention patterns associated with low-information tokens. These results suggest that trainable prefix tokens provide a lightweight way to improve retrieval representations.