Learning to Compress: Unlocking the Potential of Large Language Models for Text Representation

AI summary

Overview

Research area: Natural Language Processing — text representation / sentence embeddings, specifically unsupervised adaptation of large language models (LLMs) into text encoders.

Technical level: Intermediate. The paper is readable with a basic understanding of embeddings, contrastive learning, and LLM fine-tuning, but the dimensional-collapse analysis and KL-divergence objective require some background.

Scope: The paper proposes context compression as a new unsupervised pretext task for turning causal LLMs into text embedding models, introduces the CTKD objective, and analyzes why it outperforms token-level alternatives.

What This Paper Is About

Most LLMs are causal models trained to predict the next token, which makes them poor at producing a single holistic vector that represents a whole sequence. Prior work adapts them for text representation using token-level pretext tasks such as masked next-token prediction (MNTP in LLM2Vec) or embedding-based auto-encoding/auto-regression (EBAE/EBAR in Llama2Vec). This paper asks whether a sequence-level task — learning to compress a whole context into a few compact "memory tokens" that can substitute for the original context — is a better pretext task for representation learning.

Key Contributions

1. The authors investigate context compression as a pretext task for adapting LLMs to text representation, and identify empirically that the best-performing objective is Continuation Task with Knowledge Distillation (CTKD) — a continuation objective trained with a KL-divergence loss against the frozen LLM's own next-token distribution — rather than reconstruction (RC) or continuation with a plain NLL loss (CT-NLL).
2. They determine an appropriate number of memory tokens for downstream representation, reporting that performance is stable for token counts in the range [1, 8] except for retrieval, and that retrieval drops clearly at 16 tokens — contrasting with compression literature where token counts on the order of 100 help generation.
3. They analyze why CTKD wins, showing that models trained with CT-NLL suffer more severe dimensional collapse than CTKD-trained models, both in terms of low-rank pooled embeddings and highly correlated memory tokens.
4. They show that contrastive post-training (unsupervised contrastive learning followed by supervised contrastive learning) further alleviates dimensional collapse, producing LLM2Comp, which outperforms contemporary LLM-based encoders while using substantially less supervised data (0.36M samples vs. 1.16M for LLM2Vec).

Main Findings

• Reconstruction is weak. LLM2Comp_RC performs only slightly better than simple last-token pooling (LT), offering marginal gains. In the supervised contrastive learning stage its average score is 65.49, well below LLM2Comp_KL's 66.78.
• Continuation helps, CTKD helps most. After compression pretraining alone (Table 1), LLM2Comp_NLL reaches 46.22 average, matching Llama2Vec, while LLM2Comp_KL reaches 52.49 — beating LLM2Vec's 46.87, MetaEOL's 50.83, and all other listed baselines.
• CTKD trains more stably. Across five runs, standard deviation is 1.37 for LLM2Comp_KL versus 2.65 for LLM2Comp_RC and 5.32 for LLM2Comp_NLL. LLM2Comp_NLL can reach 51.85 but drops to 42.95 in unfavorable cases.
• CT-NLL collapses more. On embeddings from 60,000 samples of SciDocsRR, the number of effective dimensions is on the order of 10 for LLM2Comp_NLL and 100 for LLM2Comp_KL, against a total dimension of 4096. The KL divergence appears to act as a regularizer that preserves information from less frequent tokens.
• Token correlation hurts. LLM2Comp_NLL memory tokens are highly similar at 32K training samples; the issue is partly alleviated at 128K, with performance rising from 42.95 to 48.38. Removing redundant tokens via clustering (merging clusters whose minimum cross-cluster similarity exceeds 0.9) raises the 128K model from 48.38 to 51.09 — but a single-token setting performs unsatisfactorily, suggesting too much information loss.
• Contrastive stages add large gains. In the UCL stage, LLM2Comp_KL averages 58.51 versus LLM2Vec's 57.82; in the SCL stage it averages 66.78 versus LLM2Vec's 66.11, Instructor's 63.76, BGE-ICL's 65.43, and ULLME (Llama-3) 61.05. The contrastive stages matter most for retrieval and clustering.
• Greater sample efficiency and faster convergence. LLM2Comp peaks with 0.36M training samples versus 1.16M for LLM2Vec. In the SCL stage it reaches optimal performance within 200 steps and stays stable for most tasks, though retrieval declines if contrastive learning continues beyond that point.
• Effective dimensionality grows with each stage. LLM2Comp_KL + UCL has a higher effective dimension than LLM2Comp_KL alone, and adding SCL raises it further, correlating with the performance gains.

Methodology in Plain English

The authors start from a causal LLM (meta-llama/Llama-2-7b-chat-hf) and convert it into a bidirectional encoder. They expand the vocabulary by adding 8 special memory tokens, growing the vocabulary from 32,000 to 32,008, and keep the embedding layer fully trainable while adapting the rest of the model with LoRA (rank 16, alpha 32) applied to the attention query/value/output projections and the feedforward up/down/gate projections.

The encoder reads a context and produces embeddings for the memory tokens. A frozen copy of the original LLM then tries to use those memory tokens as a substitute for the context. In the reconstruction variant, the frozen LLM must regenerate the original text; in the continuation variant (CT-NLL) it must generate a held-out continuation using the negative log-likelihood from the frozen LLM; in CTKD the encoder is additionally trained to make the frozen LLM's next-token distribution conditioned on the compressed memory match its distribution conditioned on the original context, using a KL-divergence loss.

Sentences are embedded by mean pooling the memory token embeddings. Pretraining used 32,000 samples from English Wikipedia-103 (Wikitext-103), 8,000 steps at batch size 4, learning rate 1e-4, weight decay 1e-5, bfloat16, DeepSpeed ZeRO-0, and a warm-up decay schedule with minimum learning rate 1e-5; the authors report results for a stronger seed (2026) and a weaker seed (42).

Post-training follows SimCSE-style unsupervised contrastive learning with dropout rate 0.2 and batch size 128 for 1,000 steps over 128,000 Wikipedia sentences on a single H800, then supervised contrastive learning with in-batch and hard negatives (hard negatives pre-chosen by the E5 dataset from a cross-encoder) over 1,024,000 samples. Dimensional collapse is measured with SVD on the covariance matrix of embeddings from the 60,000 SciDocsRR samples, and token redundancy is measured with an averaged correlation matrix across memory tokens.

Why This Matters

Impact on research: The paper reframes the design space for LLM text encoders — from token-level pretext tasks to sequence-level compression — and supplies a mechanistic explanation (dimensional collapse) for why one compression objective beats another. It also suggests that strong embeddings can be obtained with far less supervised data than current recipes, which changes how researchers might budget annotation effort.

Real-world applications:

• Retrieval and search: the paper evaluates retrieval on datasets spanning biomedical (Bior.), medical (Medr.), Twitter (Twen.), scientific fact-checking (SciF.), news fact-checking (NFCo.), and argument retrieval (Argu.).
• Clustering and deduplication: evaluated on the MTEB clustering tasks.
• Semantic textual similarity: evaluated on STS17, SICK-R, and STSB, relevant to paraphrase detection and duplicate question matching.
• Classification and reranking: evaluated on banking and emotion classification datasets and on StackExchange, SciDocs, and SciFact reranking — relevant to content moderation, support routing, and reordering search results.

Industry relevance: Embedding models are the backbone of retrieval-augmented generation, recommendation, and semantic search. A model that is built on the widely available Llama-2 backbone, requires only a small amount of supervised data, and converges in roughly 200 contrastive steps points to lower training cost and faster iteration cycles for production embedding services.

Future Directions

• Theory for the collapse connection. The authors state that their conclusions are primarily empirical and that a deeper theoretical analysis is needed to formally establish the connection between the InfoNCE loss and dimensional collapse in both unsupervised and supervised contrastive learning.
• Better balance of information and redundancy. The single-token setting performs unsatisfactorily, and LLM2Comp_KL's effective dimensionality remains small relative to the total dimension of 4096, which the authors say leaves room for improvement.
• Contrastive learning beyond convergence. Performance on retrieval declines when contrastive learning continues beyond roughly 200 steps; the authors note that InfoNCE with fixed negative sampling becomes less effective and call this an interesting direction for future research.
• Extending the token length trade-off. Since retrieval is the task most sensitive to the number of memory tokens, and the 16-token setting already degrades retrieval, the paper implies a need for compression objectives that scale to longer memory-token counts without this penalty.

Target Audience

Researchers and engineers working on text embedding models, LLM adaptation, or retrieval systems who want to understand pretext-task design for turning generative LLMs into encoders. It is also relevant to practitioners seeking sample-efficient alternatives to large-scale contrastive training recipes, and to anyone studying representation collapse in self-supervised learning. Readers need basic familiarity with sentence embeddings, contrastive learning, and LoRA fine-tuning.

Authors’ abstract

Text representation plays a critical role in tasks like clustering, retrieval, and other downstream applications. With the emergence of large language models (LLMs), there is increasing interest in harnessing their capabilities for this purpose. However, most of the LLMs are inherently causal and optimized for next-token prediction, making them suboptimal for producing holistic representations. To address this, recent studies introduced pretext tasks to adapt LLMs for text representation. Most of these tasks, however, rely on token-level prediction objectives, such as the masked next-token prediction (MNTP) used in LLM2Vec. In this work, we explore the untapped potential of context compression as a pretext task for unsupervised adaptation of LLMs. During compression pre-training, the model learns to generate compact memory tokens, which substitute the whole context for downstream sequence prediction. Experiments demonstrate that a well-designed compression objective can significantly enhance LLM-based text representations, outperforming models trained with token-level pretext tasks. Further improvements through contrastive learning produce a strong representation model (LLM2Comp) that outperforms contemporary LLM-based text encoders on a wide range of tasks while being more sample-efficient, requiring significantly less training data. Code is available at https://github.com/longtaizi13579/LLM2Comp.