Dual-Co-State Constrained Flow Sampler
Implementation & benchmark of arXiv:2607.04513 — Constrained Flow Matching via Lagrangian Dual Flows
Source paper: Constrained Flow Matching via Lagrangian Dual Flows arXiv:2607.04513 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Augment a flow-matching or diffusion sampler with a dual variable for each equality constraint and integrate the sample and dual variables as one coupled ODE. The learned generative velocity is corrected in the constraint-normal direction using the transpose Jacobian of the constraint, while the dual state accumulates residual violations; this replaces per-step projection or nonlinear optimization.
Formulas
Mathematical statement
The paper's construction is based on the classical primal-dual differential multiplier flow for equality-constrained optimization. Let $x(t)\in\mathbb{R}^d$ be the generated sample, let $g(x(t),t)\in\mathbb{R}^m$ be a differentiable equality constraint, let $J_g(x,t)=\partial g(x,t)/\partial x\in\mathbb{R}^{m\times d}$ be its Jacobian, and let $\lambda(t)\in\mathbb{R}^m$ be the Lagrange dual co-state. For a learned unconstrained velocity $v_\theta(x,t)$, use the coupled flow $\dot{x}=v_\theta(x,t)-\alpha(t)J_g(x,t)^\top\lambda$ and $\dot{\lambda}=\beta(t)g(x,t)$, where $\alpha(t),\beta(t)>0$ are gains. The first equation is the primal velocity corrected by the gradient of the Lagrangian term $\lambda^\top g$; the second integrates constraint violation and therefore increases the normal correction when $g(x,t)\neq0$. The defining advantage over projection is that no solve of $J_gJ_g^\top\lambda=g$ and no pseudoinverse is required. For inequality constraints $h(x,t)\leq0$, replace $g$ by the violation $[h]_+$ and optionally use a nonnegative dual state with $\dot{\lambda}=\beta[h]_+$; the MVP should first test equalities. If the constraint is intended only at the terminal time $T$, set $g(x,t)=w(t)g_T(x)$ with $w(t)$ increasing toward $T$, or use a terminal boundary residual $g_T(x(T))$ in evaluation. The mathematical hypothesis required by the implementation is that $g$ is differentiable and its Jacobian-vector products are computable; exact finite-time satisfaction is not automatic for arbitrary gains, so residual and sample-quality tradeoffs must be measured explicitly.
Implementation notes
Integrate this at inference time in the ODE solver used by a flow-matching model, rectified-flow model, or probability-flow diffusion sampler. The tensor being modified is the sample velocity, immediately after the network returns $v_\theta(x,t)$ and before the solver advances $x$. Given an initial sample $x_0$ and dual state $\lambda_0=0$, execute: (1) compute $v=v_\theta(x,t)$; (2) evaluate the constraint vector $c=g(x,t)$; (3) compute the Jacobian-transpose vector product $q=J_g(x,t)^\top\lambda$ using automatic differentiation, without explicitly materializing $J_g$; (4) return the augmented ODE derivative $dx=v-\alpha(t)q$, $d\lambda=\beta(t)c$; (5) integrate the concatenated state $(x,\lambda)$ with Euler, Heun, or Dormand-Prince. In pseudocode: state=(x,lam); def rhs(t,state): x,lam=state; v=model(x,t); c=g(x,t); q=vjp(lambda z: g(z,t), x, lam); return v-alpha(t)*q, beta(t)*c; state=solver.step(rhs,t,state). For terminal constraints, use $g(x,t)=w(t)g_T(x)$ with $w(t)=(t/T)^p$, and clip or softly normalize $\lambda$ if the correction becomes unstable. Compute from the mathematics only the Jacobian-transpose correction and dual accumulation; estimate $\alpha,\beta,p$ empirically on a held-out constraint-strength sweep. Start with a 2D or 8D toy flow and then CIFAR-10 flow matching with constraints such as fixed pixel mean, prescribed classifier score, or a differentiable image statistic. Compare unconstrained sampling, projection after every step, and penalty-gradient correction at equal numbers of network evaluations. Record terminal residual $\|g_T(x_T)\|$, FID or validation likelihood, solver function evaluations, and wall-clock time. Success is a 10--100x reduction in terminal constraint residual relative to the unconstrained model without projection, while matching projection-based sample quality and requiring no additional inner optimization loop.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built a reproducible 2D nonlinear equality-constrained flow sampler using the coupled primal-dual ODE, alongside unconstrained and per-step projection baselines. The transpose-Jacobian VJP check matched finite differences with absolute error 3.1e-10, and the linear sanity test reduced residual from 2.0 to 0.00125. On the toy flow, tuned dual gains (c=10, p=1) reduced mean terminal residual from 0.5213 to 0.0160 (32.6x) with endpoint distortion 0.609 versus 0.632 for projection, but stronger gains became unstable and default c=1 achieved only 4.1x reduction; this is a promising but tuning-sensitive signal, not a demonstrated general win.
- Agent confidence
- 8/10
- Baseline
- Unconstrained: mean |g_T|=0.5213; projection: mean |g_T|=1.0e-16, distortion from unconstrained endpoint=0.6317
- Idea
- Dual flow c=1,p=1: mean |g_T|=0.1264, 4.12x reduction; tuned c=10,p=1: mean |g_T|=0.0160, 32.55x reduction, distortion=0.6088; 300 Euler steps and no inner solve
Limitations: Only a small 2D synthetic flow and one circular equality constraint were tested. No learned neural flow, image-quality metric, wall-clock/VJP-vs-projection benchmark, adaptive ODE solver, multiple constraints, or terminal-only constraint schedule was evaluated; the tuned gain sweep also showed NaNs at c=30.
How to run: python3 dual_flow_experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Implemented a tabular neural-network analogue of the dual co-state correction with a tuned baseline sweep, 8 paired seeds, permutation test, and gain-based mechanism check. The idea was marginally better (MSE 12.3887 vs 12.4054; paired delta -0.0167), but won only 2/8 pairs and was not significant (p=0.7103), so there is no demonstrated benchmark win. Tabular was selected because this constraint/regularization intervention has no direct sampler track and tabular is the closest built-in domain.
- Baseline
- Tuned Adam lr=0.003, weight_decay=0.0001; 8-seed test MSE mean=12.4054099321, std=1.7629887167.
- Idea
- Dual correction alpha=0.1, beta=0.1 with the same optimizer and budget; 8-seed test MSE mean=12.3887301683, std=1.6727843838; paired delta=-0.0166797638, wins=2/8, permutation p=0.7103.
- Mechanism
- At seed 0, increasing alpha from 0.01 to 0.2 reduced training equality residual from 0.59031 to 0.44740, but this modest signal does not quantitatively establish transfer of the stage-1 sampler effect.
This is a training-time analogue because the fixed supervised benchmark has no generative sampler track; it does not test actual flow/diffusion inference, image quality, projection baselines, VJP wall-clock overhead, multiple constraints, or terminal-only sampler constraints.
How to run: python3 stage2_bench.py
Artifacts
- dual_flow_experiment.py 4.6 KB View Raw
- report.md 1.6 KB View
- report_bench_2026-08-30T120140.md 4.2 KB View
- results.json 6.2 KB View Raw
- stage2_bench.py 5.0 KB View Raw
- stage2_results.json 2.9 KB View Raw