One-sentence Summary

Researchers from Tencent HY LLM Frontier and the University of Pennsylvania propose FlowBalance, a verifier-grounded self-improvement method that learns a normalized distribution over complete responses by calibrating token-level self-guidance scores with verifier-derived group advantage via trajectory balance to reweight a reference policy, improving mathematical reasoning over FlowRL on $\text{Qwen3-4B}$Qwen3-4B and $\text{Qwen3-8B}$Qwen3-8B while enhancing training speed and stability, avoiding response-length collapse, and increasing correct-strategy diversity in $\text{AIME24}$AIME24.

Key Contributions

• Introduces FlowBalance, a verifier-grounded self-improvement method that learns a normalized distribution over complete responses by combining a frozen privileged-hindsight self-guidance score with verifier-derived group advantages, retaining guidance on positive-advantage trajectories, reversing it on negative ones, and disabling it when no outcome preference exists.
• Fits the reference-supported target via profiled trajectory balance with a single log-partition estimate per rollout group, eliminating the need for a separate token-level imitation loss; theoretical analysis establishes within-group contrast preservation, a minimum-change reverse-KL characterization, monotonic verifier control of target reward, and an exact anti-self-confirmation correction on rejected responses.
• On mathematical reasoning, FlowBalance improves average performance over FlowRL on both Qwen3-4B and Qwen3-8B across five benchmarks, reaches the AIME24 validation threshold in fewer updates than GRPO, maintains training stability, avoids direct OPSD’s response-length collapse, and exhibits higher correct-strategy diversity in a controlled AIME24 diagnostic.

Introduction

Post-training reasoning models aim to improve by learning from their own sampled solutions, but this self-improvement loop faces two core challenges. First, reinforcement learning with verifiable rewards (RLVR) provides reliable outcome signals only as sparse, terminal feedback, which cannot capture fine-grained evidence along long reasoning paths. Second, dense self-guidance from a frozen copy of the policy, which scores each token using training-only context, is inexpensive but untrustworthy: it may favor rejected trajectories, shorten reasoning, or reinforce the model's own errors, leading to self-confirmation.

To address these issues, the authors introduce FlowBalance, a distributional self-improvement operator that combines verifier-grounded advantages with outcome-calibrated privileged-hindsight guidance. The key innovation is a trajectory-level energy function where the verifier determines the direction of improvement, and dense self-guidance refines that direction only when the verifier confirms success, reversing it on failures. This energy defines a reference-supported Gibbs target, fitted via profiled trajectory balance, which preserves all within-group probability contrasts and acts as a minimum-change update from the reference policy.

The authors provide theoretical guarantees showing that FlowBalance converts false-positive self-guidance into a verifier-grounded correction, maintains monotonic verifier control over target reward, and is information-efficient by profiling only one scalar partition per rollout group. Empirically, on Qwen3-4B and Qwen3-8B, FlowBalance outperforms GRPO, OPSD, RLSD, and FlowRL across mathematical reasoning benchmarks, reaches AIME24 validation accuracy faster than GRPO, remains stable over extended training, and preserves higher semantic strategy diversity, avoiding the response-length collapse seen under direct dense-guidance methods.

Method

FlowBalance is a self-improvement procedure that turns a group of the model’s own sampled reasoning trajectories into a normalized next-policy target. The core design combines a sparse verifier signal with a dense, privileged-hindsight self-guidance signal, and then fits the resulting distribution via trajectory balance. The authors frame the problem in the context of a trainable policy $\pi_\theta(y \mid x)$πθ​(y∣x) that factorizes over tokens, where the state at token $t$t is the prompt and previously generated tokens. At the start of each training iteration, the current parameters are snapshotted as $\theta^-$θ−, and this frozen rollout policy $\pi_{\theta^-}$πθ−​ generates experience. A fixed reference policy $\pi_{\mathrm{ref}}$πref​, which is a copy of the initial checkpoint, is used to control drift and provide support for the target distribution.

For each prompt, the frozen policy samples a group of $N$N responses. A verifier assigns a terminal reward $R_i$Ri​ to each response based on final-answer correctness. The authors compute a stopped group-relative advantage for each response:

where $\mu_R(x)$μR​(x) and $\sigma_R(x)$σR​(x) are the mean and standard deviation of the rewards within the sampled group. This normalization ensures that the advantage is relative to the model’s current performance on that specific prompt, providing a grounded but sparse outcome signal.

To obtain dense, token-level guidance, the authors introduce a privileged-hindsight view. Each training prompt is paired with training-only context $c$c, such as a reference solution or task feedback. The frozen snapshot is evaluated through two views: the rollout view $\pi_{\mathrm{roll}}(\cdot \mid s_t) = \pi_{\theta^-}(\cdot \mid x, y_{<t})$πroll​(⋅∣st​)=πθ−​(⋅∣x,y<t​), which generates the response without observing $c$c, and the privileged-hindsight view $\pi_{\mathrm{H}}(\cdot \mid s_t, c) = \pi_{\theta^-}(\cdot \mid x, c, y_{<t})$πH​(⋅∣st​,c)=πθ−​(⋅∣x,c,y<t​), which sees $c$c only after the response has been sampled. This hindsight view scores the same tokens that were already sampled but does not generate a replacement trajectory and receives no gradient. This construction is related to on-policy self-distillation, but in FlowBalance it serves a narrower role: it supplies a stopped dense feature for defining a trajectory-level target, rather than being optimized as a separate token-level imitation loss.

The dense self-guidance is quantified as a clipped token-level hindsight gain:

which is aggregated over the complete response to produce a trajectory-level feature:

This feature measures how much the privileged-hindsight view raises or lowers the average sampled-token log probability relative to the fixed reference policy. All quantities are stopped, and no token is resampled from $\pi_{\mathrm{H}}$πH​.

Because dense self-guidance can be useful while still being wrong about the verified outcome, FlowBalance uses the group-relative advantage to calibrate its direction. The trajectory energy is defined as:

where $\eta_A, \beta_G \geq 0$ηA​,βG​≥0. If $A(y) > 0$A(y)>0, positive guidance increases the trajectory energy, reinforcing the response. If $A(y) < 0$A(y)<0, positive guidance is reversed, preventing a confidently scored failure from becoming self-reinforcing supervision. If $A(y) = 0$A(y)=0, the dense branch is disabled. This outcome-calibrated energy defines an unnormalized target:

which is normalized over the realized rollout group to form the target distribution:

The reference policy retains support, the advantage supplies verified outcome direction, and the hindsight gain provides dense within-trajectory evidence. The partition term converts relative trajectory energies into a probability-conserving target, allowing self-guidance to refine the distribution within the verifier’s direction without becoming an independent local loss.

The authors fit this target through trajectory balance. The complete-trajectory balance equation is:

and the corresponding residual for a sampled response is:

At zero residual, the partition cancels in pairwise probability ratios, meaning the energy controls relative preference among the model’s own sampled experiences, while the log-partition absorbs only the common prompt-level offset. The main loss is the mean squared residual:

Gradients are taken only through the trainable-policy log probabilities; rewards, advantages, self-guidance scores, partition estimates, and sampled responses are all stopped.

The same principle can be extended to subtrajectories. Let $Z_{\mathrm{FlowBalance}}(s)$ZFlowBalance​(s) denote the continuation partition from state $s$s, with $Z_{\mathrm{FlowBalance}}(s_{\mathrm{term}}) = 1$ZFlowBalance​(sterm​)=1. A per-token shaped increment is defined as:

For an interval $i:j$i:j, the subtrajectory residual is:

Minimizing the expected squared subtrajectory residual provides a denser fitting objective along long responses, though the experiments in the paper use the complete-response implementation unless otherwise stated.

For the partition, the authors use a profiled group estimator rather than a learned prompt-conditioned estimator. Each response implies an estimate:

and the group estimate is the average:

Gradients are stopped through this estimate. The self-improvement loop is explicit: the frozen snapshot first generates experience without privileged context, then provides a training-only hindsight score on those same trajectories. The verifier grounds the direction of the score, and trajectory balance internalizes the resulting normalized distribution into the next policy snapshot.

Experiment

FlowBalance is evaluated on Qwen3-4B and Qwen3-8B for mathematical reasoning, comparing against GRPO, OPSD, RLSD, and FlowRL baselines. It consistently improves average accuracy across five benchmarks, with the largest gains on AIME24 Pass@16, while also accelerating early convergence, maintaining stability over longer training, and avoiding the response-length collapse seen in direct distillation. Ablations confirm that both verifier grounding and self-guidance strength contribute to performance. A separate LLM-judged diagnostic shows FlowBalance produces more diverse correct solution strategies than GRPO or RLSD, illustrated by a case where it finds a hidden rectangular-box embedding instead of the standard Cayley-Menger determinant approach.

FlowBalance consistently improves mathematical reasoning accuracy over GRPO, OPSD, RLSD, and FlowRL on both Qwen3-4B and Qwen3-8B backbones, with the largest gains on Qwen3-8B. The method also shows faster convergence and better stability during training, while avoiding the response-length collapse seen in OPSD. FlowBalance achieves the highest average accuracy on both backbones, improving over GRPO by about 2 points and over OPSD by over 10 points on Qwen3-4B and over 26 points on Qwen3-8B. On Qwen3-8B, FlowBalance obtains the best mean on every reported benchmark, whereas OPSD performs markedly worse than all other methods. FlowBalance reaches a given AIME24 validation accuracy in about 100 steps compared to roughly 143 for GRPO, and it maintains near-peak accuracy over 400 steps while GRPO degrades after step 180. FlowBalance sustains substantially longer reasoning trajectories than OPSD, which rapidly collapses to short responses.

The ablation study sweeps the verifier coefficient and the self-guidance coefficient around their default settings, with the highest average accuracy achieved at the largest verifier coefficient. The default configuration, which uses a verifier coefficient of 15 and a self-guidance coefficient of 1, yields the best aggregate performance among the tested values. Increasing the verifier coefficient from 5 to 15 improves the five-benchmark average from 65.65 to 67.61. The default self-guidance coefficient of 1 is paired with the best-performing verifier coefficient, indicating that the default settings are optimal within the swept ranges.

The table contrasts two correct solution strategies for an AIME geometry problem: one uses the standard Cayley-Menger determinant approach, while the other exploits a hidden rectangular box embedding to simplify volume calculation. Both arrive at the same final answer, but the different mathematical representations illustrate how FlowBalance can produce diverse correct strategies beyond a single dominant template. GRPO's correct trajectory follows the standard Cayley-Menger determinant route using pairwise distances. FlowBalance's correct trajectory identifies a hidden 4x5x8 box embedding from the edge lengths, then uses a scalar triple product for volume. Both strategies yield the same final answer, but they rely on distinct mathematical representations and tools.

FlowBalance outperforms GRPO, OPSD, RLSD, and FlowRL on mathematical reasoning across Qwen3-4B and Qwen3-8B, with the largest gains on the larger backbone, showing faster convergence, better stability, and avoiding response-length collapse. Ablations confirm that the default verifier and self-guidance coefficients are optimal, with higher verifier strength improving accuracy. A case study illustrates that FlowBalance can discover diverse correct solution strategies, such as exploiting a hidden geometric embedding, unlike GRPO which follows a standard approach.