Hierarchical Credit Assignment for RLVR on Fused Gromov-Wasserstein Geometry

1 INTRODUCTION

• Uniform advantages and local position or entropy signals fail to capture a reasoning token’s global semantic novelty relative to the current policy.The introduction also identifies a lack of principled attribution across token, step, and trajectory scales.Source: §1 · p. 1, Figure 1 · p. 2
• HARA represents rollouts using hidden states and token positions, compares same-outcome rollouts with an FGW barycenter, and reweights token advantages by novelty.Anchor-guided linearization converts the costly FGW formulation into a Wasserstein problem solvable with Sinkhorn.Source: §1 · p. 2
• Across 3 model sizes, 5 reasoning benchmarks, and 3 group-based RLVR methods, HARA avoids premature entropy collapse, expands exploration, and outperforms existing credit-assignment approaches.The experiments also report a good balance between effectiveness and efficiency.Source: §1 · p. 2
• The paper presents an FGW-barycenter view of internal reasoning, multi-scale credit assignment, approximation and complexity analysis, and performance gains across diverse settings.The authors describe their barycenter representation as a first, to their best knowledge.Source: §1 · p. 2

2.1 REINFORCEMENT LEARNING WITH VERIFIABLE REWARDS (RLVR)

• RLVR models an autoregressive response to a prompt and ground-truth answer, with a binary reward for whether the response matches the answer.The response probability is expressed as a product of its token probabilities.Source: §2.1 · p. 3

2.2 GROUP-BASED POLICY OPTIMIZATION

• Group-based policy optimization estimates relative advantages from multiple sampled rollouts grouped by verifiable rewards, without an explicit critic.The GRPO objective also uses an importance ratio and KL regularization against a reference policy.Source: §2.2 · p. 3

2.3 FUSED GROMOV-WASSERSTEIN BARYCENTER

• FGW distance combines feature and structural discrepancies between attributed graphs, while an FGW barycenter summarizes shared node-feature and graph-structure patterns.The distance weights Wasserstein and Gromov-Wasserstein terms using α.Source: §2.3 · p. 3

3.1 DISTRIBUTIONAL MODELING OF INTERNAL REASONING PATTERN

• Rollouts become attributed chain graphs over token hidden states and positions, and outcome-specific FGW barycenters summarize the policy’s internal reasoning patterns.The graph keeps a chain structure while representing token hidden states as node features.Source: §3.1.1 · p. 4, §3.1 · p. 4

3.1.1 FGW BARYCENTER AS INTERNAL REASONING PATTERN

• The barycenter-support mapping is sparse and block-diagonal, while FGW achieves higher ASL than W and GW and ASL rises monotonically with K.The mapping indicates deterministic correspondences to meaningful reasoning stages.Source: §3.1.1 · pp. 4–5, Figure 2 · p. 5

3.1.2 ANCHOR-GUIDED LINEARIZATION

• Fixing reasoning-state anchors linearizes the FGW comparison into a Wasserstein distance that can be solved efficiently with entropy regularization and Sinkhorn.The first and last rollout tokens serve as shared initial and final state anchors within a prompt and outcome.Source: §3.1.2 · p. 5

3.2 HIERARCHICAL CREDIT ASSIGNMENT

• HARA uses span contributions to linearized FGW distance to reweight group advantages at sequence and token levels, encouraging correct novel behavior while penalizing incorrect novel behavior less.The sequence- and token-level variants are named HARA-seq and HARA-tok.Source: §3.2 · pp. 5–6, Figure 3 · p. 6

3.3 THEORETICAL ANALYSIS

• The approximation-error bound depends on order-reversing pairs in the optimal transport plan; FGW’s alignment of similar normalized positions is expected to keep such pairs few.This yields an expectation of small error when the transport plan is sparse and block-diagonal.Source: §3.3 · pp. 6–7
• HARA has time complexity O(BKLLb(d + S)) and space complexity O(KLLb + KLd + Lbd), scaling linearly with rollout length L.With barycenter size Lb = 32, the paper describes the method as efficient and scalable across reasoning tasks.Source: §3.3 · p. 7

4.1 EXPERIMENTAL SETUP

• The evaluation spans five math benchmarks, three Qwen3-Base model sizes, three credit-assignment baselines, and GRPO, Dr.GRPO, and DAPO.GRPO is the default baseline; metrics are average accuracy for MATH500 and Minerva Math, and Pass@16 for AMC23, AIME24, and AIME25.Source: §4.1 · p. 7

4.2 MAIN RESULTS

• HARA consistently leads existing credit-assignment methods across datasets and LLMs, improving GRPO by up to 4.8% in average performance and baselines by up to 4.0%.HARA-tok also outperforms HARA-seq on most datasets and LLMs.Source: §4.2 · p. 7, Table 1 · p. 8
• On MATH500 with Qwen3-8B-Base, HARA converges faster, achieves better final performance, maintains higher entropy, and generates longer responses than GRPO.The figure tracks accuracy, response length, and entropy during training.Source: §4.2 · p. 8, Figure 4 · p. 8

4.3 EFFICIENCY RESULTS

• HARA adds up to 13.6% in total training time versus GRPO, yet reaches the same accuracy up to 2.64× faster and gains 8.7% accuracy in the same time.The reported accuracy-versus-time comparison is on MATH.Source: §4.3 · p. 8, Figure 5 · p. 9, Figure 6 · p. 9

4.4 FURTHER STUDIES

• Among the tested novelty metrics, FGW consistently performs best, while GW underperforms even baseline GRPO despite aligning tokens in Figure 2.The comparison includes Wasserstein distance and embedding distance from the barycenter to each rollout’s last hidden state.Source: §4.4 · p. 9, Table 2 · p. 9
• Additional studies examine the FGW trade-off parameter α and ablate the anchor-guided linearization in Appendix C.2.Source: §4.4 · p. 9

• Prior FGW-barycenter applications include attributed-graph interpolation, multivariate time-series transitions, cortical-surface alignment, and multi-network alignment.The cited works use input-dependent or dynamically weighted barycenters, fused unbalanced transport, or FGW barycenters.Source: §5 · p. 9
• RLVR credit-assignment methods include sequence-level error-diversity shaping, span-level Wasserstein distance, and token-level proximity or entropy weighting.The passage cites Liu et al. (2026), Chen et al. (2026), Parthasarathi et al. (2025), Cheng et al. (2026), and Xie et al. (2026).Source: §5 · p. 9
• The surveyed methods do not identify semantically novel reasoning behaviors at the finest token level.Source: §5 · p. 9

• HARA represents each rollout as a distribution over token hidden states and locations, then uses same-outcome FGW barycenters to capture latent reasoning patterns.Source: §6 · p. 9
• It measures a reasoning element’s semantic novelty by its contribution to the FGW distance from the barycenter and reweights token-level advantages accordingly.Source: §6 · p. 9
• Experiments show HARA improves diverse group-based RLVR methods while maintaining a balance between effectiveness and efficiency.Source: §6 · p. 9

AI USE STATEMENT

• The authors report using generative AI only to polish grammar and readability, and state that they reviewed the assisted work and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

A.1 PROOF OF LEMMA 3.1

• The proof fixes the second transport plan to pre-align reasoning anchors, addressing the original FGW formulation’s non-convex, quadratic GW term.Source: §A.1 · p. 12
• The resulting linearized formulation is a Wasserstein distance that is linear with respect to the transport plan.Source: §A.1 · p. 12

A.2 PROOF OF THEOREM 3.2

• Theorem 3.2 states that the approximation error of anchor-guided linearization is bounded for an optimal plan with uniform marginal distributions.Source: §A.2 · p. 13
• The proof evaluates the original FGW objective and its linearization at a common transport plan and rewrites GW terms using paired normalized token locations.Source: §A.2 · p. 13

A.3 PROOF OF THEOREM 3.3

• HARA’s time complexity is O(BKLLb(d + S)) and its space complexity is O(KLLb + KLd + Lbd), with K rollouts, length L, barycenter size Lb, hidden-state dimension d, and iteration counts S and B.Source: §A.3 · p. 15

B.1 DATASET SETUP

• Training uses 12,000 MATH problems after excluding problems in MATH500 from the 12,500-problem dataset.Source: §B.1 · p. 15
• Evaluation covers MATH500, AMC23, AIME24, AIME25, and Minerva Math.Source: §B.1 · p. 15

B.2 TRAINING SETUP

• Training uses 16 sampled rollouts per prompt and removes the KL penalty by setting β = 0 across group-based RLVR methods.Evaluation metrics are averaged over five runs; experiments use eight NVIDIA A100 GPUs.Source: §B.2 · p. 15

B.3 HYPERPARAMETER SETUP

• HARA uses α = 0.25, barycenter size 32, and entropy regularization coefficient 0.05 consistently across experiments.Source: §B.3 · p. 16

B.4 SETUP FOR INTERPRETABILITY STUDY IN FIGURE 2

• The interpretability study uses Qwen3-8B-Base and 16 rollouts per outcome-grounded group to compute barycenters.It sets barycenter size to 32, entropy regularization to 0.05, and α to 0.25.Source: §B.4 · p. 16
• The annotation protocol builds a compact shared superset of reasoning stages, consolidating equivalent functions while allowing rollouts to skip, repeat, or reorder stages.Each rollout is segmented into contiguous spans, every token is assigned to one stage, and genuine differences in reasoning order are preserved.Source: §B.4 · pp. 16–17

• HARA consistently improves Dr.GRPO and DAPO across all datasets, and predominantly outperforms existing credit-assignment methods on most datasets and LLMs.Source: §C.1 · p. 18

• Up to 4× speed-up per gradient update step comes with less than 1% performance variation when using anchor-guided linearization.The authors attribute the efficiency gain to strict convexity of the linearized problem.Source: Table 4 · p. 18, §C.2 · p. 19
• Less than 2% performance variation occurs for α in [0.25, 0.75]; α = 0 or 1 causes more severe degradation by omitting semantic or structural information.HARA uses α = 0.25 by default.Source: Figure 7 · p. 18, §C.2 · pp. 18–19