Reproduce “GEM: Empowering LLM for both Embedding Generation and Language Understanding” on OpenResearch

About this paper

Caojin Zhang, Qiang Zhang, Ke Li, Sai Vidyaranya Nuthalapati, Benyu Zhang, Jason Liu, Serena Li, Lizhu Zhang and 1 more

Meta researchers develop GEM (Generative Embedding large language Model), a framework that enables decoder-only LLMs to generate high-quality text embeddings while preserving their original language generation and reasoning capabilities through a computationally efficient self-supervised approach requiring only 32,000 training samples, achieving comparable MTEB embedding performance to specialized models (25.6 vs 4.38 retrieval score for Llama 3.2 1B) with minimal degradation in language understanding (MMLU), outperforming methods like Llm2vec that sacrifice generative abilities and GritLM that requires 10x more training data, by inserting special tokens into input text and modifying causal attention masks to create information bottlenecks where suffix tokens can only attend to special tokens rather than the original prefix, forcing compression of semantic information into the special tokens which are then extracted as embeddings, with mixed next-token prediction and contrastive learning objectives preventing catastrophic forgetting while enhancing embedding quality across Llama and Mistral model families from 1B to 8B parameters.

Key insights

• Information bottleneck design through attention mask manipulation effectively forces LLMs to compress semantic information into special tokens without architectural changes
• Mixed training with regular text (controlled by mix ratio p=0.8) prevents catastrophic forgetting while enabling embedding capability acquisition
• Self-supervised contrastive learning with dropout perturbations enhances embedding quality while preserving causal attention necessary for generation

Summary by alphaXiv.