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

TL;DR

LLMs are typically optimized for next-token prediction, limiting their ability to produce holistic sequence representations. This work introduces context compression as a pretext task, training an encoder to generate memory tokens that replace the full context, with CTKD (Continuation Task with Knowledge Distillation) providing the most effective objective. Post-training through unsupervised and supervised contrastive learning mitigates dimensional collapse, yielding LLM2Comp, a robust, data-efficient text encoder that outperforms LLM2Vec and Llama2Vec on diverse tasks. The approach offers practical benefits for clustering, retrieval, and semantic similarity, while requiring substantially less labeled data and achieving faster convergence than competing methods.

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.

Figures (12)

• Figure 1: The comparison of different pretext tasks: (1) MNTP (Mask Next Token Prediction) in LLM2Vec; (2) EBAE or EBAR in Llama2Vec; (3) Context Compression task in LLM2Comp (ours)
• Figure 2: Mean and standard variation of LLM2Comp$_{RC}$, LLM2Comp$_{NLL}$, and LLM2Comp$_{KL}$ across different task types, computed over five runs.
• Figure 3: LLM2Comp$_{KL}$: Token length and its Effect
• Figure 4: Comparing singular values of LLM2Comp$_{NLL}$ and LLM2Comp$_{KL}$.
• Figure 5: $\textbf{Llm2Comp}_{NLL}$ training with 32000 samples (left) and 128000 samples (right) in a bad training case.