Source-linked AI summary

Following the Autoregressive Nature of LLM Embeddings via Compression and Alignment

Jingcheng Deng, Zhongtao Jiang, Liang Pang, Liwei Chen, Kun Xu, Zihao Wei, Huawei Shen, Xueqi Cheng

arXiv:2502.11401v3cs.CL

TL;DR

LLM embeddings' autoregressive next-token orientation conflicts with contrastive learning's need for global semantic, cosine-aligned representations, limiting efficient use of LLM pretraining. AutoRegEmbed addresses this through information compression and conditional distribution alignment, and achieves strong performance with fewer samples while matching or exceeding reported baselines. Its embeddings may inherit biases from harmful or discriminatory training data, requiring careful curation and caution in sensitive domains.

  • Problem

    Autoregressive LLM embeddings emphasize next-token semantics rather than global input semantics, creating a mismatch with contrastive learning and inefficient use of LLM capabilities.

  • Method

    AutoRegEmbed combines information compression with conditional distribution alignment over embedding conditional probability distributions.

  • Results

    AutoRegEmbed outperforms traditional contrastive methods and achieves performance comparable to SOTA models with fewer training samples.

  • Takeaways & Limitations

    AutoRegEmbed provides an autoregressive-compatible contrastive approach that improves representation-learning efficiency under limited data.

  • Takeaways & Limitations

    AutoRegEmbed does not filter harmful content and may carry biases from problematic training data into embeddings, especially in sensitive applications.

Abstract

from arXiv · show

A new trend uses LLMs as dense text encoders via contrastive learning. However, since LLM embeddings predict the probability distribution of the next token, they are inherently generative and distributive, conflicting with contrastive learning, which requires embeddings to capture full-text semantics and align via cosine similarity. This discrepancy hinders the full utilization of LLMs' pre-training capabilities, resulting in inefficient learning. In response to this issue, we propose AutoRegEmbed, a new contrastive learning method built on embedding conditional probability distributions, which integrates two core tasks: information compression and conditional distribution alignment. The information compression task encodes text into the embedding space, ensuring that the embedding vectors capture global semantics. The conditional distribution alignment task focuses on aligning text embeddings with positive samples embeddings by leveraging the conditional distribution of embeddings while simultaneously reducing the likelihood of generating negative samples from text embeddings, thereby achieving embedding alignment and uniformity. Experimental results demonstrate that our method significantly outperforms traditional contrastive learning approaches and achieves performance comparable to state-of-the-art models when using the same amount of data.

1 Introduction

AutoRegEmbed addresses the mismatch between autoregressive LLM embeddings and cosine-based contrastive learning by compressing text into globally semantic embeddings and aligning their conditional distributions. It reportedly outperforms traditional contrastive methods and matches state-of-the-art performance with fewer samples.

  • LLM final-token embeddings can represent next-token semantics rather than the global semantics of the input text, limiting contrastive embedding quality.
  • AutoRegEmbed combines information compression with conditional distribution alignment to satisfy global-semantic, alignment, uniformity, and autoregressive requirements.
  • Information compression forces a frozen decoder to reconstruct targets from compressed variables alone, creating a bottleneck that preserves context and instruction semantics.
  • Conditional distribution alignment compares embedding-conditioned probabilities, increasing positive-sample likelihood while decreasing negative-sample likelihood.
  • AutoRegEmbed achieves performance comparable to SOTA models with fewer training samples and superior learning efficiency.Its Pareto frontier outperforms traditional contrastive learning methods, indicating a better efficiency–performance trade-off.

2 Related Works

Related work spans conventional encoder-only contrastive models, fine-tuned LLM embedding models, and methods that use LLMs without fine-tuning. These approaches differ in architecture, training requirements, and computational cost.

  • Text embedding research is organized into early models, fine-tuned LLMs, and LLMs without fine-tuning.
  • Early Models: Early systems such as SentenceBERT and SimCSE use contrastive learning with small encoder-only models, while instruction-based methods target broader transferability.
  • LLMs with Fine-Tuning: Fine-tuned LLM methods use strategies including last-token extraction, bidirectional attention with average pooling, latent attention pooling, in-context learning, and synthetic data.
  • AutoRegEmbed claims strong results without traditional cosine-based contrastive fine-tuning by exploiting LLMs' autoregressive potential.
  • LLMs without Fine-Tuning: Without fine-tuning, prior methods use trajectory distributions or repeated autoregressive inputs, but these approaches can incur high or doubled computational costs.

3 Method

AutoRegEmbed adapts LLM embeddings through information compression and conditional distribution alignment, shifting representations toward global semantics while preserving the models’ autoregressive nature. Its alignment objective compares conditional probability distributions rather than cosine distances between embedding points.

  • Method: AutoRegEmbed combines information compression with conditional distribution alignment to produce instruction-aware embeddings capturing global semantics.The framework first compresses context and instructions into compressed tokens, then aligns their conditional distributions.
  • Motivation: Decoder-only LLM embeddings may overemphasize next-token semantics because final-token hidden states do not encode the input’s global semantics.Average pooling does not necessarily recover global semantics, while attention pooling alters the original architecture.
  • Information Compression: Information compression appends a small set of compressed tokens and trains a decoder to reconstruct the target using only those variables, creating an information bottleneck.The encoder processes the text, instruction, and compressed tokens, while the decoder reconstructs the target from the resulting compressed representation.
  • Conditional Distribution Alignment: Conditional distribution alignment replaces cosine similarity with divergence between the conditional distributions associated with text and sample embeddings.This extends point alignment to distribution alignment while addressing the mismatch between autoregressive pre-training and conventional contrastive objectives.
  • Conditional Distribution Alignment: The final objective aligns a query’s distribution with a positive sample and increases the relative likelihood of generating positives over negatives.S1 measures the difference between positive-generation log probabilities, while S2 amplifies the positive–negative probability gap and normalizes against the initial model.

4 Experiments

AutoRegEmbed is evaluated across semantic similarity, retrieval, and efficiency settings, with ablations testing its two tasks and loss-function design. It consistently outperforms comparable contrastive-learning baselines, matches prior state-of-the-art performance with far less data, and improves as both compression and conditional distribution alignment contribute.

  • 4.2 Main Results: Compared with unsupervised LLM2Vec, AutoRegEmbed is 4.74% better despite LLM2Vec using almost 160,000 samples, while its own fair untrained baseline averages 20% lower.Prompt-optimization methods also provide limited gains while incurring additional reasoning costs.
  • 4.2 Main Results: AutoRegEmbed surpasses leading supervised contrastive methods across 10 STS datasets, outperforming LLM2Vec by 0.58 while using only 66,382 training samples versus tens of millions for prior SOTA models.The comparison uses Table 1’s STS evaluation and emphasizes data efficiency.
  • 4.3 Performance on retrieval tasks: AutoRegEmbed consistently outperforms LLaMA2-inbatch-M across MS MARCO, NFcorpus, and SCIDOCS, despite both methods using the same training data.It is second only to gte-Qwen2-7B-instruct on MS MARCO and remains competitive on NFcorpus and SCIDOCS.
  • 4.4 Ablation Study: Conditional Distribution Alignment improves performance by 9.17%, while Information Compression contributes a 16.99% improvement across the ablation tasks.Table 3 evaluates the contribution of both components.
  • 4.4 Ablation Study: The original Equation 2 loss outperforms logarithmic-similarity, KL-divergence, and JS-divergence variants despite those alternatives having more intuitive interpretations.Thus, the reported alternatives do not improve on the original loss formulation.
  • 4.4 Ablation Study: With 15,000 samples, AutoRegEmbed already exceeds the maximum performance of other contrastive-learning models, indicating higher learning efficiency under matched sample budgets.Both AutoRegEmbed and competing methods improve with more data, but AutoRegEmbed grows fastest.

5 Conclusions

AutoRegEmbed addresses the mismatch between autoregressive LLMs and traditional contrastive learning through conditional probability distributions, combining information compression with conditional distribution alignment. It achieves comparable performance to state-of-the-art models with fewer training samples and superior learning efficiency.

  • AutoRegEmbed uses embedded conditional probability distributions to capture global semantics while maintaining embedding alignment and uniformity.Its two tasks are information compression and conditional distribution alignment.
  • The method achieves comparable performance to state-of-the-art models with fewer training samples and superior learning efficiency.

6 Limitations

AutoRegEmbed lacks inherent mechanisms to detect or filter harmful content, so problematic training data can propagate bias into its embeddings. The authors recommend curated, filtered data and caution in sensitive domains.

  • AutoRegEmbed cannot inherently filter or detect malicious or harmful content in the data it processes.It therefore cannot evaluate the ethical or safety implications of its inputs.
  • Harmful or discriminatory training material may carry bias into generated embeddings and produce undesirable outcomes in real-world tasks.
  • Users should curate and ideally filter training data, while exercising caution in sensitive domains.Unsafe or biased embeddings could have significant consequences in those settings.

A Implementation details

The implementation uses separate optimization settings for information compression and conditional distribution alignment, with fixed batch sizes and bounded sequence lengths. Baseline comparisons also tune in-batch negative sampling and synchronize gradients across devices.

  • Information compression uses a 2e-5 learning rate, batch size 32, 2 epochs, and 5 compressed tokens, while alignment uses 5e-6, batch size 32, and 4 epochs.The temperature parameters are τ=0.05 and β=0.1 for both tasks.
  • Both tasks cap context, instruction, and target sequences at 512 tokens.
  • The fair baseline uses batch size 512 for in-batch negatives and propagates gradients across devices.Batch sizes 128, 256, 512, and 1024 were tested, with 512 performing best.

B Variants of Equation 2

The variants examine alternative similarity functions and distribution-distance measures for Equation 2. The authors hypothesize that log-odds ratios provide stronger targeted supervision than KL or JS divergence, while Tanh and Sigmoid perform similarly overall.

  • Similarity variants: The similarity function can replace sigmoid with Log-Sigmoid, or measure distance using KL or JS divergence.
  • Divergence variants: KL and JS divergence may underperform because they distribute supervision across large vocabularies, weakening gradients and degrading performance.The authors hypothesize that log-odds ratios instead supervise the generation probability of a specific label token.
  • Activation variants: Under identical settings, Tanh slightly outperforms Sigmoid on 7 STS datasets, whereas Sigmoid has a marginal advantage across all 10 datasets.Given the small gap, either function is considered reasonable.

D Explanation of Inext and Iself

Embedding tasks are divided into sentence-to-sentence similarity and sentence-to-document retrieval, with distinct instructions for each task.

  • STS compares sentences symmetrically, so its query and self instructions are identical.Retrieval is asymmetric and uses separate instructions for the query and document roles.

E Analysis of Alignment Strategies for Conditional Probability Distributions

AutoRegEmbed aligns conditional distributions by matching a query's generation of positive text with the positive text's self-generation, while using negative-sample feedback. Several alternative alignment strategies were tested, and simpler modifications performed better.

  • Alignment mechanism: Conditional-distribution alignment matches q generating d+ with d+ generating itself, while increasing distance from negative samples.The method uses instructions I_next and I_self to define these generation relationships.
  • Alignment mechanism: The σ function comparison evaluates alternative bounded transformations for distribution-based distances under identical settings.Sigmoid and Tanh are identified as the typical choices.
  • Experimental setup: The alignment-strategy study evaluates LLaMA2-7B on the MEDI STS split with τ and β fixed at 0.1.Table 6 reports performance under the different strategies.
  • Alignment strategies: Strategy 1 changes the generation target from d+ to q, whereas Strategy 2 averages the conditional probabilities of q and d+.For symmetric STS, the instruction roles coincide; retrieval uses an asymmetric previous-instruction example.
  • Alignment strategies: Strategies 3 and 4 introduce more complex positive- and negative-sample alignment designs, but they produce inferior performance.These strategies modify generation probabilities involving positive and negative samples.

F Analysis of Temperature Coefficients τ and β

The temperature coefficients τ and β control different aspects of the alignment loss, and grid-search results indicate generally stable performance except at extreme settings.

  • Coefficient roles: τ controls positive-negative separation strength, with lower values enforcing sharper contrast in the loss.Its role is described as aligning the objective more closely with InfoNCE.
  • Coefficient roles: β controls sensitivity to differences between conditional distributions by scaling log-probability-ratio differences.Its design is inspired by the scaling factor in DPO loss.
  • Hyperparameter analysis: The grid search evaluates τ values from 0.02 to 1.0 and β values from 0.1 to 0.4 on LLaMA2-7B using MEDI's STS split.The search identifies the best-performing coefficient pair.
  • Hyperparameter analysis: Performance improves when τ = 0.05 or β = 0.2 and remains relatively stable except for a few extreme settings.The reported exceptions include τ = 1.0 and β = 0.4.
  • Hyperparameter analysis: Applying τ = 0.05 and β = 0.2 simultaneously does not improve performance, suggesting their effects may not be orthogonal.
Loading 2502.11401v3…