Flow Map Reward Alignment

Reward alignment of generative flows through deterministic optimal control

Standard · KL reward tilt

$$\max_{\nu}\ \lambda\,\mathbb{E}_{\nu}[r] - \mathrm{KL}(\nu\,\|\,\rho_1) \;\;\Rightarrow\;\; \tilde\rho_1 \propto e^{\lambda r}\rho_1$$

Reweights the base distribution. The pathwise KL is tractable for diffusions, so the flow is treated as a diffusion.

Ours · deterministic optimal control

$$\max_{u}\ \lambda\, r(x_1^u) - \tfrac12\!\int_0^1 \lVert u_t\rVert^2 dt, \quad \dot x_t^u = b_t(x_t^u) + u_t$$

Transports individual samples toward higher reward. Solved directly on the deterministic flow, using the flow map.

More below! 👇

Flow Map Reward Guidance: few-step reward-guided samples

Inference-time guidance ICML 2026

How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance

Jerry Y. Huang^*, Justin Lin^*, Sheel Shah, Kartik Nair, Nicholas M. Boffi

^*Equal contribution

Training-free reward guidance of a pre-trained flow map, cast as deterministic optimal control, in as few as three network evaluations.

Project page arXiv Code

WTF: reward-aligned samples at few steps

Fine-tuning NeurIPS 2026 · Spotlight

WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps

Abbas Mammadov^†*, Jerry Y. Huang^†*, Justin Lin^*, Partha Kaushik, Sheel Shah, Kartik Nair, Yee Whye Teh, Nicholas M. Boffi

^†Project lead · ^*Equal contribution

Reward fine-tuning with an optimal transport regularizer, solved as deterministic optimal control with a flow map, which gives simulation-free reinforcement learning.

Project page arXiv Code