kv-embedding-training-free

name	kv-embedding-training-free
title	KV-Embedding: Training-free Text Embedding via Internal KV Re-routing in Decoder-only LLMs
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2601.01046
keywords	["Text Embeddings","Decoder-only Models","Representation Learning","KV Caching","Model Internals"]
description	Extract high-quality embeddings from frozen decoder-only LLMs by re-routing internal key-value states without training—outperforming training-free baselines by 10% on MTEB while maintaining robustness across sequences up to 4,096 tokens.

Overview

KV-Embedding solves a critical limitation: while decoder-only LLMs (Qwen, Mistral, Llama) excel at generation, extracting semantic embeddings from them is challenging. Their causal attention masks early tokens from later context, and their objective (next-token prediction) biases representations toward generation over semantic compression.

Core Innovation: Leverage internal key-value (KV) cache states that naturally encode sequence-level information. By re-routing these states as prepended prefixes, enable all tokens to access full sequence context within a single forward pass—activating latent embedding capabilities without training.

Problem Statement

Causal Attention Limitation: Standard decoder-only attention masks are causal—early tokens cannot attend to future tokens. This is efficient for generation but problematic for embeddings.

Next-Token Prediction Bias: Models optimized for generation focus on predicting the next token. Embeddings trained this way emphasize generative features (fluency, diversity) over semantic compression.

Training Complexity: Typical embedding approaches require fine-tuning with contrastive losses, adding computational overhead and reducing accessibility.

KV-Embedding Approach

Stage 1: Extract KV States

During forward pass, capture key-value cache states from each transformer layer:

def extract_kv_states(model, input_ids):
    """Extract KV cache states during forward pass."""

    hidden_states = []
    with torch.no_grad():
        for layer_idx, layer in enumerate(model.model.layers):
            # Forward through layer, capturing KV cache
            out, cache = layer(
                input_ids,
                attention_mask=...,
                past_key_values=...
            )
            hidden_states.append(cache)

    # KV states from final layer contain full sequence information
    return hidden_states[-1]

Stage 2: Re-route as Prefix

Use final layer's KV states as a prepended prefix in a second forward pass:

Key Insight: The KV states of the final token encode a compressed view of the entire sequence. By prepending these states, all tokens gain access to sequence-level context.