PaperlayerView exact revision on arXiv

Source-linked AI summary

Flat-RLVR: Sharpness-Aware Reinforcement Learning with Verifiable Rewards for Reasoning Capacity in LLMs

Yeonbum Kim, Hyun Kyu Lee, Sung Whan Yoon

arXiv:2610.29281v1cs.LGcs.AI

TL;DR

RLVR can sharpen policy distributions, suppressing rare-yet-correct trajectories and reducing sample diversity. Flat-RLVR integrates sharpness-aware minimization into group-relative RLVR; the authors report improved diversity and mathematical reasoning capacity while preserving accuracy.

  • Problem

    RLVR can concentrate probability on already-likely responses while suppressing rare-yet-correct trajectories, reducing diversity relevant to multi-sample inference-time computation.

    Source: §1 · pp. 1–2

  • Method

    Flat-RLVR integrates sharpness-aware minimization into group-relative RLVR to seek flatter reward regions and smooth policy distributions.

    Source: §1 · p. 2, §5.1 · p. 5

  • Results

    Flat-RLVR variants consistently improve output diversity and Pass@256 mathematical reasoning capacity while preserving accuracy in the reported LLM post-training adaptations.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    Flatness-seeking optimization offers a way to mitigate RLVR distribution sharpening while maintaining accuracy in the reported adaptations.

    Source: §7 · p. 9

Abstract

from arXiv · show

Reinforcement learning with verifiable rewards (RLVR), particularly through group-relative policy optimization (GRPO), has emerged as a central paradigm for the post-training of large language models (LLMs). While RLVR is highly effective at enhancing the reasoning capabilities of LLMs, it commonly induces \textit{distribution sharpening}, whereby the policy distribution becomes excessively peaked and rare-yet-correct trajectories are progressively suppressed, thus diminishing sample diversity. To this end, we promote flat-reward optima in RLVR to smooth the policy distribution. We theoretically link flat reward RLVR to reduced distribution sharpening, and empirically show improved sample diversity while maintaining single-attempt accuracy. Our method, called \texttt{Flat}-RLVR, searches for flat reward regions via sharpness-aware minimization during LLM post-training, combining this with GRPO variants to achieve better trade-offs between sampling efficiency and reasoning capacity.

1. Introduction

RLVR has become a prominent approach to LLM post-training, but it can sharpen policy distributions and suppress rare correct reasoning paths. Flat-RLVR uses sharpness-aware minimization to seek flatter reward regions, with the stated goal of improving diversity while preserving accuracy.

1. Introduction

  • Prior observations link RLVR to suppressed rare-yet-correct samples, reduced diversity, and concentration on problems already solvable by the base model (He et al., 2025).These findings are presented as evidence of distribution sharpening: reinforcing easier, already-solved instances while losing exploratory capacity.Source: §1 · pp. 1–2
  • Reduced sample diversity can undermine inference-time compute strategies that generate multiple candidates and use a verifier to select correct solutions.The passage highlights mathematics and code, where deterministic verifiers enable reliable selection.Source: §1 · p. 2
  • Flat-RLVR steers RLVR optimization toward flatter reward regions rather than steepest loss descent, aiming to smooth policy distributions and diversify generated candidates.The introduction motivates this approach as a response to distribution sharpening.Source: §1 · p. 2
  • Flat-RLVR variants consistently improve Pass@256 mathematical reasoning capacity while preserving accuracy in the authors’ LLM post-training adaptations.The framework combines sharpness-aware minimization with GRPO, DAPO, and CISPO variants.Source: §1 · p. 2

2. Related Work

The related work positions Flat-RLVR as a sharpness-aware optimization approach for group-relative RLVR. It distinguishes this intervention from reward- or token-level mitigations and draws on SAM variants and prior work connecting flatness with robustness.

2.1. Group-Relative RLVR

  • Group-relative RL methods normalize rewards across rollouts, avoiding a separate value network; this work applies sharpness-aware optimization across that method family.The passage names GRPO, DAPO, and CISPO as examples of group-relative RLVR methods.Source: §2.1 · p. 2

2.2. Distribution Sharpening in RLVR

  • Existing sharpening mitigations intervene on rewards or token distributions, whereas this work biases updates toward flat loss-landscape regions via SAM, an orthogonal approach.Examples include F-GRPO’s focal weighting, unlikeliness rewards for rare correct trajectories, and clipping-based variants such as DAPO.Source: §2.2 · p. 2

2.3. Advances in SAM

  • Flat-RLVR draws on SAM variants addressing parameter-rescaling sensitivity, noisy gradients, computational overhead, and whether optimization reaches flat minima.The work adopts periodic SAM updates inspired by LookSAM and Gradient Centralization inspired by GCSAM, while stating it is not limited to one SAM approach.Source: §2.3 · pp. 2–3

2.4. Flat-Minima and Generalization

  • Prior work links flat minima with generalization and robustness, including in RL, but how flatness-aware RL steers LLM behavior remains nontrivial.The authors describe their work as the first attempt to employ SAM in LLM post-training to relieve RLVR distribution sharpening.Source: §2.4 · p. 3

3. Preliminaries

The preliminaries describe GRPO’s group-relative clipped objective and SAM’s search for solutions with both low loss and low curvature. Together, these provide the optimization components used by Flat-RLVR.

3.1. Group Relative Policy Optimization

  • GRPO computes rollout advantages relative to a sampled group and optimizes a clipped surrogate objective; in RLVR, all tokens in a rollout share the same advantage.The presentation drops the KL regularization term, following cited findings that it hinders exploration in reasoning.Source: §3.1 · p. 3

3.2. Sharpness-Aware Minimization

  • SAM minimizes worst-case loss within a parameter-perturbation neighborhood, penalizing sharp minima and encouraging convergence to flat ones.The perturbation lies within an ℓ2 ball of radius ρ, and its inner maximization is approximated to obtain the worst-case perturbation.Source: §3.2 · p. 3

4. Analysis on Flatness in RLVR

The analysis studies how sharpness-aware RLVR changes output probabilities using a local sequence-level unclipped surrogate. It shows that flatness corrections can reduce the growth of probability gaps among correct outputs, subject to a cumulative-condition assumption.

4. Analysis on Flatness in RLVR

  • The analysis defines Flat-RLVR through a sharpness-aware loss and applies a two-step SAM update: perturb toward the worst-case direction, then descend at the perturbed point.The local loss is the negative of the sequence-level RLVR objective, so maximizing that objective corresponds to minimizing the loss.Source: §4 · pp. 3–4
  • Flat-RLVR's update combines the standard reward-improving score direction with a curvature-based flatness correction that dampens over-amplification.Under a sequence-level NTK approximation, the residual is negligible at leading order.Source: §4 · p. 4
  • For a higher-probability and a lower-probability correct path, a positive correction reduces the GRPO-induced growth of their log-probability gap.The comparison is between two verifier-correct outputs, not between correct and incorrect outputs.Source: §4 · pp. 4–5
  • When the cumulative flatness correction dominates the residual, Flat-RLVR's correct-set sharpening is no larger than that induced by the corresponding GRPO score updates.This is the condition and comparison stated in Proposition 4.5.Source: §4 · p. 5
  • The sharpening result does not imply higher probability for incorrect outputs; negative-advantage incorrect outputs can still be suppressed.The correction can reduce over-concentration among correct paths while preserving reward separation between correct and incorrect outputs.Source: §4 · p. 5

5. Proposed Method: Flat-RLVR

Flat-RLVR integrates sharpness-aware minimization into group-relative RLVR, using SAM-based updates to seek flatter regions. Its LLM post-training adaptations schedule SAM periodically, adjust perturbation radius by prompt difficulty, and apply gradient centralization to both SAM steps.

5.1. Flat-RLVR

  • Flat-RLVR integrates sharpness-aware minimization into group-relative RLVR, building on standard GRPO group-relative advantages.For each prompt, the policy samples completions and normalizes their rewards within the group to estimate advantages.Source: §5.1 · p. 5

5.2. Adaptations for LLM post-training

  • SAM is applied every n steps, making n both a scalability knob and a control on effective regularization that impacts the Pass@K trade-off.The periodic approach limits the extra gradient computation required to obtain a perturbation.Source: §5.2 · pp. 5–6
  • Figure 3 traces BASE, GRPO, and Flat-GRPO Pass@K curves from K = 1 to K = 256 on in-domain mathematical reasoning benchmarks.The caption identifies Qwen2.5-Math-1.5B for panels (a), (c), and (d), and Llama-3.2-1B-Instruct for panel (b).Source: Figure 3 · p. 6
  • The perturbation radius ρ is adjusted according to prompt difficulty to improve convergence stability.The method adopts a prompt-difficulty-aware strategy with Focal Weight.Source: §5.2 · p. 6
  • Gradient centralization is applied to both SAM ascent and descent steps; the paper reports improved exploration during post-training.The adaptation reduces gradient norm throughout optimization.Source: §5.2 · p. 6

6. Experiments

The experiments test whether Flat-RLVR preserves reasoning coverage and diversity while retaining single-attempt accuracy, across mathematical and code tasks, models, and RLVR variants. They also examine training-time evidence of reduced sharpening and the method’s computational costs.

6.2. Reasoning Capacity Evaluation

  • Flat-GRPO maintains an advantage over GRPO and the base model across Pass@K, while retaining RLVR’s single-attempt gains.As K grows, GRPO’s advantage erodes and the base model can match or surpass it; Flat-GRPO preserves exploratory capacity.Source: §6.2 · pp. 6–7

6.3. Extended Evaluations

  • Across GRPO, DAPO, and CISPO, flatness-seeking improves average Pass@256 while preserving Pass@1 on both evaluated backbones.Pass@256 gains are +2.5, +3.4, and +3.0 on Qwen2.5-Math-1.5B, and +2.0, +1.2, and +1.3 on Llama3.2-1B-Instruct, respectively.Source: Table 2 · p. 7, §6.3 · p. 7
  • Flat-GRPO exceeds F-GRPO in average Pass@256 by +2.1 on Qwen2.5-Math-1.5B and +0.1 on Llama-3.2-1B-Instruct while remaining competitive on Pass@1.The approaches operate through loss geometry and advantage magnitude, respectively.Source: Table 3 · p. 7, §6.3 · pp. 7–8

6.2. Reasoning Capacity Evaluation

  • GRPO forgets 6.3% of base-solvable problems and newly solves 5.4%, while Flat-GRPO forgets 2.8% and newly solves 6.3%.The analysis categorizes problems according to whether the base and RL-tuned models solve them.Source: Table 1 · p. 6, §6.2 · p. 7

6.3. Extended Evaluations

  • On Qwen2.5-3B-Instruct, Flat-GRPO improves Pass@256 by +2.0 over GRPO at a Pass@1 cost of -0.3.On MBPP code generation, it achieves a comparable sampling-efficiency gain to GRPO (-2.2) while better preserving reasoning capacity (+4.0).Source: Table 4 · p. 8, §6.3 · p. 8, Table 5 · p. 8

6.4. Empirical Validation of Theoretical Analysis

  • Flat-GRPO reduces the correct-set log-probability gap during training and retains more probability mass on rare-correct trajectories than GRPO.The gap tracks the difference between high- and low-probability correct outputs; Figure 4 plots the gap and lowest-probability correct-response trajectory.Source: §6.4 · p. 8, Figure 4 · p. 8
  • At perturbation scale σ = 0.001, Table 6 reports that Flat-GRPO converges to a flatter region than GRPO, as assessed using the low-pass-filter (LPF) flatness measure.The comparison uses a fixed evaluation batch of 32 MATH prompts and identical batch and perturbation seeds for both models.Source: §6.4 · pp. 8–9, Table 6 · p. 9

6.5. Analysis of Cost-Benefit Trade-off

  • With periodic SAM, training time stays within 0.1 hours of GRPO and per-GPU memory is comparable or lower; more frequent SAM boosts larger-K Pass@K but lowers Pass@1.On Qwen2.5-3B-Instruct, a SAM step adds 0.39s (+0.43%) over an ordinary step.Source: §6.5 · p. 9, Table 7 · p. 9, Figure 5 · p. 9

7. Conclusion

The conclusion presents flatness-seeking optimization as a way to mitigate distribution sharpening in group-relative RLVR.

7. Conclusion

  • Flat-RLVR integrates sharpness-aware minimization into RLVR; experiments with Flat-GRPO, Flat-DAPO, and Flat-CISPO report improved output diversity while preserving accuracy.The paper’s mathematical analysis finds that flat-minima search dampens abrupt probability increases.Source: §7 · p. 9

A.1. Proof of Lemma 4.3

The appendix derives the Flat-RLVR update by approximating the worst-case perturbation and perturbed gradient, then decomposes the change in the correct-support gap relative to GRPO. Summing the stepwise relation under the proposition’s dominance condition gives the claimed inequality.

A.1. Proof of Lemma 4.3

  • Under the one-step assumption that πθ ≈ πθold, the gradient and Hessian simplify before the worst-case perturbation is approximated by a first-order Taylor expansion.The resulting Flat-RLVR update minimizes −L(θt + ϵ∗).Source: §A.1 · p. 16

A.2. Proof of Lemma 4.4

  • The change in the correct-support log-probability gap is approximated to first order, and the Flat-RLVR gap change is expressed as the GRPO term plus a residual.The residual is Rt = −αρ Btgt/∥gt∥2 − αRT.Source: §A.2 · p. 16

A.3. Proof of Proposition 4.5

  • The proof sums the pointwise difference between Flat-RLVR- and GRPO-induced gap changes across training steps and applies the proposition’s dominance condition to obtain its inequality.Source: §A.3 · p. 17

B. Adaptations for LLM post-training

The adaptations tune SAM’s perturbation radius by prompt difficulty, use Gradient Centralization around its ascent and descent steps, and anneal the radius over training. The passages report that these choices support stability and mitigate entropy or reward collapse.

B.1. Focal Weight

  • The focal weight and perturbation radius are adjusted according to prompt difficulty, reducing perturbation influence at low success rates to prevent unstable updates while preserving a signal for rare correct trajectories at high success rates.Source: §B.1 · p. 17

B.2. Gradient Centralization

  • Applying Gradient Centralization before SAM’s descent step substantially mitigates entropy collapse compared with applying it only before ascent.Both variants in Figure 6 also use SAM-n and cosine annealing.Source: Figure 6 · p. 17, §B.2 · p. 18

B.1. Focal Weight

  • Computing the focal weight from per-batch success rate uses the training batch size and the number of correct samples in that batch.Source: §B.1 · p. 17

B.2. Gradient Centralization

  • Gradient Centralization before SAM’s ascent step reduces gradient variance and provides a more reliable perturbation direction for the sparse-reward, high-variance RLVR setting.Source: §B.2 · p. 17

B.3. Annealing SAM ρ

  • Without cosine annealing, training reward collapses to zero midway through training, halting effective learning.Source: Figure 7 · p. 18, §B.3 · p. 18
  • Cosine annealing gradually reduces ρ from its initial to minimum value, attenuating perturbations to limit early entropy collapse and promote later stable convergence.Source: §B.3 · p. 18

C. Ablation Study

The ablation reports separate effects for SAM, cosine annealing, Gradient Centralization, and Focal Weight, and concludes that stabilizing strategies help flatness-aware optimization improve Pass@k reasoning coverage. A separate appendix passage describes a scalar-cache approximation for SAM when ρ is sufficiently small.

C. Ablation Study

  • The ablation concludes that flatness-aware optimization can improve Pass@k reasoning coverage, with strategies needed to stabilize training and amplify its effect.Source: C. Ablation Study · p. 18

D. The lightweight approximation for Memory Efficiency.

C. Ablation Study

E. Further Results

The appendix presents per-benchmark Pass@K curves across five mathematical benchmarks and two backbones, and describes a solvable-problem coverage analysis based on 256 samples per problem. That analysis classifies each problem by whether it is solvable by the base model, the RL-tuned model, both, or neither.

E.1. All Pass@K curves

  • Figure 8 reports Pass@K curves from K = 1 to K = 256 for MATH500, AMC23, AIME24, AIME25, and AIME26 on both backbones.The curves complement the aggregated results in Section 6.2.Source: §E.1 · p. 19

E.2. Solvable-problem coverage

  • For solvable-problem coverage, a problem counts as solvable by a model if at least one of its 256 sampled completions is correct.Source: §E.2 · p. 19
  • The analysis assigns each problem to one of four categories: solvable by both models, forgotten by the RL-tuned model, newly solved by it, or solvable by neither.Reported percentages are averaged across benchmarks.Source: §E.2 · p. 19

F.1. Full Training Details

The appendix records training configurations for Flat-RLVR and F-GRPO, specifies the evaluation protocol, and documents dataset and model licenses. It also provides supplementary tables and Pass@K curves for the reported experiments.

F.1. Full Training Details

  • Training uses token-level loss and follows TRL’s official guidelines for GRPO variants, with hardware allocations varying across models.Full hyperparameters are summarized in Table 12.Source: §F.1 · p. 19, Table 12 · p. 21

F.2. Full Training Details: Flat-RLVR

  • Flat-RLVR training decays the SAM hyperparameter ρ and uses a lightweight approximation for memory efficiency.A lower bound on the focal-weight-modulated perturbation radius preserves a non-trivial perturbation when the focal weight is small.Source: §F.2 · p. 19, Table 11 · p. 20

F.3. Full Training Details: F-GRPO

  • F-GRPO scales each prompt group’s relative advantage by a focal weight based on its empirical success rate, down-weighting high-success prompts to preserve solution diversity.The implementation sets γ = 1 and keeps the other hyperparameters identical to the GRPO baseline.Source: §F.3 · p. 19

F.4. Evaluation Protocol

  • Table 10 categorizes per-benchmark solvable-problem coverage by whether base and RL-tuned models solve each problem, with averages corresponding to Table 1.The table reports forgotten problems for (✓/×) and newly solved problems for (×/✓), for both backbones.Source: Table 10 · p. 20
  • Evaluation generates 256 samples per problem and reports Pass@1 through Pass@256, with completion lengths adjusted by benchmark family.MATH500 and AMC23 use 1024 tokens, AIME benchmarks 2048, and GSM8K 192; all benchmarks use a single system prompt.Source: §F.4 · pp. 19–20

F.5. Licenses

  • The appendix lists licenses for datasets and models, including Apache License 2.0 for the AIME datasets and several Qwen models.It reports no explicit license for AMC23 and MBPP; Qwen2.5-3B-Instruct and Llama-3.2-1B-Instruct use their respective Qwen Research and Llama 3.2 Community licenses.Source: §F.5 · p. 20

F.4. Evaluation Protocol

  • Figure 8 plots Pass@K from K = 1 to 256 on five in-domain benchmarks, comparing the base model, GRPO, and Flat-GRPO across two backbones.Qwen2.5-Math-1.5B appears in the top row and Llama-3.2-1B-Instruct in the bottom row.Source: Figure 8 · p. 21

Related papers

Loading 2610.29281v1…