Table of Contents
Fetching ...

Hierarchical Token Prepending: Enhancing Information Flow in Decoder-based LLM Embeddings

Xueying Ding, Xingyue Huang, Mingxuan Ju, Liam Collins, Yozen Liu, Leman Akoglu, Neil Shah, Tong Zhao

TL;DR

Decoder-only LLM embeddings suffer from restricted backward information flow and over-compression of long sequences. Hierarchical Token Prepending (HTP) introduces block-level local and global summary tokens plus mean-pooling readout to create multiple backward pathways and reduce information squashing, backed by bounds showing enhanced robustness. Empirically, HTP yields consistent gains across 11 BEIR retrieval datasets and 30 general embedding benchmarks, particularly in long-context settings, and also improves finetuned embedding models like NV-Embed-v2. The method is simple, architecture-agnostic, and scalable, offering a practical route to superior long-document embeddings for retrieval and analysis tasks.

Abstract

Large language models produce powerful text embeddings, but their causal attention mechanism restricts the flow of information from later to earlier tokens, degrading representation quality. While recent methods attempt to solve this by prepending a single summary token, they over-compress information, hence harming performance on long documents. We propose Hierarchical Token Prepending (HTP), a method that resolves two critical bottlenecks. To mitigate attention-level compression, HTP partitions the input into blocks and prepends block-level summary tokens to subsequent blocks, creating multiple pathways for backward information flow. To address readout-level over-squashing, we replace last-token pooling with mean-pooling, a choice supported by theoretical analysis. HTP achieves consistent performance gains across 11 retrieval datasets and 30 general embedding benchmarks, especially in long-context settings. As a simple, architecture-agnostic method, HTP enhances both zero-shot and finetuned models, offering a scalable route to superior long-document embeddings.

Hierarchical Token Prepending: Enhancing Information Flow in Decoder-based LLM Embeddings

TL;DR

Decoder-only LLM embeddings suffer from restricted backward information flow and over-compression of long sequences. Hierarchical Token Prepending (HTP) introduces block-level local and global summary tokens plus mean-pooling readout to create multiple backward pathways and reduce information squashing, backed by bounds showing enhanced robustness. Empirically, HTP yields consistent gains across 11 BEIR retrieval datasets and 30 general embedding benchmarks, particularly in long-context settings, and also improves finetuned embedding models like NV-Embed-v2. The method is simple, architecture-agnostic, and scalable, offering a practical route to superior long-document embeddings for retrieval and analysis tasks.

Abstract

Large language models produce powerful text embeddings, but their causal attention mechanism restricts the flow of information from later to earlier tokens, degrading representation quality. While recent methods attempt to solve this by prepending a single summary token, they over-compress information, hence harming performance on long documents. We propose Hierarchical Token Prepending (HTP), a method that resolves two critical bottlenecks. To mitigate attention-level compression, HTP partitions the input into blocks and prepends block-level summary tokens to subsequent blocks, creating multiple pathways for backward information flow. To address readout-level over-squashing, we replace last-token pooling with mean-pooling, a choice supported by theoretical analysis. HTP achieves consistent performance gains across 11 retrieval datasets and 30 general embedding benchmarks, especially in long-context settings. As a simple, architecture-agnostic method, HTP enhances both zero-shot and finetuned models, offering a scalable route to superior long-document embeddings.

Paper Structure

This paper contains 51 sections, 3 theorems, 18 equations, 7 figures, 13 tables.

Key Result

Theorem 3.1

In a causal, decoder-only Transformer with $L$ layers, there exists a depth-dependent constant $K_L>0$ and a nonnegative, lower-triangular, row-stochastic mixing matrix $\mathbf{A}\in\mathbb{R}^{n\times n}$ (capturing aggregate attention+residual flow across layers) such that, for every input positi

Figures (7)

  • Figure 1: Left: Performance by embedding method across sentence lengths. Right: Performance drop of last-token and PromptEOL vs. mean at longer lengths.
  • Figure 2: Disabling second-pass backward attention significantly degrades STS performance.
  • Figure 3: The attention map shows that End-of-Sentence (EOS) tokens are capturing information with more attention lookup to previous tokens.
  • Figure 4: HTP partitions the input into blocks and creates a two-level summary. First, the hidden state of each block's final token is copied to a local summary token (<PST>). These local summaries are then aggregated into a global (<B-PST>) block at the front, making them accessible to all tokens. A mean readout with early exiting produces the final embedding.
  • Figure 5: Avg. NDCG@10 across different context lengths.
  • ...and 2 more figures

Theorems & Definitions (6)

  • Theorem 3.1: Mean vs. Last-token Embedding
  • Remark A.1: Multi‑head and projections
  • Lemma A.2: Path–sum equals matrix entry
  • proof
  • Theorem A.3: Last‑token vs mean‑token sensitivity
  • proof