Technical Summary: The Truth Lies Somewhere in the Middle (of the Generated Tokens)

Problem Statement

In autoregressive language models (LLMs), hidden states are generated sequentially under causal masking, meaning each token's representation depends on the prompt and all preceding generated tokens. A critical methodological question arises: how should these temporally distributed hidden states be collapsed into a single vector representation that accurately reflects the model's internal semantic state?

Prior work on bidirectional encoders typically employs mean pooling or final-token pooling. However, in autoregressive settings, simple averaging has been shown to perform poorly because it mixes hidden states computed under unequal contexts. Consequently, existing approaches often rely on selecting a single token (e.g., the last token) or training models with specific objectives to improve embedding quality. It remains unclear whether semantic information in standard, pre-trained autoregressive models is localized to a specific token position or distributed across the generation trajectory, and whether mean pooling can yield superior semantic representations without additional training.

Methodology

The authors evaluate the semantic quality of representations derived from generated tokens by measuring their alignment to reference spaces in three distinct domains: vision-language, reasoning, and protein structures.

Evaluation Metric: Kernel Alignment

The core metric is Centered Kernel Alignment (CKA), specifically the debiased version to avoid inflation artifacts. The authors compare the similarity structure (kernel) induced by the language model's embeddings against a reference kernel constructed from ground-truth semantic structures:

• Vision-Language: Reference kernels are derived from image embeddings (using DINOv2) corresponding to image captions. High alignment implies the model places captions of visually similar images closer together.
• Reasoning: Reference kernels are derived from gold-standard solutions to math and science problems.
• Protein: Reference kernels are derived from protein structure embeddings (using ESM-3).

Experimental Setup

• Models: The primary experiments use Qwen3-14B (a decoder-only Transformer). Comparisons are made with OLMo3-7B-Think and gpt-oss-20B.
• Pooling Strategies: The authors compare last-token pooling (selecting the final state of the trajectory) against mean-token pooling (averaging hidden states across the entire generated sequence).
• Mixing Analysis: To determine if semantic information is localized or distributed, the authors construct convex combinations of embeddings from distinct token segments (e.g., first, middle, and final thirds of the generation) and evaluate alignment across the resulting simplex.
• Control Conditions: Experiments include shuffling token orders, shuffling image-text pairings, and embedding text generated by one model using the embedding space of another to isolate the effects of generation dynamics.

Key Contributions and Results

1. Mean Pooling Generated Tokens Outperforms Single-Token Selection

Contrary to the limitations observed in prompt-side embeddings, mean pooling across generated tokens yields significantly higher semantic alignment than any individual token, including the last token.

• In vision-language tasks, mean-pooled generated representations achieved an alignment score of 0.410, compared to 0.291 for the best single token and 0.184 for prompt tokens.
• This improvement is consistent across decoding seeds, indicating the underlying semantic structure is stable despite text variation.
• The benefit holds across different pooling methods (mean pooling outperforms attention, max, and best-per-token pooling).

2. Semantic Information is Distributed Across the Generation

The paper demonstrates that semantic information is not localized to a single position but is distributed across the token trajectory.

• Mixing Effect: Convex combinations of embeddings from different token segments (e.g., mixing the first, middle, and final thirds) consistently achieve higher alignment than any single segment alone.
• Granularity: This "mixing" benefit persists even when token segments are recursively subdivided, suggesting that different parts of the generation carry complementary semantic information.
• Domain Generalization: This phenomenon is observed not only in vision-language tasks but also in reasoning (Math-500, GPQA) and protein structure tasks, indicating it is a general property of autoregressive generation.

3. Generated Tokens vs. Prompt Tokens

A counter-intuitive finding is that representations derived from generated tokens are better aligned than those from prompt tokens, even when the prompt contains the full semantic content.

• Prompt-token representations do not benefit from mean pooling; averaging them often dilutes information due to causal masking limitations.
• Generation appears to "make the same semantic content easier to recover" in the hidden state space. As the model produces a continuation, its internal states become better aligned with the reference structure tied to the original prompt.

4. Representational Phases and Internal State Dynamics

The alignment trajectory across generated tokens reveals interpretable phases of model behavior:

• Phases: The generation process exhibits distinct phases: generic preamble, prompt repetition, recall, and caption-specific response. Alignment increases as the generation moves from generic text to content-specific descriptions.
• Causal Probes: Specific generic phrases (e.g., "Let me recall what I know") induce sharp spikes in alignment, even before scene-relevant content appears. This suggests that such tokens correspond to shifts in the model's internal state toward the reference structure.
• Model Specificity: These phase structures disappear when a model embeds text generated by a different model. For instance, Qwen3's alignment trajectory degrades when processing OLMo3's generations if OLMo3 hallucinates facts (e.g., misidentifying the location of a power station). This indicates that the observed dynamics are tied to the specific model's internal computation and consistency, not just the text itself.

5. Layer Depth vs. Token Dimension

While mean pooling across the token dimension improves alignment, mean pooling across layer depth does not. Averaging across layers yields performance comparable to the best single layer (typically the final layer) but does not improve upon it. This contrasts with prompt-side embeddings, where intermediate layers often perform best.

Significance and Claims

The paper claims that hidden states formed during autoregressive generation are best collapsed into a vector representation by mean pooling across tokens, rather than selecting a single token. This challenges the prevailing view that causal masking makes mean pooling ineffective for decoder-only models.

The authors argue that:

1. Generation as Context: The model's own continuation acts as a source of context that makes the input's semantic content more accessible in hidden-state space.
2. Trajectory-Based Representation: Semantic information in autoregressive models should be viewed as distributed across the generation trajectory rather than localized to a single hidden state.
3. Interpretability: Kernel alignment serves as a causal probe for internal model states, revealing that specific token outputs (like "thinking" or "recall" phrases) correspond to measurable shifts in the model's representational geometry.

The work suggests that for tasks requiring high-quality text embeddings from autoregressive models, generating a continuation and mean-pooling the resulting hidden states is a simple, effective, and training-free method that outperforms standard prompt-based embeddings. However, the authors note limitations: this approach incurs a computational cost proportional to the generation length and produces representations that are probes of internal state rather than "soft" tokens usable for further generation.