GRACE: Generative Representation Learning via Contrastive Policy Optimization
Jiashuo Sun, Shixuan Liu, Zhaochen Su, Xianrui Zhong, Pengcheng Jiang, Bowen Jin, Peiran Li, Weijia Shi, Jiawei Han
2025-10-08
Summary
This paper introduces a new way to train Large Language Models (LLMs) to better understand and represent text, moving beyond simply creating static summaries of words. It focuses on making the LLM's reasoning process more transparent.
What's the problem?
Current methods for training LLMs to understand text treat them like 'black boxes'. They focus on getting the model to create good numerical representations of words (embeddings) but ignore the model’s ability to actually *think* about the text and explain its understanding. This means we don't know *why* the model thinks something is similar to something else, and it limits how well the model can truly grasp meaning.
What's the solution?
The researchers developed a framework called GRACE, which reframes how LLMs are trained. Instead of just minimizing errors in creating embeddings, GRACE uses a 'reward' system. The LLM is prompted to generate explanations – like short, human-readable reasons – for why two pieces of text are related. These explanations are then used to create the embeddings. The model is rewarded for generating explanations that make sense and accurately reflect the relationship between the texts. This process essentially turns the LLM into an 'agent' that can explain its reasoning.
Why it matters?
This work is important because it improves the quality of text representations created by LLMs, leading to better performance on various tasks. More significantly, it makes these models more interpretable. We can now see *why* the model makes certain connections, which builds trust and allows us to better understand and improve these powerful AI systems.
Abstract
Prevailing methods for training Large Language Models (LLMs) as text encoders rely on contrastive losses that treat the model as a black box function, discarding its generative and reasoning capabilities in favor of static embeddings. We introduce GRACE (Generative Representation Learning via Contrastive Policy Optimization), a novel framework that reimagines contrastive signals not as losses to be minimized, but as rewards that guide a generative policy. In GRACE, the LLM acts as a policy that produces explicit, human-interpretable rationales--structured natural language explanations of its semantic understanding. These rationales are then encoded into high-quality embeddings via mean pooling. Using policy gradient optimization, we train the model with a multi-component reward function that maximizes similarity between query positive pairs and minimizes similarity with negatives. This transforms the LLM from an opaque encoder into an interpretable agent whose reasoning process is transparent and inspectable. On MTEB benchmark, GRACE yields broad cross category gains: averaged over four backbones, the supervised setting improves overall score by 11.5% over base models, and the unsupervised variant adds 6.9%, while preserving general capabilities. This work treats contrastive objectives as rewards over rationales, unifying representation learning with generation to produce stronger embeddings and transparent rationales. The model, data and code are available at https://github.com/GasolSun36/GRACE.