Chance-Constrained Flow Optimization

Adapting pre-trained diffusion and flow models to optimize downstream utilities is central to their real-world deployment. Existing steering methods typically optimize average rewards or impose constraints on average properties. However, downstream requirements often concern the probability of generating samples that meet particular criteria, which cannot be represented by average rewards or constraints. To address this gap, we propose Chance-Constrained Flow Optimization (CCFO), a distributional steering framework that maximizes expected reward while explicitly controlling the probability of constraint satisfaction. For such a constrained optimization problem, we theoretically derive a closed-form expression for the optimal dual variable, which allows to directly characterize the optimal distribution. This characterization decomposes CCFO into a scalable two-stage procedure: a lightweight dual variable estimation and a single entropy-regularized fine-tuning step via a specific pseudo-reward. By decoupling dual estimation from model fine-tuning, CCFO achieves a computational cost comparable to that of standard expected-reward fine-tuning methods. Experiments on 2D illustrative examples, high-dimensional text-to-image generation, and protein sequence design establish that CCFO achieves better reward–feasibility trade-offs than state-of-the-art baseline methods.