PaperlayerView exact revision on arXiv

Source-linked AI summary

KV-Embedding: Training-free Text Embedding via Internal KV Re-routing in Decoder-only LLMs

Yixuan Tang, Yi Yang

arXiv:2601.01046v1cs.CL

TL;DR

Training-free embeddings from decoder-only LLMs face causal-context asymmetry and next-token prediction bias. KV-Embedding reroutes the final token’s KV states as a prefix and selects layers using intrinsic dimensionality; it outperforms training-free baselines on MTEB and maintains robust performance on long sequences.

  • Problem

    Training-free decoder-only LLM embeddings are challenged by causal attention that withholds later context from early tokens and next-token prediction bias in the final token.

    Source: §1 · p. 1

  • Method

    KV-Embedding reroutes final-token KV states as an internal prefix and uses intrinsic-dimensionality-based layer selection.

    Source: §1 · p. 2

  • Results

    KV-Embedding outperforms existing training-free methods by up to 10% across three backbones and maintains robust performance on sequences up to 4,096 tokens.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The framework provides an efficient, scalable way to use frozen LLMs for text embeddings without manual cross-architecture layer tuning.

    Source: §1 · p. 2

    KV-Embedding increases latency over standard pooling and may not match supervised models fine-tuned with contrastive objectives on large datasets.

    Source: Limitations · p. 9

Abstract

from arXiv · show

While LLMs are powerful embedding backbones, their application in training-free settings faces two structural challenges: causal attention restricts early tokens from accessing subsequent context, and the next-token prediction objective biases representations toward generation rather than semantic compression. To address these limitations, we propose KV-Embedding, a framework that activates the latent representation power of frozen LLMs. Our method leverages the observation that the key-value (KV) states of the final token at each layer encode a compressed view of the sequence. By re-routing these states as a prepended prefix, we enable all tokens to access sequence-level context within a single forward pass. To ensure model-agnostic applicability, we introduce an automated layer selection strategy based on intrinsic dimensionality. Evaluations on MTEB across Qwen, Mistral, and Llama backbones show that KV-Embedding outperforms existing training-free baselines by up to 10%, while maintaining robust performance on sequences up to 4,096 tokens. These results demonstrate that internal state manipulation offers an efficient alternative to input modification, and we hope this work encourages further exploration of LLM internals for representation learning.

1 Introduction

Training-free embeddings avoid the computation and datasets required by fine-tuning, but causal attention and next-token prediction create incomplete or prediction-biased representations. KV-Embedding reroutes final-token KV states as a global prefix and selects layers by intrinsic dimensionality; evaluations report gains over training-free baselines.

1 Introduction

  • Causal attention leaves early tokens unaware of later context, while next-token prediction biases the final token away from semantic compression.Mean pooling aggregates incomplete representations; last-token pooling relies on a prediction-biased vector.Source: §1 · p. 1
  • Existing fixes include prompts, input repetition, and special-token prepending, but the latter approaches can double sequence length or use out-of-vocabulary tokens.Input repetition can also incur quadratic attention complexity and lost-in-the-middle effects.Source: §1 · p. 1
  • KV-Embedding re-routes the final token’s per-layer KV states as an internal prefix, giving all tokens sequence-wide context in one forward pass.It also uses a prompt to guide the final token toward semantic compression and selects re-routing layers using intrinsic dimensionality.Source: §1 · p. 2, Figure 1 · p. 2
  • KV-Embedding outperforms existing training-free methods by up to 10% across Qwen3-4B, Mistral-7B, and Llama-3.1-8B, with robust LoCoV1 performance up to 4,096 tokens.The paper also reports improved isotropy, alignment, and uniformity.Source: §1 · p. 2

2 Related Work

Related work spans fine-tuned LLM embeddings, training-free extraction, and attention-state manipulation. KV-Embedding differs by using internal KV states for embedding extraction without training or input-level changes.

2 Related Work

  • Fine-tuned decoder-only embedding models achieve strong representations, but their compute and dataset requirements limit scaling to rapidly evolving LLM architectures.Prior approaches include contrastive fine-tuning and attention-mask changes for bidirectional context.Source: §2 · p. 3
  • Training-free methods use prompts to guide semantic compression or repeat inputs and prepend auxiliary tokens to address causal-context limits.Echo doubles sequence length, while token prepending may use out-of-vocabulary tokens with unpredictable representations.Source: §2 · p. 3
  • KV-Embedding repurposes internal attention states for representation extraction, routing final-token KV states as a prefix so early tokens can access a compressed sequence summary.Related KV approaches instead use learned prefixes for fine-tuning or cache selection for generative throughput.Source: §2 · p. 3

3 Method

The method combines a compression-oriented prompt, internal KV re-routing, intrinsic-dimensionality-based layer selection, and hybrid pooling. Re-routing makes the final token’s sequence summary available to other positions within a forward pass.

3.1 Preliminaries

  • Decoder-only causal attention computes token states without visibility into subsequent tokens, and pooling final-layer states produces a fixed-size embedding.The method section introduces KV-Embedding as a training-free framework for decoder-only LLMs.Source: §3 · p. 3, §3.1 · p. 4

3.2 KV-Embedding

  • KV-Embedding addresses causal-attention limitations through a compression-oriented prompt, internal KV re-routing, and automated layer selection..Source: §3.2 · p. 4

3.2.1 Compression-Oriented Prompting

  • The prompt frames text as context or a query and asks for compression in one word, guiding the final-token representation toward the input’s semantic essence.The paper describes this as a way to mitigate next-token prediction bias.Source: §3.2.1 · p. 4

3.2.2 Key-Value Re-Routing

  • At selected layers, the final-position KV pair is prepended to the layer’s KV matrices, allowing each token to attend to a global sequence summary.The final token’s KV pair aggregates preceding positions; the method also adds an attention bias of 1.0 for the re-routed position.Source: §3.2.2 · p. 4

3.2.3 Automated Layer Selection via Intrinsic Dimensionality

  • Layer selection uses intrinsic dimensionality to target compressed representations while avoiding shallow surface-feature layers and late prediction-biased layers.For multiple local minima, the method selects low-ID regions in the network’s middle-to-late portion and excludes the first 20% of layers.Source: §3.2.3 · p. 4
  • The final embedding averages mean pooling and last-token pooling, combining sequence-wide distributional evidence with the globally informed final state before L2 normalization.The paper reports that hybrid pooling is more consistent across models and tasks.Source: §3.2.3 · pp. 4–5

4 Experiments

The experiments evaluate KV-Embedding in zero-shot, training-free settings on MTEB and long-context LoCoV1 across three decoder-only backbones. KV-Embedding outperforms the training-free baselines on MTEB and maintains an advantage on LoCoV1 at every tested sequence length.

4.1 Experimental Setup

  • The evaluation covers seven MTEB task categories and LoCoV1 documents truncated to 1,024, 2,048, and 4,096 tokens.The three backbones are Qwen3-4B, Mistral-7B-Instruct-v0.1, and Llama-3.1-8B-Instruct, evaluated zero-shot without training.Source: §4.1 · p. 5
  • The comparison includes pooling and prompt-based baselines, Echo, Token Prepending, and an ablation that disables KV re-routing while retaining the prompt.Echo repeats the input, doubling sequence length; the ablation helps separate the effect of re-routing from the prompt.Source: §4.1 · p. 5, §4.2 · p. 5

4.2 Results

  • Retrieval shows the largest MTEB increase: Qwen3-4B reaches 0.2765, compared with 0.1857 for PromptEOL.STS and Clustering also improve consistently, while Summarization shows the smallest variance across methods.Source: §4.2 · p. 5
  • KV-Embedding achieves the best overall MTEB performance across all three backbones; on Mistral-7B, removing re-routing lowers the average from 0.534 to 0.450.The ablation retains the compression-oriented prompt, isolating the contribution of KV re-routing.Source: §4.2 · p. 5, Table 1 · p. 6
  • On LoCoV1, KV-Embedding outperforms all baselines at every tested sequence length; across Qwen3-4B and Llama-3.1-8B, it beats the best baseline by 1.3–3.5×.Mistral-7B scores above 0.18 across settings, while baselines remain below 0.10.Source: §4.2 · pp. 5–6

5 Analysis and Ablation

The analysis tests whether KV-Embedding preserves causal attention while providing global context, and examines attention patterns, embedding geometry, bias strength, layer selection, pooling, and prompt robustness. Results favor re-routing over mask removal and support the paper’s selected bias, ID-based layers, and hybrid pooling.

5.1 Mechanism Validation

  • Removing the causal mask degrades performance below the last-token baseline on Qwen3-4B and Mistral-7B, whereas KV-Embedding relocates existing KV states without changing attention structure.The paper attributes the failure to out-of-distribution states under bidirectional attention in models pretrained with unidirectional patterns.Source: §5.1 · p. 6

5.2 Embedding Analysis

  • KV-Embedding’s attention aligns with both the intent word “Looking” and functional anchors such as “examples” and “beginners,” unlike methods that miss the initial intent-bearing word.The analyzed query concerns machine-learning tutorials with practical code for beginners.Source: §5.2 · pp. 6–7, Table 3 · p. 8
  • On Mistral-7B-Instruct, KV-Embedding achieves the best alignment and uniformity values among baselines, with improved uniformity indicating a more isotropic embedding space.The analysis describes decoder-only representations as tending to cluster within a narrow cone.Source: §5.2 · p. 7, Table 5 · p. 8

5.3 Ablation of Key Components

  • Average performance peaks at attention bias 1.0 and declines above 3.0; the authors use 1.0 to balance global-context access with local specificity.Performance improves from bias 0 to 1.0 and remains stable across a wide range.Source: §5.3 · p. 7, Figure 2 · p. 7
  • ID-based selection scores 0.494 on Qwen3-4B and 0.534 on Mistral-7B; on Mistral-7B, seven selected layers outperform the 10-layer middle-third strategy.Early-layer selection scores 0.409 on Qwen3-4B.Source: §5.3 · p. 7
  • Hybrid pooling is more consistent across models and tasks than mean pooling, and the five tested prompt templates yield scores within 0.529–0.540.The hybrid averages last-token and mean representations; mean pooling tends to underperform, and the authors adopt the hybrid for the main experiments.Source: §5.3 · p. 8

6 Conclusion

KV-Embedding extracts representations from decoder-only LLMs by rerouting final-token KV states as a prefix within selected layers. Evaluations show gains over training-free methods on MTEB and LoCoV1, with robust performance up to 4,096 tokens.

6 Conclusion

  • KV-Embedding reroutes final-token KV states as a prefix within selected layers, letting all positions access a compressed sequence summary in one forward pass.Automated intrinsic-dimensionality-based layer selection supports applicability without manual tuning.Source: §6 · p. 8
  • Evaluations on MTEB and LoCoV1 show up to 10% improvement over training-free methods and robust performance on sequences up to 4,096 tokens.The conclusion frames internal state manipulation as an alternative to external input modification.Source: §6 · p. 8

Limitations

The paper identifies higher latency than standard pooling and a performance boundary relative to supervised embeddings. It also describes the selected-layer procedure and probing evidence that final-token KV states carry sequence-level information.

Limitations

  • KV-Embedding increases latency relative to standard pooling and may not match supervised contrastive models, so it is positioned as an alternative for resource-constrained scenarios rather than a full replacement.Source: Limitations · p. 9

A Algorithm Details

B Probing KV States for Sequence Information

  • Last-position KV probes outperform first-position probes on both sentiment datasets, with IMDB accuracy rising from 51.4% to 85.2%.The reported gain is over 33 absolute points; the other dataset shows a 15.6-point first-to-last gain.Source: B Probing KV States for Sequence… · p. 11

C Layer-wise Representation Analysis

Layer-wise evaluation finds different patterns for pair classification and STS. The intrinsic-dimensionality-selected range is layers 13–19, a transition zone intended to retain semantic abstraction before prediction bias dominates.

C Layer-wise Representation Analysis

  • Pair Classification peaks at layers 10–13 and declines later, while STS improves toward later layers, likely reflecting next-token prediction bias.Source: C Layer-wise Representation Analysis · p. 11
  • ID-selected layers 13–19 preserve strong semantic matching in a transition zone before later layers’ prediction bias dominates.The table caption describes this range as maintaining strong average performance while avoiding generation-biased later layers.Source: Table 7 · p. 11, C Layer-wise Representation Analysis · p. 11

D.1 MTEB Benchmark

The evaluation covers MTEB task categories and LoCoV1 retrieval at multiple context lengths, with detailed results organized by task and benchmark. Across LoCoV1 context windows and model architectures, KV-Embedding has the highest average performance.

D.1 MTEB Benchmark

  • MTEB evaluation spans seven task categories and 42 datasets, using task-specific measures including Spearman correlation, NDCG@10, MAP, accuracy, V-measure, and Average Precision.Source: §D.1 · p. 12

D.2 LoCoV1 Benchmark

  • LoCoV1 tests long-context retrieval under controlled truncation at 1k, 2k, and 4k tokens, with NDCG@10 compared across baselines.Source: §D.2 · p. 12

E Detailed MTEB Results

  • Detailed MTEB results are organized by task: tables report STS, retrieval, classification, pair classification, clustering, reranking, and summarization metrics.Source: E Detailed MTEB Results · p. 12

F Detailed LoCoV1 Results

G Attention Bias Experiment

The attention-bias experiment finds a stable performance range and uses b = 1.0 to balance global context with local detail. Geometry metrics and model-specific intrinsic-dimensionality trajectories characterize the representations and motivate adaptive layer selection.

H Geometry of the Embedding Space

  • Alignment measures closeness for semantically similar pairs, while uniformity measures distribution across the unit hypersphere; together they assess local clustering and global dispersion.The metrics use 1,000 F2LLM sentence pairs and ℓ2-normalized embeddings.Source: H Geometry of the Embedding Space · p. 16

I Intrinsic Dimensionality Trajectories

  • Figure 3 shows distinct ID trajectories for Mistral-7B and Qwen3-4B, with ID rising again near the final layers in both models.Mistral has a stable low-ID region around layers 13–19; Qwen3-4B descends more gradually, reaching its global minimum around layers 17–19.Source: Figure 3 · p. 16, I Intrinsic Dimensionality Trajectories · p. 16
  • Because the layers of maximal compression differ by model, ID-based selection adaptively identifies where information is most condensed rather than relying on a fixed layer heuristic.The passages report different locations for the models’ minimum ID: early for Mistral-7B and in the middle layers for Qwen3-4B.Source: I Intrinsic Dimensionality Trajectories · pp. 16, 18

G Attention Bias Experiment

  • 0.519 to 0.534: MTEB performance stays within ±1.5% across a wide bias range, but declines above b = 3.0; the method sets b = 1.0.The authors attribute the decline to possible over-reliance on the global summary at the expense of local details.Source: G Attention Bias Experiment · p. 16, Table 18 · p. 17

J Analysis of Layer Selection

ID-selected layers achieve the best average performance with fewer layers than uniform strategies. The analysis also examines causal-attention constraints, prompt variation, and pooling choices after KV re-routing.

K Impact of Attention Constraints

  • Forcing bidirectional attention gives frozen models future-position KV states unseen during pre-training; performance collapses across both backbones, with Retrieval scores near zero (0.02–0.03).KV-Embedding instead prepends states from the final token, which has attended to preceding positions, without changing the causal attention mask.Source: K Impact of Attention Constraints · p. 18

L Prompt Design

  • Mistral-7B-Instruct averages remain between 0.529 and 0.540 across five prompt variations, indicating stable performance across instruction phrasings.Prompt 0 leads on STS, Classification, and Summarization; Prompt 2 has a slight edge on Retrieval and Clustering.Source: L Prompt Design · p. 18, Table 20 · p. 20

J Analysis of Layer Selection

  • ID-selected layers achieve the highest average performance using fewer layers; on Qwen3-4B, layers 12–21 score 0.494, outperforming the broader middle-third range 12–23.The early-layer range 0–11 scores 0.409 on Qwen3-4B.Source: J Analysis of Layer Selection · p. 18, Table 19 · p. 20

M Pooling Strategy

  • Hybrid pooling of the last-token and mean representations achieves the highest average scores, while mean pooling performs worst across backbones.The authors suggest that averaging can retain early-token lexical noise, whereas hybrid pooling combines sequence-wide evidence with the final state's globally informed summary.Source: M Pooling Strategy · pp. 18–19, Table 21 · p. 20

Related papers

Loading 2601.01046v1…