Why this matters
Most on‑policy self‑improvement loops for reasoning models suffer from a tension: terminal verifiers give sparse but reliable rewards, while dense same‑model guidance can accidentally reinforce false confidence or collapse solution diversity. FlowBalance targets that tension by making dense within‑trajectory guidance conditional on verifier outcomes and converting the composite signal into a normalized next‑policy target via trajectory balance.
Key Findings
- Outcome‑calibrated dense guidance: token‑level log‑probability gains from a frozen privileged‑hindsight view are aggregated per trajectory and sign‑gated by the verifier advantage (retain on positive advantage, reverse on negative, disable when advantage is zero), preventing confident failures from becoming self‑reinforcing.
- Profiled trajectory balance: the composite energy (verifier + guided gain) is converted into a normalized distribution over complete responses with one log‑partition estimate per rollout group, keeping within‑group probability contrasts and limiting drift from a reference policy.
- Empirical gains on math reasoning: improves average performance over FlowRL on Qwen3‑4B and Qwen3‑8B (core four‑benchmark average +1.67 points on Qwen3‑4B and +1.98 on Qwen3‑8B), reaches 0.50 AIME24 validation accuracy in ~100 steps on Qwen3‑8B versus ~143 for GRPO, and avoids OPSD’s response‑length collapse while increasing correct‑strategy diversity.
Who it's for and tradeoffs
Great fit if you work on iterative self‑improvement for reasoning LLMs and need a principled way to combine sparse terminal verification with dense within‑trajectory evidence without imposing a separate token‑level imitation loss. It is particularly relevant for research or training setups that can generate rollout groups and run verifiers (e.g., mathematical or code reasoning benchmarks).
Look elsewhere if your setting lacks reliable terminal verifiers, cannot store or compute privileged hindsight scores, or if per‑rollout grouping and log‑partition estimates are impractical for your compute budget. FlowBalance shapes which response distribution is learned via a dense branch that is not optimized as a standalone imitation objective, so its benefits depend on verifier quality and rollout group diversity.
Method snapshot
FlowBalance builds a composite trajectory energy E(y) = η_A·A(y) + β_G·G_H(y)·sgn(A(y)), where A(y) is the verifier advantage and G_H(y) the privileged‑hindsight trajectory gain. That energy reweights a reference policy into a Gibbs target; profiled trajectory balance then fits the normalized target with one log‑partition per rollout group, yielding a minimum‑change reverse‑KL displacement from the reference at the attained energy level.