ConicFlow: Constrained Inference-Time Optimal Control for Flow Matching Models

Reward-guided generation with hard terminal constraints via receding-horizon, primal–dual optimal control of pretrained flow matching models.

Shervin Khalafi^*Kasra Fallah^*Sung Cho Dongsheng Ding James Anderson Alejandro Ribeiro

University of Pennsylvania · Columbia University · University of Tennessee, Knoxville

Equal contribution

Read overview ↓Paper (arXiv) ↗BibTeX

CONSTRAINED STEERING IN FLOW MATCHING — SCHEMATIC

Schematic of flow matching from source distribution p₀ to data distribution p₁. Receding-horizon control perturbs the pretrained velocity field to satisfy terminal requirements. Illustration only, not measured experimental output.

01 / Problem formulation

Reward optimization subject to terminal feasibility

A pretrained flow matching model specifies a velocity field \(b_t\) that transports an initial sample to the data distribution. ConicFlow introduces a control input during inference, leaving the model parameters unchanged. It seeks a terminal state with high reward while penalizing the magnitude of interventions and imposing hard output constraints.

\[\begin{aligned}\min_{u(t)}\quad &J_C(u)=\int_0^1\frac{\|u(t)\|^2}{2\lambda}\,dt-r(x_1^u)\\[2pt]\text{subject to}\quad &\dot x_t^u=b_t(x_t^u)+u(t),\quad x_0\text{ given},\\&h(x_1^u)\le 0.\end{aligned}\]

Problem (P-C), manuscript §1–2. Here \(r\) is a terminal reward and \(h:\mathbb R^d\to\mathbb R^m\) represents hard terminal requirements. Smaller \(\lambda\) increases the penalty on intervention.

02 / Receding-horizon primal–dual control

Short control horizons, full terminal objectives

Optimizing the entire trajectory in path space is expensive for neural flow models. ConicFlow solves a sequence of shorter constrained optimization problems and executes only the first transition before replanning. The central methodological question is where to evaluate terminal reward and constraints.

RH-H · Horizon endpoint

Truncated terminal evaluation

Evaluate reward and constraints at the end of the local planning window, \(x_{j+H_j}\). This reduces the optimization horizon but introduces a mismatch with the original terminal-time objective.

\[\Phi_j(x)=x,\quad g_j=h(x_{j+H_j}).\]

RH-T · Terminal rollout

Flow-map terminal prediction

Propagate the locally optimized endpoint to time \(N\) using the pretrained, uncontrolled flow map \(X_{j+H_j,N}\). Evaluate reward and constraints at this predicted terminal state.

\[\hat x_N=X_{j+H_j,N}(x_{j+H_j}),\quad g_j=h(\hat x_N).\]

What changes when we predict the terminal state?

Reward and constraints are evaluated at a predicted terminal sample.

Interactive conceptual timeline, not a numerical trajectory generated by the algorithm.

For either local problem, ConicFlow alternates a primal minimization of the local Lagrangian with projected dual ascent, \(\mu^{(s+1)}=[\mu^{(s)}+\eta_\mu g_j(x_j^{(s)})]_+\). After \(K\) dual updates, it executes the first optimized transition and warm-starts the next window. See Algorithm 1 in the manuscript.

03 / Theory

Approximation and convergence guarantees

The analysis connects continuous-time constrained inference to finite-step discretization, finite receding horizons, and inexact local primal–dual solutions. These bounds are conditional on the specific assumptions stated in the paper.

Lemma 1

Discretization error

\(O(1/N)\)

Optimal-value difference between continuous and discrete constrained problems under Assumption 1, including feasibility repairability.

Lemmas 2–3

Horizon bias

\(O((N-H)/N)\)

Constrained RH-H and RH-T under affine flow, bounded amplification, convex reward/constraints, and uniform feasibility. Constants differ.

Theorem 1

Inexact local solver

\(O(K^{-1/2})\)

Trajectory and terminal-violation bounds relative to exact receding-horizon solutions, with sensitivity-dependent factors.

Important qualification: RH-T is not proven uniformly better than RH-H. Remark 1 gives a sufficient condition for its tighter constrained-horizon constant (\(\kappa_A\le\sqrt{2}\)). Under these assumptions, a smaller bound constant is obtained.

04 / Experiments

Reward–energy trade-offs and constraint satisfaction

The paper evaluates both unconstrained reward-guided generation and hard-constrained sampling. In low-dimensional experiments, a flow-map model is trained on two-moons data. Image experiments use FLUX.1-dev with a flow-map LoRA and edit rewards from CLIP, ImageReward, and PickScore.

Figure 2 · Main image editing result

Matched reward, substantially less deviation

Original paper Figure 2 image-editing comparisons: uncontrolled sample, ConicFlow, HardFlow and FMRG at approximately matched rewards.

Comparison of unconstrained image edits at approximately matched reward. ConicFlow achieves the edit with lower control energy. Source: original paper, Figure 2.

Figure 3 · Hard constraints

Feasibility and energy as thresholds increase

ConicFlow and HardFlow energy and feasibility plotted against the threshold multiplier.

Constraint feasibility versus required PickScore threshold and corresponding control energy; plotted from the paper.

Figure 3 · Multi-constraint editing

Simultaneous desired attributes

Original multi-constraint image comparisons from ConicFlow paper.

Image-editing comparisons under multiple terminal constraints, where quality and feasible generation must be maintained jointly.

Figure 1 · Low-dimensional results

Two-moons reward geometry

Two moons data geometry and reward regions from the ConicFlow paper.

Uncontrolled two-moons samples and two terminal reward functions, including an off-manifold reward.

Figure 2 · High-dimensional results

Reward–energy Pareto frontier

ConicFlow versus baselines: terminal reward against control energy for image edits.

Reward plotted against integrated control energy for unconstrained text-guided image editing.

Figure 3 · Single-constraint editing

Constraint strength and visual fidelity

Constrained image editing examples for increasing threshold.

Constrained edits across thresholds with ConicFlow and baseline comparisons.

Unconstrained generation

Comparable rewards, less intervention

ConicFlow improves the reward-versus-control-energy trade-off across reported image-editing tasks relative to FMRG and HardFlow. The manuscript's Figure 2 shows comparisons at approximately matched terminal reward.

Lower-dimensional comparisons demonstrate that larger receding horizons can improve performance for off-manifold rewards; RH-H and RH-T do not consistently dominate one another.

Constrained generation

Hard terminal feasibility

With PickScore threshold constraints, ConicFlow maintains feasibility close to 100% across the evaluated thresholds, whereas the reported HardFlow feasibility decreases as constraints become harder.

The paper reports up to a 20-fold reduction in control energy relative to HardFlow in the constrained image experiments; this is not a uniform bound over all settings.

Table 1 · From the manuscript

Mean reward gain over a shared control-energy budget

Reported mean ± variation across \(n=4\) (prompt, seed) pairs per reward, relative to uncontrolled samples. Higher is better. Values use each reward's native units.

Method	CLIP ↑	ImageReward ↑	PickScore ↑
ConicFlow	0.130 ± 0.018	2.048 ± 0.455	6.158 ± 0.728
HardFlow	0.118 ± 0.021	1.279 ± 0.374	1.736 ± 0.472
FMRG	0.065 ± 0.008	1.278 ± 0.504	2.670 ± 0.073

Values are reported in Table 1 of the ConicFlow paper. The original reward–energy curves and image panels are in Figures 1–3 of the paper.

05 / Resources

Paper, implementation, and citation

ConicFlow: Constrained Inference-Time Optimal Control for Flow Matching Models. Shervin Khalafi, Kasra Fallah, Sung Cho, Dongsheng Ding, James Anderson, Alejandro Ribeiro. arXiv:2610.14307 (2026).

Read paper on arXiv ↗Project website source ↗

@misc{khalafi2026conicflow, title={ConicFlow: Constrained Inference-Time Optimal Control for Flow Matching Models}, author={Khalafi, Shervin and Fallah, Kasra and Cho, Sung and Ding, Dongsheng and Anderson, James and Ribeiro, Alejandro}, year={2026}, eprint={2610.14307}, archivePrefix={arXiv}, primaryClass={cs.AI}, url={https://arxiv.org/abs/2610.14307} }