Provably Safe Generative Sampling with Constricting Barrier Functions

TL;DR

This work proposes a safety filtering framework that acts as an online shield for any pre-trained flow-based generative model, and proves that this mechanism guarantees safe sampling while minimizing the distributional shift from the original model at each sampling step, as quantified by the KL divergence.

Abstract

Flow-based generative models, such as diffusion models and flow matching models, have achieved remarkable success in learning complex data distributions. However, a critical gap remains for their deployment in safety-critical domains: the lack of formal guarantees that generated samples will satisfy hard constraints. We address this by proposing a safety filtering framework that acts as an online shield for any pre-trained generative model. Our key insight is to cooperate with the generative process rather than override it. We define a constricting safety tube that is relaxed at the initial noise distribution and progressively tightens to the target safe set at the final data distribution, mirroring the coarse-to-fine structure of the generative process itself. By characterizing this tube via Control Barrier Functions (CBFs), we synthesize a feedback control input through a convex Quadratic Program (QP) at each sampling step. As the tube is loosest when noise is high and intervention is cheapest in terms of control energy, most constraint enforcement occurs when it least disrupts the model's learned structure. We prove that this mechanism guarantees safe sampling while minimizing the distributional shift from the original model at each sampling step, as quantified by the KL divergence. Our framework applies to any pre-trained flow-based generative scheme requiring no retraining or architectural modifications. We validate the approach across constrained image generation, physically-consistent trajectory sampling, and safe robotic manipulation policies, achieving 100% constraint satisfaction while preserving semantic fidelity.

Figures (7)

• Figure 1: Constrained generative sampling with a constricting safety tube. The flow-based sampling process transforms noise sample $x(T)$ into data $x(0)$. Unconstrained sampling process can produce unsafe samples that violate constraints, landing outside the required safe set. Our control barrier function-guided sampling augments the generative dynamics with feedback control inputs $u$ that maintain the sample within the constricting safety tube $\tilde{\mathcal{C}}(t)$ (green region) throughout the sampling process, guaranteeing safe samples and minimal perturbation of the learned sampling process.
• Figure 2: Lorenz system trajectory generation with CBF guidance. (a) Phase portrait showing the true ODE solution (dashed black). Our CBF-guided diffusion model (blue) closely tracks the true dynamics, while unconstrained sampling (red) produces physically inconsistent trajectories. (b) Evolution of the constricting barrier $\tilde{h}(x,t)$ (green) and relaxation $\epsilon(x(T),t)$ (blue) during sampling. The constriction goes from $\epsilon_0 \approx 16$ to 0, accommodating the large initial physics violation. The barrier remains non-negative throughout, verifying reverse invariance. (c) Control effort $\|u\|^2$ over sampling time, concentrated at the onset of sampling when the initial noise sample must be corrected, then rapidly diminishing as the trajectory becomes physics-consistent.
• Figure 3: Spatially constrained image generation with DDPM-bedroom-256. (a) Reference window patch ($50 \times 70$ pixels). (b) Generated image with 50 sampling steps. (c) Generated image with 200 sampling steps. Both (b) and (c) preserve the reference window exactly at the specified location while generating bedroom context. Higher sampling steps improve generation quality in unconstrained regions, but constraint satisfaction is guaranteed in both cases.
• Figure 4: Regional color intensity constraints. All images enforce different RGB values in the lower third. (a) Moderate intensity: The target pixel color is black, and the spatial mask $v(\mathbf{p})$ goes from $0.5 \to 0$, allowing the diffusion model to produce clean furniture. (b) Projection: Projection-based constraint enforcement SZ-etal:2025, constraint are same as in (a) with target pixel color black in the bottom one-third of the image. Projection at every sampling step causes intense constraint enforcement, which leads to loss of semantic information. (c) Low intensity: The target color is brown, and the spatial mask $v(\mathbf{p})$ goes from $0.2$ at the bottom to $0$ at the lower-third boundary, producing a semantically meaningful image of a room with a brown carpet.
• Figure 5: Smooth action generation for Push-T robotic manipulation. (a) The Push-T environment: a planar robot arm must push the T-shaped block (gray) to the target pose (green). (b) End-effector trajectories for a representative episode. Unconstrained Diffusion Policy (DP, dotted red) exhibits sharp directional changes during pushing. DP-DDIM (dashed blue) produces more erratic motion due to fewer sampling steps. Our CBF-guided sampling (solid black) generates a smooth trajectory that satisfies the acceleration constraint at every step while achieving the same task reward.

Theorems & Definitions (5)

• Definition 1: Constricting barrier function
• Theorem 4.1: Reverse invariance
• Remark 1: Comparison to other approaches to extending CBFs to stochastic systems
• Remark 2: Feasibility of the QP \ref{['eqn:qp-control']}
• Theorem 4.2: Distribution shift for guided sampling