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

TL;DR

Hierarchical Token Prepending (HTP) improves information flow in decoder-based large language models by partitioning inputs and using mean-pooling, leading to better long-document embeddings and performance across multiple benchmarks.

Contribution

HTP introduces a hierarchical approach with block-level summaries and mean-pooling to address information flow bottlenecks in decoder-based LLM embeddings, improving long-context performance.

Findings

Consistent performance gains across 11 retrieval datasets.

Improved long-document embedding quality.

Effective for both zero-shot and finetuned models.

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.