Improving a reasoning model using its own on‑policy experience can yield large gains, yet the inner loop is fragile. Terminal verifiers give reliable but sparse supervision, while dense same‑model guidance may reinforce false confidence or over‑focus learning on a narrow solution mode.
FlowBalance addresses this by learning a normalized distribution over complete responses. For each on‑policy trajectory, a frozen training‑time view of the policy leverages privileged context to compute token‑level log‑probability gains $\Delta\log p_t$, which are summed into a trajectory‑level self‑guidance score $G(\tau)$:
$$G(\tau)=\sum_{t=1}^{|\tau|}\Delta\log p_t$$
The score is then calibrated with the verifier‑derived group advantage $A_g$: positive‑advantage trajectories keep the score, negative‑advantage trajectories reverse its sign, and when the rollout group shows no preference the guidance is disabled. The calibrated score becomes an energy term $E(\tau)= -A_g\,G(\tau)$, which exponentially reweights a reference policy $\pi_{ref}$:
$$\pi_{new}(\tau) \propto \pi_{ref}(\tau)\,e^{-E(\tau)}$$
Only a single log‑partition estimate $\log Z_g$ per rollout group is required; trajectory balance then fits the normalized target distribution without a separate token‑level imitation loss.
Our analysis establishes:
- Within‑group contrast preservation;
- A minimum‑change reverse‑KL characterization;
- Monotonic verifier control of the target reward;
- An exact correction against false‑positive self‑guidance on rejected responses.
Empirically, FlowBalance improves average performance over FlowRL on both Qwen3‑4B and Qwen3‑8B for mathematical reasoning, accelerates training, enhances stability, avoids the response‑length collapse seen in direct OPSD, and shows higher correct‑strategy diversity in a controlled AIME24 diagnostic.
Review: FlowBalance’s verifier‑grounded energy reweighting successfully merges sparse reliable supervision with dense self‑guidance, boosting reasoning ability while preserving training robustness.