PaperlayerView exact revision on arXiv

Source-linked AI summary

Distributional Reward Calibration and Budget-Neutral Credit Assignment for Multimodal RLVR

Lang Chen, Zongxin Xu, Borui Wang, Anran Liu, Shitong Chen

arXiv:2610.40250v1cs.AI

TL;DR

GRPO advantage estimators typically address either reward baselining or token-level credit assignment, although affine baselines preserve task reward-law shape and token weights may alter total verifier credit. TACE uses rank-based optimal-transport calibration for rewards and rank-calibrated, mass-preserving token weights; it performs strongest on a majority, but not all, multimodal reasoning benchmarks.

  • Problem

    GRPO variants typically address only reward baselining or token-level credit assignment, while affine baselines preserve reward-law shape and token weights need not conserve total verifier credit.

    Source: §1 · p. 1

  • Method

    TACE combines a one-dimensional optimal-transport reward baseline with within-trajectory rank-calibrated, clipped token weights renormalized to mean one.

    Source: §1 · p. 2

  • Results

    TACE obtains the strongest results on MMMU, MathVista, and MathVision, while ZeroBench and WeMath remain stronger under existing baselines.

    Source: §5.2.1 · p. 8

  • Takeaways & Limitations

    On multimodal reasoning, distributional calibration does not by itself resolve sparse verifier outcomes or weak token-level evidence.

    Source: §6 · p. 9

    The analysis is confined to verifier-only RLVR with scalar outcome rewards on multimodal reasoning.

    Source: §6 · p. 9

Abstract

from arXiv · show

We revisit advantage estimation in Group Relative Policy Optimization (GRPO) and observe that any token-level advantage admits a two-stage factorization: a \emph{sample-level baselining} step that converts scalar verifier rewards into a comparable sequence-level signal, and a \emph{token-level credit assignment} step that distributes that signal over the trajectory. Existing estimators typically address only one of these stages: affine baselines preserve the shape of each task's reward law, so heavy tails and binary cliffs propagate into the gradient, while token-level credit rules tie the size of each token's update to raw confidence scales and need not conserve the total credit issued by the verifier. We propose \textbf{TACE} (Two-stage Advantage and Credit Estimation), a two-stage estimator that addresses both stages. At stage~(S1), we replace linear standardization with a one-dimensional optimal-transport map to a standard normal, which gives a common population reward law across tasks and caps finite-sample outliers deterministically. At stage~(S2), we rank-calibrate a stop-gradient evidence contrast within each trajectory and renormalize the clipped token weights to have mean one, so credit is redistributed across tokens while its total mass and sign are conserved. We prove population-level inter-task moment equity, a deterministic tail bound, and mass-conservation, boundedness, and sign-anchoring properties that affine baselines and unnormalized token weights do not satisfy. On a multimodal reasoning suite, TACE obtains the strongest results on a majority of the benchmarks but does not dominate every setting.

1 Introduction

GRPO advantage estimation combines reward baselining with token-level credit assignment, but existing approaches typically revise only one stage. TACE addresses both with distribution-aware reward calibration and budget-neutral token credit.

1 Introduction

  • Existing GRPO variants typically refine either affine reward baselines or token-level signals, leaving the other stage unchanged.Affine baselines preserve reward-law shape, while magnitude-based token weights can amplify or attenuate the verifier’s total signal.Source: §1 · p. 1
  • TACE combines distribution-aware rank calibration at the sample level with budget-neutral credit reallocation at the token level.Its sample baseline is distribution-aware rather than affine, and its token channel redistributes credit without inflating it.Source: §1 · p. 2
  • The paper proves inter-task moment equity and a deterministic outlier cap for stage S1, plus mass-conservation, boundedness, and sign anchoring for stage S2.On a multimodal reasoning suite, TACE achieves the strongest results on a majority of benchmarks.Source: §1 · p. 2

2 Related Work

Related work addresses reward normalization, token-credit allocation, privileged-context signals, and optimal transport in language-model training. TACE positions its two-stage estimator against those lines of work and multimodal RL post-training recipes.

2.1 Reward Normalization Inside GRPO

  • Affine GRPO normalization variants cannot reshape reward laws, so heavy tails, binary cliffs, and inter-task scale gaps can remain in gradient magnitudes.TACE instead uses rank-based transport to a fixed reference distribution.Source: §2.1 · p. 2

2.2 Advantage Reallocation in Verifier-Only RL

  • Verifier-only token-scoring methods can emphasize tokens without guaranteeing the total credit assigned to a trajectory.TACE uses within-trajectory ranks and mean-one renormalization to make token learning reallocation rather than amplification.Source: §2.2 · pp. 2–3

2.3 Privileged Context as a Magnitude Signal

  • TACE confines privileged evidence to a stop-gradient magnitude signal, then rank-calibrates, clips, and renormalizes it so verifier decisions remain the source of aggregate reinforcement.A monotone change in teacher confidence cannot alter the allocation.Source: §2.3 · p. 3

2.4 Optimal Transport for Language Model Training

  • Unlike prior language-model uses of optimal transport as a distance or projection, TACE uses the one-dimensional Brenier map as the GRPO sample-level baseline.The rank-CDF composition gives distribution invariance, outlier robustness, and order preservation.Source: §2.4 · p. 3

2.5 Multimodal RL Post-Training

  • TACE is presented as orthogonal to multimodal RL recipes based on data curation, auxiliary perceptual objectives, or external captioners.Its stage-S1 distribution invariance is intended to handle perception, mathematics, and chart rewards without per-task scale tuning.Source: §2.5 · p. 3

3 Background and Diagnosis

The background frames GRPO as a verifier-reward method that broadcasts a standardized sequence-level advantage uniformly across response tokens. It also introduces privileged-context distillation, whose gradient-direction role conflicts with TACE’s direction-anchoring requirement.

3.1 Preliminaries

  • GRPO samples candidate responses, standardizes their scalar verifier rewards into sequence-level advantages, and reuses each advantage for every token.The resulting per-token credit within a trajectory is uniform.Source: §3.1 · p. 3
  • Privileged-context distillation uses a reference policy conditioned on extra information, but letting that reference determine gradient direction violates TACE’s direction-anchoring requirement.Figure 1 depicts TACE instead using non-privileged and privileged contexts to form a token magnitude weight, then multiplying it by the sequence advantage.Source: §3.1 · p. 4, Figure 1 · p. 4

4 Method

TACE separates sequence-level reward baselining from token-level credit assignment and imposes a distinct invariant on each stage. A rank-based normal baseline and rank-calibrated, mass-preserving token weights compose into the GRPO objective.

4.1 A Two-Stage View of GRPO Advantage Estimation

  • TACE separates reward baselining from token credit assignment, requiring distribution invariance for the first stage and verifier-direction anchoring for the second.Existing variants typically modify only one stage, leaving the other stage’s distinct failure mode unresolved.Source: §4 · p. 4, §4.1 · pp. 4–5, Table 1 · p. 5

4.2 Stage 1: Quantile-Normalized Sample Baseline

  • Stage S1 maps each task’s reward distribution to N(0, 1) through a one-dimensional optimal-transport rank map, using zero advantage for groups with no reward variation.Tied raw rewards receive identical credit through the average of their quantile images.Source: §4.2 · p. 5

4.3 Stage 2: Rank-Calibrated, Mass-Preserving Token Credit

  • Stage S2 starts from stop-gradient log-evidence comparing privileged and non-privileged contexts, then rank-calibrates evidence within each trajectory.This removes dependence on raw evidence scale and is invariant to strictly increasing recalibration.Source: §4.3 · p. 6

4.4 Theory: Equity, Tails, and Conservation

  • For continuous task reward laws, stage S1 gives standard-normal advantages with equal absolute moments; its finite-group outlier cap is 1.53 versus 2.65 at N=8 and 2.89 versus 15.97 at N=256 against affine standardization.The bounds hold irrespective of reward scale.Source: §4.4 · p. 6
  • Stage S2’s weights are positive and bounded, have trajectory mean one, preserve total advantage mass, and keep every token’s sign anchored to the verifier’s advantage.The composed estimator’s token evidence ranks are invariant to strictly increasing transformations.Source: §4.4 · pp. 6–7

4.5 Composed Estimator

  • The composed advantage is the monotone reward-rank baseline multiplied by positive mean-one token weights, so S2 reallocates credit without changing its mass or sign.TACE adds one group sort, one trajectory rank sort, and one privileged-context forward pass beyond GRPO.Source: §4.5 · p. 7

5 Empirical Evaluation

The evaluation compares TACE with existing methods across five multimodal reasoning benchmarks and examines its setup, component ablations, and a qualitative WeMath example. TACE leads on three benchmarks, while removing either stage degrades results and the example illustrates coordinated visual and numerical reasoning.

5.1 Setup and Baselines

  • The study trains Qwen3-VL-8B-Instruct on MMFineReason-123K and evaluates five suites, using eight rollouts per prompt and a 256-prompt effective optimizer batch.The evaluation follows the RLSD multimodal protocol for the corpus, backbone, benchmark suite, context length, rollout settings, and quoted baseline rows.Source: §5.1 · p. 7

5.2.1 Five-Benchmark Reasoning Results

  • TACE obtains the strongest results on MMMU, MathVista, and MathVision, while existing baselines remain stronger on ZeroBench and WeMath.Table 2 reports accuracy after post-training Qwen3-VL-8B-Instruct; the baseline rows quote RLSD and the TACE row is the authors’ result.Source: §5.2.1 · p. 8, Table 2 · p. 8

5.2.2 Component Ablation

  • Removing either stage degrades results on the first three benchmarks, with the largest drop on MathVision: 6.05 points without S1 and 3.82 points without S2.Removing S1 hurts more than removing S2 on every benchmark; the ablation compares affine group standardization and uniform token weights with the corresponding full-TACE stages.Source: §5.2.2 · p. 8, Table 3 · p. 9

5.2.3 Qualitative Case Study

  • In the WeMath example, the model reads OA = 3 km and AB = 4 km, converts both to meters, preserves their order, and selects option C.The example requires matching each length to its segment and applying the kilometer-to-meter conversion consistently.Source: §5.2.3 · pp. 8–9

6 Conclusion and Limitations

The conclusion presents TACE as a two-stage response to reward-law shape and indiscriminate uniform token credit, with improvements on three benchmarks but regressions on two. The appendix states that it proves the section’s three results and adds finite-sample analyses.

6 Conclusion and Limitations

  • TACE improves results on MMMU, MathVista, and MathVision but regresses on ZeroBench and WeMath; the study is confined to verifier-only RLVR with scalar outcome rewards for multimodal reasoning.The conclusion identifies multi-marginal optimal transport for vector-valued rewards as a natural extension and reports sorting plus one privileged-context pass as additional computation.Source: §6 · p. 9
  • The conclusion attributes TACE’s two-stage design to reward-law calibration and trajectory-level credit reallocation, while noting sparse verifier outcomes or limited token evidence can still constrain performance.The paper describes its added computation as sorting and one privileged-context pass.Source: §6 · p. 9

A Deferred Proofs

  • The appendix proves the three results from Section 4.4 and adds analyses of binary verifiers, empirical-rank accuracy, transport-map optimality, and a deterministic token-signal bound.It defines Φ as the standard normal CDF, Φ−1 as its quantile function, and zα = Φ−1(α).Source: A Deferred Proofs · p. 15

A.1 Proof of Proposition 1 (Distributional Equity)

The proof maps rewards through their task-specific CDF and then the standard-normal quantile function. This produces a standard-normal advantage across tasks, unlike affine standardization, which preserves higher-moment shape differences.

A.1 Proof of Proposition 1 (Distributional Equity)

  • For a continuous reward law, applying its CDF to a reward produces a uniform variable, whose standard-normal quantile is the calibrated advantage.The proof uses the probability-integral transform and the strict monotonicity of Φ.Source: §A.1 · p. 15
  • The transformed advantage has distribution N(0, 1), so every absolute moment is task-independent.This follows from the standard-normal law of the transformed variable.Source: §A.1 · p. 15
  • Affine standardization matches the first two moments but preserves higher-moment differences between the skewed and symmetric reward laws considered.The comparison is stated for mean-zero, unit-variance laws and p > 2.Source: §A.1 · p. 15

A.2 Proof of Proposition 2 (Outlier Bound)

The proof bounds empirical quantiles using ranks, independently of reward scale, yielding a deterministic cap on the transformed outlier. In the one-outlier example, the affine statistic reaches N −1, while TACE applies the stated normal-quantile cap.

A.2 Proof of Proposition 2 (Outlier Bound)

  • The mid-rank convention keeps empirical quantiles within a deterministic interval bounded away from the endpoints, independently of reward scale.The interval’s lower endpoint is given as 1/(2N), and the bound follows because the empirical CDF depends only on ranks.Source: §A.2 · p. 15
  • For one reward equal to c > 0 and all others zero, the affine standardized outlier reaches N −1, while TACE caps it at Φ−1(1 − 1/(2N)).The stated TACE cap for this configuration is independent of c.Source: §A.2 · p. 15

A.3 Proof of Theorem 1 (Conservation, Boundedness, Sign)

The proof establishes properties of TACE’s normalized token weights and shows that the two estimator stages preserve cross-task equity and within-trajectory conservation together. For binary rewards, the rank scores depend on group size and success count, while remaining bounded where affine scores can become unbounded.

A.3 Proof of Theorem 1 (Conservation, Boundedness, Sign)

  • TACE’s token weights are positive, preserve the verifier advantage’s sign, and have a two-sided bound determined only by the clipping parameter.The proof derives positivity from clipping and normalization, then bounds each normalized weight between (1 − ϵw)/(1 + ϵw) and (1 + ϵw)/(1 − ϵw).Source: §A.3 · p. 15

A.4 Composition Property

  • Because the sample-level and token-level maps act on separate indices, cross-task equity and within-trajectory conservation hold simultaneously.Source: §A.4 · p. 15

B.1 Binary and Tied Rewards

  • For binary rewards, conditional score geometry depends only on group size and success count, while affine-score magnitudes can grow without bound near the count boundaries.The rank map trades exact centering for a deterministic, scale-free tail cap.Source: §B.1 · p. 16

B.2 Finite-Sample Accuracy of the Rank Baseline

The appendix bounds finite-sample rank-map error, establishes the population transport map’s quadratic optimality, and derives a deterministic bound on the composed token-level signal. These results cover approximation accuracy and the scope of the resulting signal bound.

B.2 Finite-Sample Accuracy of the Rank Baseline

  • The empirical mid-rank quantile map has a nonasymptotic accuracy guarantee whose failure probability depends on sample size and η, not reward scale, support, or task.The bound is intentionally conservative in the Gaussian tails.Source: §B.2 · pp. 16–17

B.3 One-Dimensional Transport Optimality

  • Among maps sending a continuous reward law to a standard normal, the monotone population map minimizes expected squared transport cost and is unique almost surely.The empirical mid-rank map is its plug-in approximation.Source: §B.3 · p. 17

B.4 Deterministic Bound on the Composed Signal

  • The composed TACE signal is bounded independently of reward magnitude, reward support, trajectory length, privileged-evidence calibration, and task identity.The bound controls both the sample-level advantage and the token-level weight; the policy ratio remains governed separately by the PPO clip.Source: §B.4 · p. 18

C Hyperparameter Sensitivity

This short-horizon, late-training sweep uses 14 continuation runs to probe sensitivity, with a 262-example WeMath pilot. Seed variation limits what single-run changes can establish: short-run reward does not reliably rank external WeMath performance, and the sweep supports no clear directional recommendation for warm-up or teacher-refresh settings.

C Hyperparameter Sensitivity

  • The study comprises 14 late-training continuation runs, each branching from a shared checkpoint for 40 steps on two GPUs, plus a 262-example WeMath choice-CoT pilot.Three runs repeat the default configuration with different data seeds; the other eleven change one factor at a time, with group-size arms keeping trajectory budgets comparable.Source: C Hyperparameter Sensitivity · p. 18, Table 4 · p. 19
  • Group sizes 4 and 16 reduce final-ten reward by 1.48 and 1.29 seed standard deviations, while warm-up and teacher-refresh effects are non-monotone and support no clear directional recommendation.The sensitivity table reports final-ten reward, response-token clipping ratio, and mean wall-clock time per continuation step.Source: C Hyperparameter Sensitivity · p. 18, Table 5 · p. 19
  • Short-run reward does not reliably rank external WeMath performance: the highest-reward arms stay near the repeated-control mean, while a control seed has the best WeMath score.The authors use the sweep to motivate assessment across seeds and external benchmarks, not hyperparameter changes after observing validation outcomes.Source: C Hyperparameter Sensitivity · p. 18, Table 5 · p. 19
  • Repeated controls span 0.0586 final-ten reward units and 4.20 WeMath points; neither ϵw = 0.1 nor S1-only improves the WeMath pilot over the repeated-control mean or best seed.Their reward gains are 1.02 and 0.90 seed standard deviations, respectively; changes smaller than the control ranges should not be read as reliable improvements.Source: C Hyperparameter Sensitivity · p. 18, Table 6 · p. 19
Loading 2610.40250v1…