MixGRPO — Unlocking Flow-based GRPO Efficiency with Mixed ODE-SDE¶
| Field | Value |
|---|---|
| arXiv | 2507.21802 |
| Submitted | 2025-07-29 (revised 2026-03-20) |
| Venue | — (preprint) |
| Authors | Junzhe Li, Yutao Cui, Tao Huang, Yinping Ma, Chun Fan, Yiming Cheng, Miles Yang, Zhao Zhong, Liefeng Bo |
| GitHub | https://github.com/Tencent-Hunyuan/MixGRPO |
| Paradigm | Policy Gradient — SDE + gradient confined to a sliding window; ODE everywhere else |
| Cites | FlowGRPO (2505.05470), DanceGRPO (2505.07818), DDPO, GRPO |
| Cited by | HunyuanImage 3.0 (production training pipeline) |
Context¶
FlowGRPO and DanceGRPO established that applying GRPO to flow matching requires converting the ODE to an SDE for every denoising step — giving a tractable importance ratio but forcing the entire \(T\)-step trajectory through SDE sampling. MixGRPO is the first work to ask: do all \(T\) steps actually need to be stochastic? The answer is no: only the steps that contribute to the importance ratio and gradient need SDE treatment. This file builds on the FlowGRPO background; see flow_grpo.md for the ODE→SDE derivation.
Problem 1 — Full \(T\)-step SDE is slow and blocks high-quality ODE solvers¶
Issue: FlowGRPO and DanceGRPO run SDE sampling for all \(T\) denoising steps. Two compounding costs arise: (1) SDE steps are first-order Euler-Maruyama — they cannot use high-order ODE solvers (DPM-Solver++, Heun) that achieve the same quality with \(2\text{–}3\times\) fewer steps; (2) gradient tracking through all \(T\) SDE steps requires storing all \(T\) intermediate activations — \(O(T \cdot d)\) peak memory.
Idea: Confine both SDE sampling and gradient tracking to a contiguous sliding window \(\mathcal{W}(l) = \lbrace{}l, l{+}1, \ldots, l{+}w{-}1\rbrace\) of \(w\) denoising steps — here \(l\) is the window's left boundary (its starting step index) and \(w = |\mathcal{W}|\) is the fixed window size (number of optimised steps). Outside the window, use a standard deterministic ODE with no gradient tracking. Slide the window through the trajectory over training iterations (advancing \(l\)), so different segments accumulate gradient updates over time. The paper's intuition: early denoising steps shape global structure and later steps fine detail, so the window enacts a curriculum that optimises high-impact early steps first.
Why this works: The GRPO importance ratio \(\rho_t\) (per-step ratio of the new vs. old policy's transition density at step \(t\); defined below) is meaningful only at SDE steps (deterministic ODE steps have \(\rho_t = 1\) trivially, contributing nothing to the gradient). Confining SDE to \(w\) steps reduces memory to \(O(w \cdot d)\), and allows the non-window steps to use fast ODE solvers — improving trajectory quality at the tail (high detail) and head (coarse structure). Sliding the window ensures the entire trajectory is eventually covered across iterations, so the full reward landscape is optimised even though only \(w\) steps are updated at each iteration.
Result: MixGRPO achieves nearly 50% lower training time than DanceGRPO while outperforming it across multiple dimensions of human-preference alignment (abstract) — i.e. the sliding window removes most of the SDE/gradient cost without sacrificing reward.
Three trajectory zones¶
For a trajectory with \(T\) denoising steps, window starting at step \(l\). Here \(x_t\) is the noisy latent at flow time \(t\), with \(x_T\) (equivalently \(x_1\)) pure Gaussian noise and \(x_0\) the clean image; the trajectory runs \(x_T \to \cdots \to x_0\):
Outside window (ODE, no grad), where \(v_\theta(x_t, t, c)\) is the model's velocity prediction conditioned on prompt \(c\) and \(\Delta t\) is the step size:
Inside window (SDE with score correction, with grad), where \(\sigma_t\) is the SDE diffusion (noise) coefficient at step \(t\) and \(\epsilon_t\) is freshly sampled Gaussian noise:
where \(\hat{x}_0 = x_t - tv_\theta(x_t,t,c)\) is the Tweedie estimate of the clean image. The per-step transition inside the window is Gaussian, with mean \(\mu_\theta(x_t,t,c)\) equal to the deterministic part of the update above:
Window sliding — every \(\tau\) training iterations the left boundary advances by stride \(s\) (clamped so the window stays inside the trajectory):
Problem 2 — Window width is fixed; extreme compression still needs high-order ODE¶
Issue: Even with a window of \(w=4\) steps out of \(T=25\) (the paper's optimum), the 21 ODE steps outside the window still use first-order Euler. This is acceptable for moderate \(T\), but for larger \(T\) or video (where \(T\) may reach 100+), Euler ODE quality degrades. A variant that uses a high-order ODE solver outside the window would compound the speedup.
Idea — MixGRPO-Flash: Shrink the window to \(w=1\) step and replace all ODE steps with DPM-Solver++ (a high-order ODE solver) at compression ratio \(\tilde{r} < 1\) — the fraction of ODE steps retained after solver compression:
Why this works: DPM-Solver++ is a high-order (second/third order) solver that achieves Euler-equivalent quality at \(\sim 2\text{–}3\times\) fewer steps. Combining \(w=1\) with \(\tilde{r} \approx 0.3\text{–}0.4\) gives a total speedup:
Result: MixGRPO-Flash further reduces training time by 71% (abstract) with almost no degradation in reward performance — the largest efficiency gain in the policy-gradient chain.
Training Objective¶
GRPO-clip loss restricted to window steps only; ODE steps contribute nothing. The outer average is over the group of \(N_g\) samples generated per prompt (group size), and the inner average is over the \(|\mathcal{W}| = w\) steps in the window:
Here \(\hat{A}^{(i)}\) is the group-relative advantage of sample \(i\) — its reward standardised within the group, \(\hat{A}^{(i)} = (r^{(i)} - \mathrm{mean}(r)) / \mathrm{std}(r)\) — and \(\epsilon\) is the PPO clip range (the clip caps the ratio to \([1{-}\epsilon, 1{+}\epsilon]\) so a single update cannot move the policy too far).
Importance ratio \(\rho_t^{(i)}\) (window steps only) — the ratio of the current policy \(\pi_\theta\) to the rollout policy \(\pi_{\theta_\text{old}}\) for taking the realised transition \(x_{t-\Delta t}^{(i)}\), written out from the Gaussian log-densities (means \(\mu_\theta\), \(\mu_{\theta_\text{old}}\) under current and frozen rollout parameters):
Algorithm¶
Input: pretrained v_θ, reward r, prompt dist p_c, N_g,
window width w, slide interval τ_iter, stride s, total steps T
Initialize: l ← 0
Repeat:
Every τ_iter iterations: l ← min(l + s, T - w)
1. Sample prompts {c_j}
2. For each c_j, generate N_g trajectories (mixed ODE/SDE):
x_T^(i) ~ N(0,I)
ODE (no grad): x_T^(i) → x_{l+w}^(i) # steps T ... l+w+1
SDE (with grad): x_{l+w}^(i) → x_l^(i) # steps l+w ... l+1 (window)
For t = l+w, ..., l+1:
x̂_0 ← x_t - t·v_{θ_old}(x_t^(i), t, c)
μ_{θ_old} ← x_t - v_{θ_old}·Δt + (σ_t²Δt/2t²)·(x̂_0 - x_t)
x_{t-Δt}^(i) ← μ_{θ_old} + σ_t√Δt·ε, ε ~ N(0,I)
ODE (no grad): x_l^(i) → x_0^(i) # steps l ... 1
3. Compute R^(i) = r(x_0^(i), c_j)
4. Â^(i) = (R^(i) - mean) / std
5. For K gradient steps:
For t ∈ W(l), i = 1,...,N_g:
μ_θ ← x_t - v_θ·Δt + (σ_t²Δt/2t²)·(x̂_0 - x_t) # current θ
ρ_t^(i) = exp(-(‖x_{t-Δt}^(i) - μ_θ‖² - ‖x_{t-Δt}^(i) - μ_{θ_old}‖²) / (2σ_t²Δt))
L ← -mean[ min(ρ·Â, clip(ρ, 1-ε, 1+ε)·Â) ]
θ ← θ - η ∇_θ L
6. θ_old ← θ
MixGRPO-Flash variant: replace ODE steps outside window with DPM-Solver++ at rate r̃;
shrink window to w=1; total effective steps T̃ = 1 + (T-1)·r̃.
Efficiency Comparison¶
| Method | SDE steps / traj | Grad steps | ODE solver outside | Reported speedup |
|---|---|---|---|---|
| FlowGRPO | All \(T\) | All \(T\) | Euler (N/A) | 1× (baseline) |
| FlowGRPO-Fast | 1–2 (branch point) | 1–2 | Euler | ~\(T/2\)× |
| MixGRPO | \(w\) (sliding) | \(w\) | Euler | \(T/w\)× |
| MixGRPO-Flash | 1 | 1 | DPM-Solver++ | 71% time reduction |
Recommended hyperparameters (ablated in paper): \(w=4\), \(\tau_\text{iter}=25\), \(s=1\) for \(T=25\).
Limitations¶
| Problem | Note |
|---|---|
| Window width \(w\) is a hyperparameter | Sensitivity studied in paper; \(w=4\) recommended |
| SDE noise inside window still causes artifacts | CPS can replace the SDE step inside \(\mathcal{W}(l)\) |
| Ratio imbalance within window persists | GRPO-Guard is compatible |
| Still requires SDE for any gradient | Fully solver-agnostic: see DGPO, AWM |