Table of Contents
Fetching ...

LMK > CLS: Landmark Pooling for Dense Embeddings

Meet Doshi, Aashka Trivedi, Vishwajeet Kumar, Parul Awasthy, Yulong Li, Jaydeep Sen, Radu Florian, Sachindra Joshi

TL;DR

The paper addresses the pooling bottleneck in dense embeddings, where CLS pooling biases representations toward early positions due to RoPE long-term decay and mean pooling dilutes salient local signals. It proposes Landmark (LMK) pooling, which inserts landmark tokens between fixed-size chunks and forms the final embedding by mean-pooling the LMK token embeddings, formalized as $\mathcal{P}_{\text{LMK}}(H^{enc}) = \frac{1}{|\mathcal{L}|} \sum_{l \in \mathcal{L}} h_l$. Through extensive English and multilingual experiments, LMK matches or exceeds existing pooling methods on short-context tasks and yields substantial improvements on long-context benchmarks, with robustness across granularities and compatibility with pretraining approaches like RetroMAE. The work demonstrates that LMK provides a practical, scalable pooling mechanism that balances local salient features with global information, enabling better long-context extrapolation in dense embeddings for retrieval and beyond.

Abstract

Representation learning is central to many downstream tasks such as search, clustering, classification, and reranking. State-of-the-art sequence encoders typically collapse a variable-length token sequence to a single vector using a pooling operator, most commonly a special [CLS] token or mean pooling over token embeddings. In this paper, we identify systematic weaknesses of these pooling strategies: [CLS] tends to concentrate information toward the initial positions of the sequence and can under-represent distributed evidence, while mean pooling can dilute salient local signals, sometimes leading to worse short-context performance. To address these issues, we introduce Landmark (LMK) pooling, which partitions a sequence into chunks, inserts landmark tokens between chunks, and forms the final representation by mean-pooling the landmark token embeddings. This simple mechanism improves long-context extrapolation without sacrificing local salient features, at the cost of introducing a small number of special tokens. We empirically demonstrate that LMK pooling matches existing methods on short-context retrieval tasks and yields substantial improvements on long-context tasks, making it a practical and scalable alternative to existing pooling methods.

LMK > CLS: Landmark Pooling for Dense Embeddings

TL;DR

The paper addresses the pooling bottleneck in dense embeddings, where CLS pooling biases representations toward early positions due to RoPE long-term decay and mean pooling dilutes salient local signals. It proposes Landmark (LMK) pooling, which inserts landmark tokens between fixed-size chunks and forms the final embedding by mean-pooling the LMK token embeddings, formalized as . Through extensive English and multilingual experiments, LMK matches or exceeds existing pooling methods on short-context tasks and yields substantial improvements on long-context benchmarks, with robustness across granularities and compatibility with pretraining approaches like RetroMAE. The work demonstrates that LMK provides a practical, scalable pooling mechanism that balances local salient features with global information, enabling better long-context extrapolation in dense embeddings for retrieval and beyond.

Abstract

Representation learning is central to many downstream tasks such as search, clustering, classification, and reranking. State-of-the-art sequence encoders typically collapse a variable-length token sequence to a single vector using a pooling operator, most commonly a special [CLS] token or mean pooling over token embeddings. In this paper, we identify systematic weaknesses of these pooling strategies: [CLS] tends to concentrate information toward the initial positions of the sequence and can under-represent distributed evidence, while mean pooling can dilute salient local signals, sometimes leading to worse short-context performance. To address these issues, we introduce Landmark (LMK) pooling, which partitions a sequence into chunks, inserts landmark tokens between chunks, and forms the final representation by mean-pooling the landmark token embeddings. This simple mechanism improves long-context extrapolation without sacrificing local salient features, at the cost of introducing a small number of special tokens. We empirically demonstrate that LMK pooling matches existing methods on short-context retrieval tasks and yields substantial improvements on long-context tasks, making it a practical and scalable alternative to existing pooling methods.
Paper Structure (23 sections, 7 equations, 5 figures, 17 tables, 2 algorithms)

This paper contains 23 sections, 7 equations, 5 figures, 17 tables, 2 algorithms.

Figures (5)

  • Figure 1: MLDR (English) dataset long context retrieval performance for different pooling strategies using modernbert-base finetuned on MSMarco Passages.
  • Figure 2: Normalized attention distribution over long MLDR documents ($>8192$ tokens) for mmBERT-base fine-tuned with CLS and LMK pooling with varying maximum sequence lengths (MSL).
  • Figure 3: Binned distribution of passage token lengths in the MS MARCO Document Ranking training set, computed using the ModernBERT-base tokenizer.
  • Figure 4: Retrieval Hits@k between LMK token embeddings and left, right, or any (left or right) chunk embeddings for sequences of length 8,192, with LMK granularity 32 ($\geq$248 chunks).
  • Figure 5: Long-context retrieval performance across datasets for different pooling strategies using ModernBERT-base fine-tuned on MSMarco Passages.