PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2510.04506v1cs.CLcs.AIcs.IR

TL;DR

Existing contrastive training can treat LLMs as opaque encoders, leaving their reasoning inaccessible. GRACE turns contrastive signals into rewards for a rationale-generating policy, and reports broad MTEB gains while preserving general capabilities.

  • Problem

    Prevailing contrastive methods train LLMs as opaque encoders, producing static embeddings without exposing the reasoning behind similarity judgments.

    Source: §1 · p. 2

  • Method

    GRACE trains an LLM policy to generate explicit rationales, uses them to form representations, and treats contrastive signals as rewards for policy-gradient optimization.

    Source: §1 · p. 2, §3.1 · p. 4, §3.2 · p. 4

  • Results

    GRACE reports broad MTEB gains across supervised and unsupervised settings, including an 11.52% increase in MTEB average over the Base model in the supervised setting.

    Source: §4.3 · p. 8

  • Takeaways & Limitations

    GRACE makes the semantics behind its representations inspectable through rationales while delivering cross-category MTEB gains and preserving general capabilities.

    Source: §7 · p. 15

    The optimization is on-policy, so the method omits importance sampling.

    Source: §3.4 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Conventional contrastive training treats LLMs as opaque encoders, limiting access to their reasoning and generative capabilities. GRACE instead uses contrastive signals to guide a rationale-generating policy, producing interpretable explanations and representations.

1 INTRODUCTION

  • Conventional contrastive approaches optimize static embeddings in an opaque latent space, suppressing LLMs’ structured reasoning and natural-language generation.The resulting similarity judgments do not reveal why texts were judged similar or which semantic features were prioritized.Source: §1 · p. 2
  • GRACE reframes contrastive objectives as rewards that guide an LLM policy to generate interpretable understandings of input texts.These understandings serve as human-readable explanations and are encoded as representations.Source: §1 · p. 2
  • The model generates explicit rationales, derives embeddings from them by mean pooling, and uses policy-gradient optimization to increase query–positive similarity and decrease query–negative similarity.The reward-based optimization is intended to teach the model to generate faithful rationales while learning effective representations.Source: §1 · pp. 2–3
  • GRACE improves the MTEB average by 11.5% over baseline models.The paper also reports that the method preserves general-purpose capabilities without performance degradation.Source: §1 · p. 3

2 PRELIMINARIES

The preliminaries describe contrastive embedding learning, policy-gradient optimization, and GRPO. Together, these sections introduce the representation objective and the policy-optimization framework used to train generative models.

2.1 CONTRASTIVE LEARNING FOR TEXT EMBEDDINGS

  • Contrastive embedding learning trains an encoder to minimize InfoNCE loss, commonly using in-batch negatives, with effectiveness depending on hard-negative mining.Hard negatives are challenging samples close to the decision boundary.Source: §2.1 · p. 3

2.2 POLICY GRADIENT OPTIMIZATION

  • Policy-gradient optimization maximizes expected reward over responses sampled from a policy, with a baseline used to reduce gradient-estimate variance.The cited policy-gradient expression is described as an unbiased estimator of the true gradient.Source: §2.2 · p. 3

2.3 GROUP RELATIVE POLICY OPTIMIZATION

  • GRPO samples a group of responses for each input and computes each response’s advantage relative to that group.Its objective uses an importance ratio and is presented without KL divergence.Source: §2.3 · p. 4

3 GENERATIVE REPRESENTATION LEARNING VIA CONTRASTIVE POLICY OPTIMIZATION

GRACE treats contrastive signals as rewards guiding a generative policy that produces interpretable rationales, rather than as losses applied to a passive encoder. It converts those rationales into representations and optimizes them with contrastive rewards, including in an unsupervised setting.

3 GENERATIVE REPRESENTATION LEARNING VIA CONTRASTIVE POLICY OPTIMIZATION

  • GRACE reframes contrastive signals as rewards guiding a generative policy that turns the LLM into an active agent producing interpretable rationales.The framework replaces the passive encoder view with a policy-guided approach to representation learning.Source: §3 · p. 4

3.1 RATIONALE-GENERATING POLICY

  • For each input, the policy generates an explicit reasoning trace that identifies salient semantic features, key concepts, and relations grounded in domain knowledge.The rationale is intended to provide a transparent trace supporting downstream similarity assessment.Source: §3.1 · p. 4

3.2 FROM RATIONALE TO REPRESENTATION

  • The model conditions hidden states on the instruction-augmented input and rationale, then uses masked mean pooling to form a representation while excluding system-prompt tokens.The resulting representation is described as semantically rich and interpretable.Source: §3.2 · p. 4

3.3 CONTRASTIVE REWARDS AS POLICY GUIDANCE

  • GRACE uses a policy-based optimization approach, illustrated with GRPO, while remaining agnostic to the specific policy-gradient algorithm.The approach is motivated partly by practical difficulties and inefficiencies in training value models.Source: §3.3 · p. 5

3.3.1 ROLLOUT STRATEGY

  • Positive documents receive K stochastic rollouts, while queries and negative documents receive single-sample generations that serve as fixed reward-computation anchors.Similarity-based rewards support advantage estimation in GRPO and policy updates intended to improve text-understanding quality and diversity.Source: §3.3.1 · p. 5

3.3.2 REWARD DESIGN

  • The composite reward combines query–relevant alignment, within-document rollout consistency, and hard-negative discrimination, with hyperparameters balancing consistency and hard-negative contributions.Hard negatives are the most challenging distractors, selected by maximum similarity across another instance’s rollouts; temperature scaling sharpens the reward distribution.Source: §3.3.2 · pp. 5–6

3.4 POLICY OPTIMIZATION OBJECTIVE

  • GRPO advantages are computed relative to a group baseline without standard-deviation normalization, and the on-policy objective maximizes advantage-weighted likelihood without importance sampling.The text identifies the optimization as on-policy.Source: §3.4 · p. 6

3.5 UNSUPERVISED LEARNING EXTENSION

  • For raw text without query–document pairs, the unsupervised extension aligns each rollout with its text’s anchor interpretation and retains within-text consistency and in-batch hard-negative terms.When enabled, its overall objective uses the same weighted combination as the supervised reward design.Source: §3.5 · p. 6

4 EXPERIMENT

GRACE improves supervised and unsupervised MTEB averages across four backbones, while the experiments examine benchmark coverage, reward design, optimizer portability, and general-domain capabilities. Results also report near-zero average shifts on general tasks for GRACE, unlike the deterioration under contrastive fine-tuning.

4.1 DATASETS AND EVALUATION METRICS

  • MTEB evaluation covers seven task categories and 56 datasets, with mean pooling that excludes instruction tokens.The benchmark includes retrieval, reranking, clustering, pair classification, classification, semantic textual similarity, and summarization.Source: §4.1 · p. 6

4.3 MAIN RESULTS

  • 11.52% over Base in supervised MTEB and 6.85% over Base in unsupervised MTEB are the average improvements across four backbones.The supervised gains span seven task families; unsupervised results improve stepwise from Base through SimCSE to GRACE across all four backbones.Source: Table 1 · p. 7, §4.3 · p. 8, Table 2 · p. 8

4.2 BASELINES

  • Supervised comparisons include four training settings, while unsupervised comparisons include open models, Base, SimCSE, and GRACE.Experiments use four decoder-only backbones; supervised training uses 1.5M E5 samples, while unsupervised training follows SimCSE without additional unlabeled data.Source: §4.2 · p. 7

4.4 ABLATION STUDIES

  • A representative 16-task subset supports compute-efficient ablations and analysis while covering seven MTEB task families.The subset is selected to balance domain diversity, supervision formats, and task difficulty.Source: §4.4 · p. 8, §4.5 · p. 9

4.5 SUBSET OF MTEB BENCHMARK

  • The ablation subset contains 16 datasets spanning retrieval, reranking, clustering, pair classification, classification, STS, and summarization.The family counts are 3, 2, 3, 1, 3, 3, and 1, respectively.Source: Table 3 · p. 9, §4.5 · p. 9

4.5.1 ABLATION ANALYSIS OF REWARD FUNCTION DESIGN

  • With both reward weights zero performance is poor; hard-negative-only settings outperform consistency-only settings, and performance is more sensitive to λ2.The 5 × 5 grid varies consistency and hard-negative-mining weights across 0.0, 0.3, 0.5, 0.7, and 1.0.Source: §4.5.1 · p. 9, Figure 3 · p. 10

4.5.2 COMPARISON WITH ALTERNATIVE RL ALGORITHMS

  • GRPO is most effective in the reported setting, while REINFORCE++, ReMax, and DAPO yield smaller gains; the framework improves across RL optimizers.The comparison supports portability across different optimization schemes.Source: §4.5.2 · pp. 9–10, Table 4 · p. 10

4.5.3 GENERALIZATION TO GENERAL DOMAIN TASKS

  • Across four backbones, GRACE yields near-zero average shifts on general-domain tasks, while contrastive-learning fine-tuning deteriorates sharply.The evaluation spans mathematics, knowledge and reasoning, and code generation.Source: §4.5.3 · p. 10, Table 5 · p. 11

5 ANALYSIS

Across the analyses, GRACE improves embedding quality while largely preserving generative competence, and its training and representation choices shape both performance and output structure. The reported results also characterize efficiency costs, hyperparameter effects, and changes in model responses during training.

5.1 EFFICIENCY ANALYSIS

  • At matched decoding budgets, GRACE delivers better accuracy than Base without extra latency, while generative methods cost more latency than encoder-style methods.Generation dominates generative-pipeline latency, with diminishing returns beyond a practical decoding-budget knee around G≈256.Source: §5.1 · p. 11

5.2 JOINT ANALYSIS OF GENERATIVE AND EMBEDDING PERFORMANCE

  • GRACE-1.5B and GRACE-3B improve their base models’ embedding strength by more than 15% on average without sacrificing generative ability.GRACE-4B also substantially enhances Qwen-3-4B, placing it in a previously unattained balanced-performance regime.Source: §5.2 · p. 12

5.3 TRAINING PROGRESSION ANALYSIS

  • As training steps increase, subtask accuracy and response length rise, with later responses showing richer information density and more explicit reasoning chains.The progression figure reports steadily improving subtask accuracy alongside increasing response length.Source: §5.3 · p. 12, Figure 5 · p. 12

5.4 EFFECTS OF VARIOUS REPRESENTATION APPROACHES

  • Mean pooling from the last or penultimate layer consistently outperforms EOS-token and max pooling across model variants in supervised and unsupervised settings.Last-layer and penultimate-layer mean pooling perform similarly, with slightly higher performance in unsupervised models.Source: §5.4 · p. 13, Figure 6 · p. 13

5.4.1 HYPERPARAMETER SENSITIVITY ANALYSIS

  • Performance improves monotonically with larger batch size and more rollouts; bs = 64 and n = 8 give the best observed performance.The authors identify bs≥32 and n≥4 as a practical Pareto choice under typical compute budgets.Source: Figure 7 · p. 13, §5.4.1 · p. 14

5.5.1 OUTPUT EVOLUTION DURING TRAINING

  • Outputs evolve from brief topic recognition to structured lists, then broader contextual coverage and a coherent narrative by step 1000.The progression is described as increasing output length, contextualization, and information density.Source: §5.5.1 · p. 14, Table 6 · p. 14

5.5.2 IN-DEPTH ANALYSIS OF MODEL RESPONSE PATTERNS

  • For long documents, responses form compact outlines; for minimal factual queries, they provide precise answers with concise, high-signal context.The paper presents these as representative response patterns across input types.Source: Table 7 · p. 14, §5.5.2 · p. 15

6 RELATED WORK

Related work situates GRACE among LLM embedding approaches and reinforcement-learning methods for reasoning. The paper builds on both lines by treating contrastive learning as a reward that guides a generative policy.

6.1 LLM AS EMBEDDING

  • LLM embedding research includes tuning-free studies of instructions and tuning-based methods that adapt models, with recent state-of-the-art approaches increasingly relying on tuning.BGE and GTE strengthen semantic representations through contrastive learning on large-scale text pairs.Source: §6.1 · p. 15

6.2 REINFORCEMENT LEARNING IN REASONING

  • Reinforcement-learning approaches such as PPO and GRPO optimize reasoning policies using rewards, and GRPO removes the critic model and uses group-wise reward normalization.The paper draws inspiration from this policy-optimization line to recast contrastive learning as a reward for representation learning.Source: §6.2 · p. 15

7 CONCLUSION

The conclusion presents GRACE as a policy-gradient framework that uses contrastive rewards to shape interpretable rationales and their resulting embeddings. It reports cross-category MTEB gains in supervised and unsupervised settings while preserving general capabilities on non-embedding tasks.

7 CONCLUSION

  • GRACE reframes contrastive signals as rewards for a generative policy, using policy-gradient updates to shape the reasoning behind embeddings.The rationales make representation semantics inspectable.Source: §7 · p. 15
  • GRACE delivers consistent cross-category MTEB gains across multiple backbones in both supervised and unsupervised settings.The conclusion characterizes the gains as empirical results across multiple backbones.Source: §7 · p. 15
  • The conclusion reports that GRACE preserves general capabilities on non-embedding tasks while improving representations.Source: §7 · p. 15

A.1 TRAINING ALGORITHM

Algorithm 1 describes GRACE as a multi-phase training procedure that generates rationales, pools representations, computes contrastive rewards, and optimizes the policy.

A.1 TRAINING ALGORITHM

  • The procedure starts with an initial policy and hyperparameters, then generates query, negative, and sampled positive rationales for each batch.Source: §A.1 · p. 21
  • It encodes the rationale-conditioned inputs and applies masked mean pooling to obtain semantic representations.Source: §A.1 · p. 21
  • It computes multidimensional contrastive rewards, including hard-negative identification, and then optimizes the policy via GRPO.Source: §A.1 · p. 21

A.2 THEORETICAL PERSPECTIVE ANALYSIS

The paper connects contrastive learning and reinforcement learning through their shared reliance on learning from feedback rather than absolute ground-truth labels. It presents this shared basis as enabling a natural integration of their objectives.

A.2 THEORETICAL PERSPECTIVE ANALYSIS

  • Contrastive learning and reinforcement learning are connected by their shared ability to learn from feedback rather than absolute ground-truth labels.Source: §A.2 · p. 21
  • This shared reliance on feedback provides a principled connection between the two learning frameworks.Source: §A.2 · p. 21
  • The shared basis enables a natural integration of contrastive and reinforcement learning objectives.Source: §A.2 · p. 21

A.2.1 UNIFIED LEARNING WITHOUT GROUND TRUTH

The appendix frames contrastive preferences as reward signals for policy-gradient learning without explicit labels, relates the approach to generative InfoNCE, and describes regularization and implementation choices.

A.2.1 UNIFIED LEARNING WITHOUT GROUND TRUTH

  • Contrastive learning distinguishes positive from negative samples, allowing their relative preferences to serve as rewards for policy-gradient optimization without explicit labels.Source: §A.2.1 · pp. 21–22

A.2.2 CONNECTION TO INFONCE

  • The framework is understood as implicitly optimizing a generative variant of the InfoNCE objective.Source: §A.2.2 · p. 22

A.2.3 CONVERGENCE ANALYSIS

  • The convergence analysis attributes variance reduction to contrastive learning and policy gradients, with a consistency reward regularizing and bounding policy updates.Source: §A.2.3 · p. 22

A.2.4 ADDITIONAL TRAINING DETAILS

  • Length regularization is used to discourage excessively long responses without meaningful content, with Lmax setting the maximum sequence length and γ the over-length penalty coefficient.Source: §A.2.4 · p. 22

A.2.5 ADDITIONAL IMPLEMENTATION DETAILS

  • Training uses AdamW at 1×10−6, eight GRPO rollouts per positive document, and excludes instruction tokens from pooling during representation extraction.Source: §A.2.5 · p. 23

Related papers

Loading 2510.04506v1…