Gauge Flow Matching: Efficient Constrained Generative Modeling over General Convex Set and Beyond
Authors: Xinpeng Li
Enming Liang
Minghua Chen
Venue: ICLR 2026 Poster (2026)
Keywords: constraint
generative models
homeomorphism
convex constraints
gauge mapping
Primary area: generative models
TL;DR: An efficient gauge-mapping based constrained generative model over general convex sets and beyond.

Abstract: Generative models, particularly diffusion and flow-matching approaches, have achieved remarkable success across diverse domains, including image synthesis and robotic planning. However, a fundamental challenge persists: ensuring generated samples strictly satisfy problem-specific constraints — a crucial requirement for physics-informed problems, safety-critical applications, watermark embedding, etc. Existing approaches, such as mirror maps and reflection methods, either have limited applicable constraint sets or introduce significant computational overhead. In this paper, we develop gauge flow matching (GFM), a simple yet efficient framework for constrained generative modeling. Our GFM approach introduces a novel bijective gauge mapping to transform generation over arbitrary compact convex sets into an equivalent process over the unit ball, which allows low-complexity feasibility-ensuring operations such as reflection or projection. The generated samples are then mapped back to the original domain for output. We prove that our GFM framework guarantees strict constraint satisfaction, with low generation complexity and bounded distribution approximation errors. We further extend our GFM framework to two non-convex settings, namely, star-convex and geodesic-convex sets. Extensive experiments demonstrate that GFM outperforms existing methods in both generation speed and quality across multiple benchmarks, including synthetic data, time series, and image generation.

Published as a conference paper at ICLR 2026
GAUGE FLOW MATCHING: EFFICIENT CONSTRAINED GENERATIVE MODELING OVER GENERAL CONVEX SET AND BEYOND
Xinpeng Li', Enming Liang*., Minghua Chen
1 Department of Data Science, City University of Hong Kong 2 School of Data Science, The Chinese University of Hong Kong, Shenzhen
ABSTRACT
Generative models, particularly diffusion and flow-matching approaches, have achieved remarkable success across diverse domains, including image synthesis and robotic planning. However, a fundamental challenge persists: ensuring gener- ated samples strictly satisfy problem-specific constraints - a crucial requirement for physics-informed problems, safety-critical applications. watermark embedding, etc. Existing approaches, such as mirror maps and reflection methods, either have limited applicable constraint sets or introduce significant computational overhead. In this paper, we develop gauge flow matching (GFM), simple yet efficient frame- work for constrained generative modeling. Our GFM approach introduces a novel bijective gauge mapping to transform generation over arbitrary compact convex sets into an equivalent process over the unit ball, which allows low-complexity Ceasibility-ensurii operations such as reflection projection. The generated samples are then mapped back to the original domain for output. We prove that our GFM framework guarantees strict constraint satisfaction, with low generation com- plexity and bounded distribution approximation errors. We further extend our GFM framework two non-convex settings, namely, star-convex and geodesic-convex sets. Extensive experiments demonstrate that GFM outperforms existing meth- ods in both generation speed and quality across multiple benchmarks, including synthetic data, time series, and image generation.
INTRODUCTION
Generative models have emerged powerful tools for learning complex data distributions, achieving remarkable success in diverse applications ranging from image generation to scientific simulation.
Recent advances, particularly diffusion models and flow-matching approaches, have further pushed
the boundaries of what's possible in areas such as photorealistic image synthesis, molecular design,
and robotic trajectory planning (Ramesh et 2022 Betker et al. 2023 Sun & Yang 2023 Chi
et al.
2023 Abramson et al. 2024] et al. 2025).
However, many real-world applications necessitate generation under specific constraints. For instance, protein synthesis requires adherence to structural constraints within amino acid chains. Image generation may demand precise watermark placement or consistency with physical laws. Robotic manipulation must respect joint limits and ensure obstacle avoidance. These constraints are not merely optional considerations but fundamental requirements of their respective problem domains. Generated samples must strictly satisfy these constraints to be both meaningful and practically useful within their intended applications.
Existing approaches to constrained generative modeling face significant limitations (see Table Dp. They either have limited applicable constraints (c.g., box and simplex) or lack a strict fcasibility guarantee for generated samples. To date, developing efficient framework for constrained generation with feasibility guarantees over general compact sets, convex or not, remains largely open. This work proposes Gauge Flow Matching (GFM), addressing these challenges with the following contributions:
"Coresponding authors: Enming Liang momming.cityu@ggmii... Minghua Chen (minguua wechtaaachhhkkeeee edu..cn

---

Published as a conference paper at ICLR 2026
In Sec. 4 we propose bi-Lipschitz bijective gauge mapping, generalized from the one in (Tabas Zhang 2022a), to transform generation compact convex sets to an equivalent process over unit ball, which allows low-complexity feasibility- ensuring operations such as reflection or projection.
dx//dt 12
general
X
X1
FM
dz,/de 92
we prove that the GFM framework guarantees strict constraint satisfaction and bounded distribution approximation error while incurring sig- nificantly lower computational complexity compared to SOTA constrained generative models such as reg- ular reflection projection- -assen methods.
In Sec. 5
Zo
Z1
B
GFM
Figure 1: Gauge flow matching framework.
In Sec. 6 we further extend GFM to certain non-conver constraints, in particular star-convex and geodesic-convex sets, while inheriting the computational merits in the convex setting.
Sec. 7 we provide extensive empirical studies to demonstrate our framework's efficiency in terms of feasibility, approximation capability, and inference complexity compared to SOTA constrained generative models.
2 RELATED WORK
Ensuring feasibility is fundamental constrained optimization and neural network design. While extensive work exists constrained optimization algorithms and NN output feasibility methods (see Appendix A), constrained generative modeling presents distinct challenges. Traditional generative models like VAES and GANs can directly apply existing NN feasibility techniques. However, diffusion and flow-based models require specialized approaches because they generate samples through forward integration with NN-approximated score functions or vector fields rather than direct NN outputs. Table summarizes these specialized approaches.
Reflected Processes leverage reflection mechanisms to constrain trajectories within feasible regions. Methods differ in training: RDM" uses implicit score matching (Fishman et al. 2023), RDM employs approximated denoising score matching (Lou Ermon 2023), and RSB utilizes iterative proportional fitting (Deng et al., 2024). RFM extends this to flow-based generation over convex sets (Xie et al. 2024). While effective, reflection-basse based methods incur substantial computational overhead from boundary localization and reflection calculations during forward integration. Recent work uses Metropolis sampling to reduce this burden (Fishman et al. 2024), though without strict fcasibility guarantees. Additionally, these methods require the initial distribution's 's support match the target set C, introducing complexity since sampling from general convex sets is non-trivial.
Guided Generation incorporates auxiliary terms to enforce constraints during generation. Q-Bridge employs Doob's 's h-transffrrm with time-dependent force terms (Liu & Wu 2023), while log-barrier diffusion models enforce feasibility through barrier functions (Fishman al. 2023). Direct projection
onto the constraint set or incorporating feasible directions during forward integration for intermediate noisy samples provides another approach (Chung al. 2022 Christopher et al. 2024 Zampini et al.) 2023 Naderiparizi et al. 2025 Liang et 2025 Cardei et al. 2025). Rather than applying these operations intermediate noisy samples, several methods apply them the estimated target samples, followed by resampling or interpolation to update subsequent noisy samples (Song et al., 2203a Cheng et al. (2244) Zhang et al. 2025 Utkarsh et al., 2025 Graikos et al. 2025] Narasimhan et al. 2025 While these methods can improve feasibility, exact projection is computationally expensive, and approximated methods lack error analysis or feasibility guarantees.
methods incorporate constraint violations as penalty terms in training objec- tives (Li et al. 2024) or use Lagrangian- based training with dual variable updates (Khalafi et al.) 2024). Post-training fine-tuning with reward functions provides another pathway to improve constraint
Training/Fine-tunn
satisfaction (Fan & Lee 2023 (ehara et al. 2024 Domingo-Enne Enrich et al. 2024).
Bijective Maps transform complex constrained domains into unconstrained spaces simple sets.
RDM maps simplexes to unit cubes for scalable score matching (Lou & Ermon 2023). MDM uses closed-form mirror maps for simple convex sets such balls and simplexes (Liu al. 2024b
2

---

Published as a conference paper at ICLR 2026
Table 1: Existing study on constrained diffussion/low-matciiie models over continuous domain.
Constraint
Initial support set
Feasibility
Wasserstein bound
Low generation
Method (ref. in Sec.
set C
guarantee (2)
complexity 8)
DM/FM
R"
O(NFE n2²)
RDM
Convex
C
RSB
Smooth
C
RFM
Convex
C
MDM
Ball/Sim./Poly.
R
MFM
Convex
R"
NAMM
(Non)-Convex
R"
Barrier-based
Convex
Projection-
(Non)-Convex
R
Penalty-bused
General
R"
Gauge Flow Matching
Convex
B
The support set of initial/prior distribution matters when preparing the initial samples for training and generation, since sampling from the general convex set, even following simple uniform distribution, is computationally expensive (Kook & Vempala, 2024).
2 The low generation complexity of constrained generative models indicates the complexity is matched with regular DM/FM models, dominated by neural network evaluation (O(n )) and scaled linearly with the number of function evaluations (NFE).
and MFM refines MDM using modified Lipschitz mirror map (Guan et al. 2025 However, the inverse the mirror map for general convex domain does not admit a closed-fome expression. The mapping between convex polytopes and the unit ball or unconstrained space, as applied in Diederen & Zamboni (2025), enables efficient flow matching in the transformed domain. NAMM generalizes this approach using neural networks to approximate invertible maps (Feng al.. 2024), though it lacks theoretical guarantces for feasibility and distribution approximation. Bijective mappings have also been applied continuous embeddings of discrete categorical data (Davis al. 2024 Williams et al. 2025). Nevertheless, constructing explicit-fom bijective mappings with efficient computation for complex constraints, even general convex ones, remains non-trive non-trivial.
In summary, existing approaches either have limited applicability or lack performance guarantees. We propose gauge mapping-hase approach that, while conceptually related to mirror map methods, offers broader applicability, theoretical guarantees, and improved computational efficiency.
3 PROBLEM STATEMENT
We consider flow matching-based generative modeling for data distribution Pdata over a general compact convex [ set C " The vanilla flow-matching model (Lipman al., 2022 Liu al. 2022b) is trained by matching the designed conditional flow (e.g., linear flow) as:
min ((vg) = Ex.. [10 t) - (x1 =o>1l]
(1)
where = (1 t) + tz with Zi ~ Po. matching loss ( yields vector field that transforms simple initial distribution, e-g., Gaussian Po = (0, T), into the target data distribution P1 Pdata (Liu et al. 2022b). In practice, the vector field is parameterized by a neural network Ua and optimized using samples from the target distribution according to (T. Sample generation is achieved through forward integration = To So vo(#,, l)dl, initializingg from a Gaussian sample xo and following the learned vector field vg.
Pl. and ~ 24((0, 1)). The minimizer of the flow
However, the generated samples often exhibit deviations from the constraint set C due to a phenomenon known as error propagation (Li al. 2023c). This occurs when the approximation errors of NN-based vector fields accumulate throughout the discretized integration process, ulti- mately resulting in significant deviation of the generated samples from the constraint sets. Existing
Open issue:
Compact convex set includes linear-equalijy and convex-inequality constraints. In this work, we consider the convex- convex-inequaliy inequality without loss of generality. For linear- equarrey,ullty, it can be embedded in an unconstrained subspace by selecting independent variahles and reconstructing the dependent variables via closed-form equality
solving (Donti et al. 2020 Liang et al. 2023 Ding et al. 2023), see Appendix B for details.
3

---

Published as a conference paper at ICLR 2026
approaches addressing feasibility issues suffer from either limited applicability or high computational complexity (see Table I).
GAUGE FLOW MATCHING OVER CONVEX SETS
To address these limitations of existing methods and enable efficient constrained generative modeling. we introduce our GFM framework. It employs gauge mapping an explicit bijective mapping between two convex sets-to transform complex constrained generative modeling into an equivalent modeling over simple unit ball. The framework (i) builds the flow-matching model for transformed data distribution over a unit ball through inverse gauge mapping: and (ii) generates samples over a unit ball via a closed-form reflection and transforms them back to the original space through forward gauge mapping.
4.1 GENERALIZED GAUGE MAPPING BETWEEN CONVEX SETS
We first introduce bijective mapping between two compact convex sets in Euclidean space, known
as gauge mapping (Tabas & Zhang 2022a b Liang & Chen 2025 Li et al.
2025b] Liu et al. 2025):
Definition 4.1 (Gauge mapping). Let (cla, x
E A(C x²)) be the Gauge/Minkossi funtcciti function (Blanchini & Miani 2008) given an interior point x" E int(C). The gauge
inf (A
0
mapping B C can be defined between a unit p-norm ball and compact convex set:
||=1|e
7C x², xx))
Φ (2) (0(2, x²)
+ x", V2 E B,
(x)
(x - x x²), Vx EC,
(2)
x
x. llp
As shown in Fig. 2 the gauge mapping (-)) es- tablishes a continuous bijective correspondence (i.e., homeomorphismn) between any compact convex set and a unit p-norm ball: = 2(B) and B- = -- (C). Intuitively, the gauge map- ping transforms a unit ball into convex set by first translating the unit ball to the interior point (0 > 2)), then scaling along every radial direction from this interior point such that the boundary of the ball becomes aligned with the boundary of the convex set (e.g., 21 + 11). As a result, all level sets of B are mapped to level sets of (e.g., 22 + x22). The inverse mapping is similarly constructed by inverse scaling and translation back to the origin.
= ((x)
21
x; 4(2))
= (0)
0
B
C
Figure 2: Gauge mapping illustration.
Remark (Literature of Gauge function/map). Gauge/Minkowski functions YC have been exten- sively studied in the control and online learning communities for Lyapunov function design and efficient online optimization (Blanchini] 1995 Raković & Lazar 2014;; Mhammedij 2022 Lu et al., 2023). While gauge mapping & was proposed in (Tabas & Zhang 2022a), specifically between polytope and cube, we generalize it arbitrary pairs of compact convex sets and provide efficient computation methods. Specifically, in the gauge mapping calculation, (i) gauge function xe (xx xx) has closed-form expressions for common convex sets (e.g., linear, quadratic, and convex cones) and can be efficiently computed via bisection methods for general convex constraints; (ii) the interior point x: is obtained once for all by solving a convex feasibility problem offline (see details in Appendix (B)
Further, we characterize the bi-Lipschitz property of gauge mapping, as will be demonstrated in Sec. 5, which is essential for establishing the regularity of transformed data distribution and controlling the distribution approximation error in our GFM framework.
Proposition 4.1 (Bi- -Bi-hipsc Property of the Gauge Mapping). Let C C. R be a compact convex set and let x E iC) be an interior point. Define the inner and outer radii with respect to x² as
ri := (r (r 2 0: B2(x=, (e(x",r) CC,
To := in 2 0 : C C B2(x") (a(x,,r))
such that B2(x",r,) cCC B2(x", ro). Then the gauge mapping between C and 2-norm ball B2 satisfies the following bounds:
Forward Lipschitz:
T < 2r. + +///
Inverse Lipschitz:
La-1 2/n
(3)
4

---

Published as a conference paper at ICLR 2026
The bi-Lipschitz property of gauge mapping depends on interior point selection (Liu al. 2025 When a near- boundary interior point selected (i.e., ri - 0), the bi-Lipschitz constant approaches infinity, which severely distorts the data distribution and creates significant challenges for generative modeling. In practice, we aim to identify a "central" interior point (where ro is close to ri) to minimize the bi-Lipschitz constant. Such interior point be approximately determined through constraint residual minimization (Tordesillas et al., 2023), which involves linear optimization over the target convex set and can be solved in polynomial time. While we derive explicit Lipschitz bounds for gauge mapping from a 2-norm ball to convex set, these bi-Lipschitz bounds can be extended to arbitrary pairs of compact convex sets. For a more detailed introduction gauge mapping and its properties, we refer readers to Appendix B
Algorithm Training in GFM
Algorithm 2 Generation in GFM
Input: Data samples: æ1 Pdata over compact convex set and gauge mapping between C and B.
Input: Prior samples 2o LB) NN vector field va. and gauge mapping between C' and B.
I:
Inverse transformation: 21 Φ (x1)
1:
ODE solver with reflection over B in Eq. Forward transformation: 21 (x1)
Regular flow matching over in Eq.
2:
Output: Trained NN vector field
Output: Generated feasible samples 271.
4.2 TRAINING PHASE OF GFM
Given the gauge mapping & between the convex set C and unit ball B, a flow matching model is trained on the transformed space as:
min ((ve) =
Ezo... =
[110(2e,))) (21 - 20)1
(4)
where the initial samples 2o Yo are sampled simple prior distribution to (e.g., uniform) supported on unit ball B, the terminal samples 21 q1 are transformed from the samples from the original data distribution via inverse gauge mapping as 21 = 4b- - (z1). and 2 (1 t)2o + t21. In essence, GFM models the transformed data distribution as Qdata & # Pdata over unit ball B, where # is the push-forward operator. We then leverage the regular flow matching training approach to train a neural network vector field ve following (4)
Remark 2 Comparison to Mirror Maps). Mirror map-based generative models also employ bijective mapping to transform constrained distributions to the unconstrained dual space (Liu et al.) 2024bb. However, it admits only flosed-form computation for simple convex sets (e.g.. balls and sim- plices), and it maps boundaoy samples to infinity in the dual spaces, challenging the transformed generative modeling both theoretically and practically. In contrast, gauge mapping is computationally efficient for any compact convex set and maintains bounded Lipschitz constants for both sides (Prop. 4.1). These properties crucial for the regularity of transformed data distribution and efficient flow-mtchhing training, as will be analyzed in Sec.
4.3 INFERENCE PHASE OF GFM
After training, we generate samples within the unit ball B following the NN vector field uo. To constrain the generation trajectory within B, we apply an additional reflection term (Xie et al. 2024):
20 L.'' (Vo(EL, t)dt + dL.),
21
where zo is sampled from unit ball following a prior distribution (e.g., uniform), and L, is the reflection term when 2; hits the constraint boundary (Xie et al.) 2024). Finally, we recover the sample to the original space following the forward gauge mapping as I1 ((1)).
Remark 3 Comparison to Regular Reflection). Reflection- mechanisms are well-established in previous works to keep the generated samples within the constraint set (Lou Ermon 2023 Fishman et al. 2023 Deng et al. 2024 Xie et al. 2024). However, they face several limitations: (i) the reflection term is computationally expensive beyond simple sets (e.g., ball and simplex); (fi) existing reflection generative models require a prior distribution within the target constraint set C, which is challenging to sample from (even for uniform distribution) over general convex sets during training
generation (Kook & Vempala 2024). These computational issues prevent reflection methods
5

---

Published as a conference paper at ICLR 2026
from being applied more complex sets. In contrast, after transforming the data distribution over a unit ball through inverse gauge mapping, we can easily sample from a unit ball (e.g.. uniformly), implement closed-form reflection term with (nn) algorithmic complexity, and batch computation for multiple samples, thus ensuring efficient sample generation within the ball and strict feasibility after mapping the original data space.
PERFORMANCE ANALYSIS OF GFM
Proposition 5.1 (Regularity of Transformed Data Distribution). Assume that the original data distribution Pdata satisfies regularity conditions (Assumption I in (Wan et al. 2024(). The transformed data distribution qdata = $ Pasta by bi-Lipschitz homeomorphism (e.g., gauge mapping) also satisfies those regularity conditions.
We first assume that the original data distribution Pdata satisfies regularity conditions following (Wan cct al., 2024), ensuring well well-posedness of standard flow matching over the data distribution. Veritying the regularity of the transformed distribution qdata 4 Pdata is essential for establishing the existence and well-posednese of our gauge flow matching model. Under the bi-Lipschitz properties of gauge mapping established in Prop. 4.1 we can verify that these regularity conditions are preserved. It is worth noting that theoretical analyses of flow matching models in the literature assume various regularity conditions for data distributions (Benton et al. 2023 Gao et al.] 2024a b). We verify a general condition from recent work (Wan et al. (024). Detailed proofs and further discussion
other common regularity conditions are provided in Appendix D
Proposition 5.2 (Wasserstein Bound of GFM). Let the NN approximation error be ej Pt and Pt is the probubility density at time 1 driven by the target vector field u. Assume that vg is Li-Lisschitz for z E B and t E [0, Denote the induced probability distribution p under with reflected generation, the Wasserstein-2 dis- tance between the data distribution Pdata and the approximated distribution pe is bounded by W2(piaaaa 1og")) S Leg2/22Leen
B,,,,, ([00 uu(z, 1)13,, where
The Wasserstein error of GFM is bounded by the Lipschitz of gauge mapping multiplied by the distribution error in the transformed space under reflected generation. To reduce the distribution approximation error, we can regularize the neural network Lipschitz constant La or optimize the loss funetion such that eg is minimized, following standard flow matching training pipelines (Liu et al.) 2022b). Specific to our GFM framework, we can further optimize the Lipschitz constant of the gauge mapping Lap As discussed following Prop. 4.1 we can select "central" interior point for the gauge mapping by solving convex optimization problems, thereby reducing the bi- Lipsschitz constants.
Proposition 5.3 (Inference Complexity for GFM). Consider a compact convex set C TR" defined by constraints g((x) < O, for i 1,2,.... m. The generation complexity of GFM (forward integration
and gauge mapping calculation) is O(NFE n² + m C), where C. = maax {Ci} is for gauge
15i5
function calculation and varies by constraint type:
(i) For linear constraints gi (z) aT S C. =; O(n): (ii For quadratic constraints
(x(x) "Qrta" x-b D 0, C, = O(n)); (iii) For second-order cone constraints (x()) = ||A A x-
pll2 (a x bx+bb) S O, C, ((k)); (iv) For matrix cone constraints ge(x)) Σ}= I.j F2+ + F. ł
C. = Omk22 + k; where E R", b R, Q S4x", BAxxm,AC AE Bunk p E R and F3 E k
For general convex function g: (z), C, Occ log (ot) using bisection, where c; is the complexity to evaluate g:(-) given a point and Ebis is the ervor tolerance in bisection.
The forward integration complexity of our model aligns with regular flow matching approaches, requiring NFE (Number of Function Evaluations) multiplied by the evaluation complexity of og. The additional reflection over a unit ball incurs negligible overhead with (n) complexity compared to the NN forward calculation (O(n")). For gauge mapping computation, given that the interior point is pre-computed offline, (i) for common convex sets, it can be solved efficiently (Fig. 8). and (ii) in general case, the bisection algorithm achieves linear convergence with minimal per-iteration complexity, merely requiring calculating the constraint function without solving any optimization problem (Mhammedi 2022).

---

Published as a conference paper at ICLR 2026
6 GAUGE FLOW MATCHING BEYOND CONVEX SETS
We extend the gauge mapping principle to two important classes of non-convex constraint sets, broadening GFM's applicability to richer geometrie settings. We include formal definitions in Appendix C and provide empirical studies in Sec.
Star-convex sets are bounded regions where the entire boundary is visible from designated interior point (c.g., fo.. noom ball). These sets arise in robotic planning and ftance-constrainee optimization (Charnes & Cooper 1959 Hansen et al. 2020 Liu al. 2022a). The key insight is that compact star-convex sets preserve the essential properties needed for gauge function construction both the gauge function and bijective mapping to the unit ball extend naturally from convex case (Licht 2024).
>> Geodesic-conver sets are regions on Riemannian manifolds where geodesics connecting any two points are unique and remain entirely within the set. These constraints widely arise in geometric learning applications (Chen & Lipman 2023; Miller et al., 2024; Zaghen et al. (025). The key insight is that connexic- ensures the exponential map at any interior point provides local diffeomorphism (Lee) 2006), allowing gauge mapping construction in the tangent space.
7 EMPIRICAL STUDY
We conduct extensive simulations across illustrative examples, robotics benchmarks, and high- dimensional constrained sampling tasks to demonstrate the effectiveness of our GFM framework. Detailed experimental settings and data descriptions are provided in Appendix E
Baselines: We compare against the following constrained generative models, selected for their applicability to diverse constraint types: (i) FM: Vanilla flow matching from Gaussian to target
distribution with linear conditional flow (Lipman et al., 2022 Liu et al. 2022b); (ii) DM: Vanilla diffusion model with variance preserving diffusion process (Ho et al. |020] Song et al. 2020); (iii)
Reflection: Reflection term applied when generated samples the constraint boundary (Xie) et al. 2024]; (iv) Metropolis: Metropolis sampling for approximating reflection-based generation (Fishman et al., 2024); (v) Projection: Orthogonal projection are applied when generated samples
violate constraints (Christopher et al. 2024); (vi) GFM: Our framework in Sec.
Metrics: We evaluate those baselines based on (i) constraint satisfaction (i.e., feasibility) ratio (%) of 10,000 generated samples, (ii) distribution approximation error, which is measured by Maximum Mean Discrepancy (MMD) between data samples and generated samples (Fishman et al. 2023). (iii) average per-epoch training time, including prior sampling, data transformation (in GFM only), and NN training, and (iv) inference time for generating 1,000 batched samples (unless otherwise specified, including prior sampling, forward integration, and data transformation (in GFM only).
7.1 SYNTHETIC EXAMPLES: CONVEX, STAR-CONVEX, AND GEODESIC-CONVES SET
We first evaluate GFM's performance in convex, star-convex, and geodesic-conver domains. As shown in Figures 3 we observe: vanilla diffusiion/llow-aaaee matchiinn models fail to guarantee sample feasibility, especially when probability density concentrates near boundaries. Existing reflection, projection, and Metropolis methods face limitations with complex constraints due to: (i) expensive prior distribution sampling, increasing both training and inference time; (ii) constraint-speciffc implementation requirements incurring longer inference time and limited constraint settings.
Our GFM framework transforms generative modeling over diverse convex into simpler process over a unit ball, enabling efficient prior sampling and computational efficiency comparable vanilla FM. After generation within the unit ball, samples mapped back to the target constraint set via low-complexity gauge mappings, ensuring strict feasibility. Further, due to the bi-Lipschitz property (Prop. 4.11 of the gauge mapping, the approximation error of GFM remains comparable to vanilla FM. This property ensures that the distortion introduced by our transformation is bounded and controlled, preserving the fidelity of the generated samples while guaranteeing their feasibility within the target constraint set.
reflection calculations with
7.2 CONSTRAINED CONFIGURATIONAS MODELING OF ROBOTIC ARMS
We apply GFM to robotic arm control tasks following (Fishman et al. 2023 2024). This involves learning distributions over joint locations and manipulability ellipsoids represented by symmetric
7

---

Published as a conference paper at ICLR 2026
Fcasibility
MMDJ (x10
Training
Inference
(%)
(8)
(8)
DDPM
95.0
4.79
0.17
0.59
(a) FM
(b) DDPM
(c) Refection
FM
95.9
8.57
0.18
0.29
Refection
100
25.9
6.40
14.0
Metropolis
100
130
6.40
6.12
Projection
7.12
GFM
100
3.50
0.18
0.63
(d) Metropolis
(e) Projection
(f) GFM
Reflection, Metropolis, and Projection models chare the same NN. velocity model, and thus have the same training time per opoch.
Figure 3: Performance over joint linear and quadratic convex sets
Feasibility (%)
MMDJ
Training Inference
(x10*
(s)
(8)
DDPM
91.1
6.03
0.20
0.70
FM
93.4
4.92
0,22
0.39
(a) DDPM
(b) FM
(c) Reflection
Reflection
100
7.96
5.58
12.4
Metropolis
100
24.6
5.58
8.72
Projection
100
5,58
GFM
100
5.01
0.22
0.74
Sampling from distribution time-consuming for reflection and Metropolis methods.
domain
(d) Metropolis (e) Projection
(f) GFM
Projection onto non convex incurs significant computational
cost.
Figure 4: Performance over star-convex sets
Feasibility (%)
MMDJ
Training Inference
(x10-3
(s)
DDPM
93.4
2.84
0.22
0.73
(a) DDPM
FM
FM
96.1
1.20
0.21
0.35
Metropolis
100
48.3
0.37
3.36
GFM
100
1.29
0.23
1.44
Sampling prior distributions over constrained maniffllss computationally expensive, resulting long Metropolis inference times.
Metropolis
(d) GFM
High rejection rates further increase Metrypolis inference time. Reflection and projection methods for constrained manifolds are not implemented computationl
Figure 5: Performance over geodesic-convex sets
positive-defintte matrices with trace constraints. Figure TT (Appendix E.6 shows generated velocity manipulation ellipsoids and trajectories. GFM successfully models this distribution while maintaining constraint satisfaction.
7.3 CONSTRAINED TIME SERIES GENERATION FOR TRAFFIC DATA
Time series prediction is crucial for engineering applications, yet real-world data often suffers from imbalance and scarcity due sensor failures privacy concerns. Synthetic data provides promising solution for augmenting limited datasets (arasimhan et al., 2025). We evaluate GFM generating realistic traffic sequences from REMS-BAYS eLi al. 2018) subject to physical (e.g.. capacity bounds) and statistical (e.g., average volume) constraints that define convex second-order cone, ensuring physically plausible and statistically consistent sequences. See details in Appendix F.8
Table d shows that GFM achieves 100% constraint satisfaction versus 88.5% for vanilla flow matching, demonstrating the necessity of constraint-aware generation. Projection baselines also achieve 100% feasibility but increase inference time from 0.3 to 49.7s (160x slower), while GFM requires only 0.43s. Critically, GFM preserves superior distribution quality with a KS statistic of 0.35 (p-value 0.42) compared projection methods. Figure 6 GFM captures characteristic traffic patterns periodic fluctuations and sharp drops-whin while strictly satisfying constraints, achieving optimal balance between feasibility, fidelity, and efficiency for practical traffic data augmentation.
8

---

Published as a conference paper at ICLR 2026
Part of PEMS-MSSSEE BAY time series
Metric
FM
Proj.
GFM
Feasibility (%)
88.5
100
100
2017 01 13
2017 01 16
2017 01
2017 01
KS +
0.30
0.37
0.35
p-value 1
0.43
0.20
0.42
Training (8)
0.07
0.08
0.09
Inference (s)
0.31
49.7
0.43
The Seonggoroo--Siiiroo (KS) statistic messures the maxi- distance between empirical cumulative distribution functions. KS value indicates identical distribations. The p-value
probability of observing at least as large as the computed value under the identical distributions hypothesis. Larger p-values suggest stronger evidence of distributional similarity. Inference time measured single generation with batch size 600.
statistic
V
2017-61-11
2017-01-28
2017-02-27
201 7-42-06
Figure 6: Performance over constrained conditional time series generations.
(a) FM
(b) Projection
(c) Reflection
(d) GFM
FID (50K)
Feasibility (%)
Epoch training
5K images inference
Final transformation
Methods
(8) 4
FM
3.57
76.66
76.7
155.2
Projection
6.88
100
100.4
556.6
Reflection 6.06
100
101.3
183.9
GFM
5.85
100
81.2
167.6
0.005
We select pixels watermarks constrained polytope. The selected pixels in the original datasct projected polytope, while other public pixels remain unchanged in the watemmarked. Minor map implementation roquires an orthogonal basis polytope Liu et 2024b) and excluded randomly sampled polytopes with general coefficients.
Figure 7: Constrained Image Generation with Embedded Watermarks
7.4 CONSTRAINED IMAGE GENERATION WITH EMBEDDED WATERMARKS
Watermarking generative models is critical for provenance tracking and copyright protection in Al Al-generated content. Following (Liu al. 2024b), we partition pixels [0. 1 into unconstrained public pixels X₁ and watermark pixels X2 E C, where C. is a random polytope for watermarking. Baselines are distilled from publicly available checkpoint checkpoolcheckpointto of flow matching models (Salimans & Ho 2022 Tong et al., 2024), sharing the same U-Net architecture (34M), distilling hyperparameters (100 epochs), and sampling settings (100-step Euler method) on CIFAR-10. See details in Appendix E.9
Fig. shows that GFM maintains 100% constraint satisfaction, achieves better FID than other constrained generative models (projection or reflection), while having similar trainingfifferncee complexity vanilla methods. This demonstrates that GFM enables efficient, high nigh-quality watermark embedding without compromising visual fidelity or constraint guarantees.
7.5 SOLUTION GENERATION FOR (RELAXED) COMBINATORIAL PROBLEMS
We further evaluate our approach on -ige--imensional (n =10, 1, 000) solution generation for relaxed combinatorial problems (Kook Vempala| 2024). The target distribution follows a log-concave density and is constrained by positive semidefinite cone and set of linear inequalities. This target distribution encapsulates several important classes of semidefinite relaxations for classical combinatorial optimization problems (e.g, max-cut and minimum- minimum-volume covering problem). We prepare the dataset following this distribution described (Kook & Vempala 2024) and train our GFM model. As shown in Table 7 our method achieves 100% feasibility rate while vanilla models fail to satisfy the constraints. GFM also maintains generation quality comparable to standard flow- matching models. Regarding inference time, the gauge mapping transformation has explicit-form

---

Published as a conference paper at ICLR 2026
Table Solution sampling for relaxed combinatorial problems.
10
10
50
50
100
100
Metric
DDPM
FM
GFM
DDPM
FM
GFM
DDPM
FM
GFM
Feasibility (%) f
47.2
0
100
0
100
0
100
MMD (x 10 4)4
5.35
9.71
9.70
43.4
43.4
43.5
85.8
85.8
85.8
Training (s)
0.26
0.26
0.27
0.30
0.31
0.53
0.69
0.72
0.87
Inference (8)
0.97
0.65
1.17
2.01
1.66
11.78
5.93
5.25
29.12
computation over PSD constraints, with computational cost primarily stemming from the final-step matrix calculation (e.g., torch. linalg. eigvalsh) for the gauge function for 1, C00 samples batch. This overhead can be further reduced through the application of advanced linear algebra packages (Van de Geijn 1997). We do not implement reflection methods since sampling the prior distribution in such high-dimensional constraint set is computationally expensive (e.g., if we apply Ball walk, it will incur On log n) membership//valuuatee For the projection methods, it will solve 10s projection problems over the PSD cone at most (1, 000 samples and 100 integration steps), which is unaffordable
complexity (Kook & Vempala 2024
computing device.
7.6 SCALABILITY TESTS AND ABLATION STUDY
Scalability of gauge mapping: We evaluate GFM's scalability in high- high-dimensional settings by measuring the computational cost of gauge function calculations. For closed-form gauge calculation on four common convex sets (Fig. 8), our approach maintains efficiency across di- mensions up to 3, 000, demonstrating practical
Linear
Quadenic
LMI
10'
Ganpe Computriion
100
applicability for high-dimensional tasks. For bisection- gauge calculation on polyno- mial constraints (sum-of-squares formulation, Table 12), a degree-4 polynomial with 50 vari- ables containing 1, 758, ,, 276 monomial terms can efficiently compute gauge functions for 000 samples in 0. 49492 seconds. Details are provided in Appendix 10
10"
10
10"
10'
16"
10"
10'
10'
10
Comoraint Dimmmiii
No
Comattrnnt
Figure 8: Gauge function computation.
Impacts of the interior point (IP) selection: Following Prop. 4.1 the IP selection determines the
bi-Lipschitz property of gauge mapping and consequently affects generation quality. We compare central versus saar-boundary IPs demonstrate this claim, with results in Table [3]
Comparison to Mirror Map: We compare the Mirror Map and Gauge Map in terms of generative modeling over the simplex and 50 dimensions, with results presented in Appendix E.5
Impacts of generation strategies: While we adopt reflection to keep samples within the unit ball (Sec. 4.3) alternatives like projection can also be efficiently implemented. Table 13 shows that projection is faster but leads to larger MMD, which motivates theoretical analysis in future works.
CONCLUDING AND LIMITATIONS
We introduced Gauge Flow Matching (GFM), a framework that transforms generative modeling over arbitrary compact convex sets into equivalent process over the unit ball through generalized gauge mapping. This approach guarantees strict constraint satisfaction with low computational complexity and bounded distribution approximation errors, and can be extended to important non-convex settings, Our experiments demonstrate that GFM outperforms existing methods in both speed and sample quality across multiple benchmarks. Despite these advances, several limitations warrant future investigation. (i) Extending the framework to more general non-convex sets presents an important direction, which could realized through decomposition the non-convex set into disjoint convex or star-convex subsets, followed by conditional generative modeling over those subsets. (ii) While our primary focus on continuous domains, extending GFM to discrete generation through relaxation embedding techniques represents promising avenue for future work (Davis et al. 2024). (iii) Adapting GFM to constrained one-step generation models could further improve generation efficiency beyond the current multi-step approach (Song al. 2023b
Frans et al., 2024
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
This work is supported in part by a General Research Fund from Research Grants Council, Hong Kong (Project No. 11200124), a Collaborative Research Fund from Research Grants Council, Hong Kong (Project No. C1049-24G), an InnoHK initiative, The Government of the HKSAR, Laboratory Al-Powered Financial Technologies, a Shenzhen-Hong Kong-Macau Science Technology Project (Category C, Project No. SGDX202205301 1203026), and Start-up Research Grant from The Chinese University of Hong Kong. Shenzhen (Project No. UDF01004086), The authors would also like thank the anonymous reviewers for their helpful comments. This work has previously been presented at the ICLR 2025 DeLTa workshop (Li et al. 2025b).
CODE AVAILABILITY
Code is available at Github
LLM USAGE
Large Language Models (LLMs) were used to aid in the writing and polishing of the manuscript.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneherger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):4935000) 2024.
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances neural information processing systems, 32, 2019.
Brandon Amos and Zico Kolter. Optnet: Differentiable optimization as layer neural networks. In International Conference on Machine Learning, pp. 136-145. PMLR, 2017.
Claude JP Bélisle, Edwin Romeijn, and Robert I. Smith. Hit-and-run algorithms for generating multivariate distributions. Mathematics of Operations Research, 18(2):255-26) 1993.
Joe Benton, George Deligiannidis. and Amnad Doucet. Error bounds for flow matching methods. arXiv preprint arXiv:2305. 2.0.::23055666// /6860, 2023.
James Betker, Gabriel Goh, Jing. Tim Brooks, Jianfeng Wang, Linjie Li, Long Ouyang, Juntang Zhuang, Joyce Lee, Yufei Guo, et al. Improving image generation with better captions. Computer Science. hups://cdn. openai. com/papers/dall-e-3, pdf. 2(3):8, 2023.
Franco Blanchini. Nonquadratic lyapunov functions for robust control. Automatica, 31(3):451-461 1995.
Franco Blanchini and Stefano Miani. Ser-theoretic methods in control, volume 78. Springer, 2008.
Michael Cardei, Jacob K Christopher, Thomas Hartvigsen, Bhavya Kailkhura, and Ferdinando Fioretto. Constrained discrete diffusion. arXiv preprint arXiv:2503.0979n 2025.
Abraham Charnes and William W Cooper. Chance-constrainedd programming. Management science, 6(1):73-79, 1959.
Bingqing Chen, Priya L Donti, Kyri Baker, Zico Kolter, and Mario Bergés. Enforcing policy feasibility constraints through differentiable projection for energy optimization. In Proceedings of the Twelfth ACM International Conference on Future Energy Systems, pp. 199-210, 2021.
Ricky TQ Chen and Yaron Lipman. Flow matching on general geometries. arXiv preprint arXiv:2302. 03660, 2023.
11

---

Published as a conference paper at ICLR 2026
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang, and Wotao Yin. Learning to optimize: A primer and a benchmark. Journal of Machine Learning Research, 23(189):1-533 2022.
Chaoran Cheng, Boran Han, Danielle C Maddix, Abdul Fatir Ansari, Andrew Stuart, Michael W Mahoney, and Yuyang Wang. Gradient-free generation for hard-constrained systems. arXiv preprint arXiv:2412.01786. 2024.
Richard Cheng, Gábor Orosz, Richard M Murray, and Joel W Burdick. End-to-end safe reinforcement learning through barrier functiins for safety-critice continuous control tasks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3387-3395, 2019.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, pp. 02783649241273668 2023.
Jacob K Christopher, Stephen Baek, and Ferdinando Fioretto. Constrained synthesis with projected diffusion models. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35:25683 2222666665666, 2022.
Oscar Davis, Samuel Kessler, Mircea Petrache, Ismail IIkan Ceylan, Michael Bronstein, and Avishek Joey Bose. Fisher flow matching for generative modeling discrete data. arXiv preprint 2rXiv:2405.14400 2024.
Wei Deng, Yu Chen, Nicole Tianjiao Yang, Hengrong Du, Qi Feng, and Ricky TQ Chen. Reflected schrödinger bridge for constrained generative modeling. arXiv preprint arXiv:2401.03228, 2024.
Tomek Diederen and Nicola Zamboni. Flows on convex polytopes, 2025. URL https://arxxii
org/abs/ 2503. 10232
Frederik Diehl. Warm-starting ac optimal power flow with graph neural networks. In 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), pp. 1-6, 2019.
Shutong Ding, Jingya Wang, Yali Du, and Ye Shi. Reduced policy optimization for continuous control with hard constraints. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Carles Domingo-HHi Ennichh Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprini 1rXiv:2409.08881, 2024.
Priya L. Donti, David Rolnick, and Zico Kolter. De3: learning method for optimization with hard constraints. In International Conference on Learning Representations, 2020.
Adam N Elmachtoub and Paul Grigas. Smart "predict, then optimize". Management Science, 68(1): 9-26, 2022.
Ying Fan and Kangwook Lee. Optimizing ddpm sampling with shorteut fine-tuning. arXiv preprint 203.122301.3366. 2023.
Berthy T Feng, Ricardo Baptista, and Katherine Bouman. Neural approximate mirror maps for constrained diffusion models. arXiv preprint 2rXiv:2406.122 12816, 2024.
Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael Hutchinson. Diffusion models for constrained domains. arXiv preprint arXiv:2304.05364 2023.
12

---

Published as a conference paper at ICLR 2026
Nic Fishman, Leo Klarner, Emile Mathieu, Michael Hutchinson, and Valentin De Bortoli. Metropolis sampling for constrained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. arXiv preprini arXiv:2410./25533 2024.
Thomas Frerix, Matthias Nießner, and Daniel Cremers. Homogeneous linear inequality constraints for neural network activations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp. 748-749, 2020.
Yuan Gao, Jian Huang, and Yuling Jiao. Gaussian interpolation flows. Journal of Machine Learning Research, 25(253): 2024a.
Yuan Gao, Jian Huang, Yuling Jiao, and Shurong Zheng. Convergence of continuous normalizing flows for learning probability distributions. arXiv preprint arXiv:2404.005S1, 2024b.
Alexandros Graikos, Nebojsa Jojie, and Dimitris Samaras. Fast constrained sampling in pre-trained diffusion models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL BR.
Yunrui Guan, Krishnakumar Balasubramanian, and Shiqian Ma. Mirror flow matching with heavy- tailed priors for generative modeling on convex domains. arXiv preprint arXiv:2510.089/ 2025.
Yuhang Guo, Zicheng Su, Hai Yang, Enming Liang, Chen Zhong, and Wanjing Ma. A smart predict-then-opiimize framework for vehicle rebalancing problem. Transportation Research Part B: Methodological, 206:103411 2026.
Guillermo Hansen, Irmina Herburt, Horst Martini, and Maria Moszyúska. Starshaped sets. Aequa- tiones mathematicae, 94: 14:1000-1092, 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6855 6851, 2020.
Noémie Jaquier, Leonel Rozo, Darwin G Caldwell, and Sylvain Calinon. Geometry-aware manipula- bility learning, tracking, and transfer. The International Journal of Robotics Research, 40(2-3): 624-650, 2021.
Shervin Khalafi, Dongsheng Ding, and Alejandro Ribeiro. Constrained diffusion models via dual training, arXiv preprint arXiv:2408.1 15094, 2024.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.69888 69980098,. 2014.
Andrei V Konstantinov and Lev V Utkin. A new computationally simple approach for implementing neural networks with output hard constraints. In Doklady Mathematics, volume 108. pp. S233- S241. Springer, 2023.
Yunbum Kook and Santosh Vempala. Gaussian cooling and dikin walks: The interior-point method for logconcave sampling. In The Thirty Seventh Annual Conference on Learning Theory, pp. 3137-3240. PMLR, 2024.
Anastasis Kratsios, Behnoosh Zamanlooy, Tianlin Liu, and Ivan Dokmanić. Universal approxima- tion under constraints is possible with transformers. International Conference on Learning Representations, 2021.
Dohyun Kwon, Ying Fan, and Kangwook Lee. Score-based generative modeling secretly minimizes the wasserstein distance. arXiv preprint arXiv:2212.06333) 2022.
John M Lee. Riemannian manifolds: an introduction to curvature, volume 176. Springer Science & Business Media, 2006.
13

---

Published as a conference paper at ICLR 2026
Anjian Li, Zihan Ding, Adji Bousso Dieng, and Ryne Beeson. Efficient and guaranteed-safe non- convex trajectory optimization with constrained diffusion model. arXiv preprint arXiv:2403.05577, 2024.
Ke Li and Jitendra Malik. Learning optimize, 2016. URL DRLLttpp:://araa..cz/////
01885
Meiyi Li and Javad Mohammadi. Toward rapid, optimal, and feasible power dispatch through generalized neural mapping. arXiv preprint
2023.
Meiyi Li, Soheil Kolouri, and Javad Mohammadi. Learning to solve optimization problems with hard linear constraints. IEEE Access, 2023a.
Ruizhe Li, Enming Liang, and Minghua Chen. On the expressivity of gnn for solving second order cone programming. In NeurIPS Workshop on GPU-Accelerated and Scalable Optimization, 2025a.
Ruizhe Li, Enming Liang, and Minghua Chen. On the universality and complexity of gnn for solving second-order cone programs. In The Fourteenth International Conference Learning Representations. ICLR, 2026.
Xinpeng Li, Enming Liang, and Minghua Chen. Gauge flow matching for efficient constrained generative modeling over general convex set. In ICLR 2025 Workshop Deep Generative Model in Machine Learning: Theory, Principle and Efficacy, 2025b.
Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffusson convolutional recurrent neural network: Data-driven traffic forecasting. In International Conference Learning Representations (ICLR '18), 2018.
Yang Li, Jinpei Guo, Runzhong Wang, and Junchi Yan. T2t: From distribution learning in training to gradient search in testing for combinatorial optimization. Advances in Neural Information Processing Systems, 36:50020-50040, 2023b.
Yangming Li, Zhaozhi Qian, and Mihacla van der Schaar. Do diffusion models suffer error prop- agation? theoretical analysis and consistency regularization. arXiv preprint arXiv:2308.0500 050211, 2023c.
Enming Liang and Minghua Chen. Generative learning for solving non-convex problem with multi- valued input-solution mapping. In 12th International Conference on Learning Representutions (ICLR 2024). ICLR, 2024.
Enming Liang and Minghua Chen. Efficient bisection projection to ensure neural-network solution feasibility for optimization over general set. In Proceedings of the Forty-second International Conference on Machine Learning. PMLR, 2025.
Enming Liang, Minghua Chen, and Steven H Low. Low complexity homeomorphic projection to ensure neural-network solution feasibility for optimization over (non-) convex set. In Proceedings of the 40th International Conference on Machine Learning. PMLR, 2023.
Enming Liang, Minghua Chen, and Steven H Low. Homeomorphic projection to ensure neural- network solution feasibility for constrained optimization. Journal of Machine Learning Research, 25(329):1 155,1 2024.
Jinhao Liang, Yixuan Sun, Anirban Samaddar, Sandeep Madireddy, and Ferdinando Fioretto. Chance chante-constrained flow matching for arXiv:2500.251157 2025.
constraint- generation. arXiv preprint
Martin Werner Licht. On the geometry of star domains and the spectra of hodge-laplace operators under domain approximation. arXiv preprint arXiv:2411.01655. 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint a0X.v:22 2022
Chenghao Liu, Enming Liang, and Minghua Chen. Characterizing resnet's universal approximation capability. In Proceedings of the Forty-first International Conference on Machine Learning. PMLR, 2024a.
14

---

Published as a conference paper at ICLR 2026
Chenghao Liu, Enming Liang, and Minghua Chen. Fast projection-free approach (without optimiza- tion oracle) for optimization over compact convex set. The Thirty-nimi Anmual Conference on Neural Information Processing Systems, 2025.
Guan-Hou Homg Liu, Tianrong Chen, Evangelos Theodorou, and Molei Tao. Mirror diffusion models for constrained and watermarked generation. Advances in Neural Information Processing Systems, 36, 2024b.
Tianyu Liu, Qianhao Wang, Xingguang Zhong. Zhepei Wang, Chao Xu. Fu Zhang, and Fei Gao. Star-convex constrained optimization for visibility planning with application to aerial inspection. In 2022 International Conference Robotics and Automation (ICRA), pp. 7861-7867. IEEE, 2022a.
Xingchao Liu and Lemeng Wu. Learning diffusion bridges constrained domains. In international conference on learning representations (ICIR), 2023.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03004 2022b.
Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning, pp. 22675-22701. PMLR, 2023.
Zhou Lu, Nataly Brukhim, Paula Gradu, and Elad Hazan. Projection-free adaptive regret with membership oracles. In International Conference on Algorithmic Learning Theory, pp. 1055-1073.
PMLR, 2023.
Jayanta Mandi, James Kotary, Senne Berden, Maxime Mulamba, Victor Bucarey, Tias Guns, and Ferdinando Fioretto. Decision-focusee learning: Foundations, state of the art, benchmark and future opportunities. Journal of Artificial Intelligence Research, 80: 10:1112177,, 2024.
Zakaria Mhammedi. Efficient projection- prrcctiion-fre online convex optimization with membership oracle. In Conference on Learning Theory, pp. 5314-5390. PMLR, 2022.
Benjamin Kurt Miller, Ricky TQ Chen, Anuroop Sriram, and Brandon M Wood. Flowmm: Generating materials with riemannian flow matching. arXiv preprint (xXi:22406.04113 2024.
Saeid Naderiparizi, Xiaoxuan Liang, Berend Zwartsenberg, and Frank Wood. Constrained generative modeling with manually bridged diffusion models. Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 19607-19615, 2025.
Sai Shankar Narasimhan, Shubhankar Agarwal, Litu Rout, Sanjay Shakkottai, and Sandeep P. Chinchali. Constrained posterior sampling: Time series generation with hard constraints. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL
https://epeenrevieest wet.e/weeeeeeeeeee.eeeeeeeee d 1tt.
Rahul Nellikkath and Spyros Chatzivasilciadis. Physics-informen neural networks for minimising worst-case violations in de optimal power flow. In 2021 IEEE International Conference Communications, Control, and Computing Technologies for Smart Grids (SmartridConmm) pp- 419-424. IEEE, 2021.
Xiang Pan, Tianyu Zhao, Minghua Chen, and Shengyu Zhang. Deepopf: deep neural network approach for security-constraiie de optimal power flow. IEEE Transactions Power Systems, 36(3):175511755, 2020.
Xiang Pan, Minghua Chen, Tianyu Zhao, and Steven H Low. Deepopf: A feasibility-optimizee deep neural network approach for ac optimal power flow problems. IEEE Systems Journal, pp. 42-47, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, segghrerformance deep learning library. Advances in neural information processing systems, 32, 2019.
15

---

Published as a conference paper at ICLR 2026
Sasa V Raković and Mircea Lazar. The minkowski ninkowski-lyapunoo equation for linear dynamics: Theoreti- cal foundations. Automatica, 50(8):2015-200 2014.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25,, 1())3, 2022.
Tyrrell Rockafellar. Second-order convex analysis. J. Nonlinear Convex Anal, 1(1-16):84, 1999.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffiusion models. In International Conference on Learning Representations, 2022. URL httpps://ppenrrvieer
net /
Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen. Solving inverse problems with latent diffusion models via hard data consistency. arXiv preprint arXiv:2307.08/23, 2023a.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.0255: 2020.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023b.
Zhiqing Sun and Yiming Yang. Difusco: Graph-based diffusion solvers for combinatorial optimization. arXiv preprint arXiv:2302.08224 2023.
Daniel Tabas and Baosen Zhang. Computationally efficient safe reinforcement learning for power systems. In 2022 American Control Conference (ACC), pp. 3303-3310. IEFF, 2022a.
Daniel Tabas and Baosen Zhang. Safe and efficient model predictive control using neural networks: An interior point approach. In 2022 IEEE 61 Conference on Decision and Control (CDC), pp. 1142-1147. IEEE, 2022b.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, pp. 1-34, 2024.
Jesus Tordesillas, Jonathan How, and Marco Hutter. Rayen: Imposition of hard convex constraints on neural networks. arXiv preprint 23X1v:2307.08330 2023.
Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous- time diffusion models as regulay-regularized control. arXiv preprint arXiv:2402. 15194, 2024.
Utkarsh Utkarsh, Pengfei Cai, Alan Edelman, Rafael
and Christopher Vincent Rackauckas. Poysicc-constrained flow matching: Sampling generative models with hard constraints. In The Thith Annual Conference Neural Information Processing Systems, 2025. URL https: //tpr////peeeevieww net. weeet/forrmmllddffeetew
Robert A Van de Geijn. Using PLAPACK PLAPACK-paralle! Linear Algebra Package. MIT press, 1997.
Zhengchao Wan, Qingsong Wang, Gal Mishne, and Yusu Wang. Elucidating flow matching ode dynamics with respect to data geometries. arXiv e-prints, pp. arXiv-2412, 2024.
Runzhong Wang, Yunhao Zhang, Ziao Guo, Tianyi Chen, Xiaokang Yang, and Junchi Yan. Linsatnet: the positive linear satisfiability neural networks. In International Conference on Machine Leurning, pp. 36605-36625, PMLR, 2023.
Bernardo Williams, Victor M Yeom-Song, Marcelo Hartmann, and Arto Klami, bijections for categorical flow matching. arXiv preprint arXiv:2510.274880 2025.
Tianyu Xie, Yu Zhu, Longlin Yu, Tong Yang, Ziheng Cheng, Shiyue Zhang, Xiangyu Zhang, and Cheng Zhang. Reflected flow matching. arXiv preprint arXiv:2405./6577, 2024.
16

---

Published as a conference paper at ICLR 2026
Jiaqi Yang, Enming Liang. Zicheng Su, Zhichao Zou, Peng Zhen, Jiecheng Guo, Wanjing Ma, and Kun An. Dff: Recision-focuseee fine-tuning for smarter predict-then oottmiee with limited data. In Proceedings of the AAAI Conference Artificial Intelligence, volume 39, pp. 36868-26877 2025.
Jinsung Yoon, Daniel Jarrett, and Mihaela der Schaar. Time-series generative adversarial networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Olga Zaghen, Floor Eijkelboom, Alison Pouplin, and Erik Bekkers, Towards variational flow matching on general geometries. arXiv preprint wxxv:2502.11988. 2025.
Stefano Zampini, Jacob K Christopher, Luca Oneto, Davide Anguita, and Ferdinando Fioretto. Training-free constrained generation with stable diffusion models. arXiv preprint arXiv:2502.05625, 2025.
Claudio Zeni, Robert Pinsler, Daniel Zügner, Andrew Fowler, Matthew Horton, Xiang Fu, Zilong Wang, Aliaksandra Shysheya, Jonathan Crabbé, Shoko Ueda, et al. A generative model for inorganic materials design. Nature, pp. 1-3, 2025.
Jiawei Zhang, Ziyuan Liu, Leon Yan, Gen Li, and Yuantao Gu. Improving diffusion- inverse algorithms under few-step constraint via linear extrapolation. In The Thitty Annual Confer- ence on Neural Information Processing Systems, 2025. URL /tLLtttt///ppeeeeeeeee
Forum? LF IXhI
Liyuan Zheng, Yuanyuan Shi, Lillian Ratliff, and Baosen Zhang. Safe reinforcement learning of control-affine systems with vertex networks. In Learning for Dynamics and Control, pp. 336-347, PMLR, 2021.
Min Zhou, Enming Liang, Minghua Chen, and Steven H Low. Partially permutation-invariant neural network for solving two-stage stochastic ac-opf problem. IEEE Transactions on Power Systems, 2025.
17

---

Contents
Related Work on NN Feasibility
19
Gauge Mapping over General Convex Sets
19
B.1
Handling Linear Equality Constraints
20
B.2
Gauge Mapping for Inequality Constraints
20
B.2.1
Gauge Function and Properties
20
B.2.2
Gauge Mapping and Properties
21
B.2.3
Computation Methods
23
C
Gauge Mapping for Non-convex Constraints
24
Star-convex Set]
24
Geodesic convex Set
Geodesie-convex
25
More General Non-convex Set
26
Theoretical Analysis GFM
26
D.I
Assumption and Preliminary
26
................................
D.2 Proof of Regularity of GFM in Prop 55
................................
27
D.3 Proof of Error Bound of GFM in Prop. 5.2
................................
28
D.4 Proof of Generation Complexity GFM in Prop. 5.3
28
Experimental Settings
28
E.
Constrained Generative Models Baselines
29
E.2
Reflection computation
................................
29
E,3
Performance Evaluation Metries)
30
Constraint Formulations in Illustrative Examples
30
E.5
Mirror Map Gauge Map Simplex Domain
31
................................
E,6
Configuration Generation in Robotic Control Benchmark
31
E.7
Solution Generation for Relaxed Combinatorial Problems
................................
32
E
Constrained Conditional Time Series Generation
33
E.9
Watermarked Image Generation Tasks
34
E.
10
Setting of Scalability Tests of Gauge Mapping
36
Ablation Study
36
18

---

Published as a conference paper at ICLR 2026
A RELATED WORK ON NN FEASIBILITY
Machine learning has been increasingly integrated with optimization problems to enable efficient decision-makinn Several paradigms have emerged for this integration, each addressing different aspects of the optimization process but sharing common architectural challenges. Learning to optimize trains models to iteratively refine solutions or to learn update rules that accelerate convergence (Li Malik, 2016, Chen et al., 2022). Decision-focused learning leverages downstream optimization tasks to refine upstream prediction models (Elmachtoub & Grigas 2022 Mandi et al. 2024; Yang cct al., 2025] Guo et al.) 2026). End-to-end learning trains neural networks to directly predict solutions to constrained optimization problems, serving either warm starts for solvers
approximate optima (Dehhl 2019 Pan et al. 20200] 2022, Li et al. 2025a 2026). Generative modeling approaches learn solution distributions for hard optimization problems (Sun & Yang/ 2023
Liang & Chen 2024
Li et al. 2023b).
core challenge across these paradigms is ensuring that neural network outputs satisfy the feasibility constraints of the underlying problems. Research on ensuring neural network feasibility can be categorized into several approaches:
Basic Constraint Handling, Specialized activation functions (e.g., Sigmoid, Softmax) address basic
constraints such as box simplex constraints (Pan et al., 2022 Donti et al., 2020). Penalties for
output constraint violations can be incorporated into loss functions to improve feasibility (Cheng
al.
2019
Pan al. 2020 Nellikkath & Chatzivasileiadis 202I).
Strict Satisfaction Methods. For exact equality constraint satisfaction, prediction-thent
reconstruction or completion techniques can be applied (Donti et al., 2020) Pan al. 2022; Liang
et al. 2023). For more general constraint enforcement, orthogonal or L2 projection is often employed.
However, solving the projection problem, either via an iterative solver or through equivalent opti-
mization layers (Amos& & Kolter 2017 Agrawal et al. 2019; Chen et al.) 202 Wang et al. 2023),
computationally intensive for real-time applications. More efficient disection-based projection
methods can be applied at the cost of minor optimality loss (Liang et al. 2023
2024
Liang & Chen
2025
Zhou et al. 2025).
Set Parameterization Approach. To guarantee feasibility, an inner approximation of the original constraint set can be constructed. For linear constraints, vertex networks employ convex combinations of vertices and rays ensure policy feasibility (Frerix al., 2020 Zheng al. (021). For general compact but fixed constraint sets, the probabilistic transformer utilizes feasible samples to ensure feasibility (Kratsios al. 202Tb. However, scalability remains challenge due to the exponential growth in required samples as problem dimensionality increases.
Gauge Functions.
These methods utilize gauge functions (Blanchini & Miani 2008
constrain
neural network outputs. A closed-form bijection, known as gauge mapping. between a hypercube and polytope is used bound the network output within the polytope (Tabas & Zhang 2022a b)
Li et al. 2023a). Gauge mapping can be generalized to arbitrary convex sets (Li et al. 2025b Liu
al. (225). For fixed convex constraints, several works apply gauge functions find feasible
boundary solutions via scaling operations (Tordesillas et al. 2023 Konstantinov & Utkin 2023 &
Mohammadi 2023).
These approaches are typically designed for the end-to-end neural network structure, which can not be directly applied the diffusion or flow-matching based generative model due to the forward integration calculation.
B GAUGE MAPPING OVER GENERAL CONVEX SETS
general compact convex set encompasses both linear equality and convex inequality constraints.
C fx E R Ar = b, 91 (x) < 0, .............................. .m (x) < 0}.
(6)
where A E R b R", and (z), ..... 9m (I:) are convex functions.
This section presents a systematic approach to handling such sets by first eliminating linear equality constraints, followed by computing gauge mappings for the remaining inequality constraints.
19

---

Published as a conference paper at ICLR 2026
B.1 HANDLING LINEAR EQUALITY CONSTRAINTS
Without loss of generality, assuming rank(A) - r, we partition the decision variable x into æ1 E Re- and 22 E RR. Accordingly, we partition matrix A into A = [A1, A2], where A E (R=x(e-r) and A2 Rexr, The equality constraint Ax b can then be written as:
Aạo A2æ2 = b
By choosing the partition such that A₂ has full rank т, we can express 22 explicitly in terms of ø1:
T2 = )(x1) = A A121)
(8)
This transformation reduces the original set one with only inequality constraints:
C {क R"~ [[[ - 0}
Therefore, only consider the incquality constraints in the main body of this work.
B.2 GAUGE MAPPING FOR INEQUALITY CONSTRAINTS
Without loss of generality, we consider C = { 91 (z) 50.....9((1)
(x) S 50}.
We define the following metries for compact convex set.
Definition B.1 PPoit-to-boundall distance and its inverse (Tordesillas et al. (023)). ) Let C c R" be
a compact convex set and x E int(C) an interior point. For any unit vector e E 1n51_= = (u E IR"
||u|| = 1}, define the interior-pointe seerioo-ppinntto-ouundaa distance function de int(C) x $ R. along direction as
de(z", = supp> 2 0 x + Au C).
(10)
The inverse distance function Re: int(C) x s--1 + R is defined as ne(x", v) == 1/d
For a compact convex set, the minimum and maximum interior point-to-bouddrry distances are bounded and are defined as:
the minimum hointto-bbundar distance: = arg an CC.
the maximum point-to-boundary distance: To arg arginf,zole C B2(x",
thus, we have B2 ,mi S CB22"" ro))
B.2.1 GAUGE FUNCTION AND PROPERTIES
We then introduce the gauge function:
Definition B.2 (Gauge/Minkowski function (Blanchini & Miani 2008)) Let C C IR" he a compact
convex with non-empty interior. The Gauge/Minkowski function % IR" x int(C) R4 is defined as
Ye(r, x0) = inf ((-=)))))))))
(11)
where E int(C) is an interior point of C.
The Gauge function generalizes the concept of a norm. For set that is symmetric about the origin, the gauge function (c, 0) defines a norm. In particular, when C = B, = {i E IR"" ||ا|| I 1] is the unit ball of the p-norm, we have Y6, (I,, 0) ||=]|pp- More generally, the gauge function satisfies the following properties for all x, E 18 and a 2 0:
Lemma (Properties of gauge function). The gauge function (2, with respect to a compact convex set C and interior point x E int(CC satisfies the following properties:
Non-negativity: ne(x, x²)
Positive homogeneity: xe ra( x²) = ave(a, for a 2 0.
Subadditivity: (0 y, v, S 22) oe)
20

---

Published as a conference paper at ICLR 2026
Convexity: induced by positive homogeneity and subadditiviry.
Differentiabilltyy under convexity, the gauge function is twice differentiable almost everywhere by Alexandrov's 's theorem (Rockafellar 1999).
Equivalent formulations based on (inverse) distance function:
Tc(x, x²) = ne(x", (=((==)////))))) ||=|||| ||3|||
(12)
de(x>xx//10) x
re <c(x/|<=]||,²") = (=0==,=//(<ll]] de(x;; x//1|z11)
(13)
Upper/lower bounds: the gauge function is bounded as: re(x, ve(2xx") x²)
|||=:1//r |
Lipschitz: the gauge function has Lipschitz constant as [nelar")= 11
Proof. The non-negativity, positive homogeneity, and subadditivity are provided in (Blanchini & Miani 2008). Convexity and differentiability are naturally implied by those properties. The equivalent formulation is derived from their definitions.
The upper and lower bounds are derived as:
(e(2, x²)
x<(x",//||=1] E [|læ|]/ro [xx||//r]]
(14)
The Lipschitz constant of the gauge function is derived as:
(c(x, S 70(x )") + vc xv,x))
(15)
Te xx) - (y, x²) S %e( - y, x²) S || y yll/r;
(16)
where the inequality 115) is derived by the positive homogeneity. Similarly, we have:
(y - ye(x, x²) S 7e x, xx x) < || yll/r;
(17)
Combining two inequalities, we have the Lipschitz as:
|,, x²) ic (y, x))
1
5 x v
(18)
B.2.2 GAUGE MAPPING AND PROPERTIES
Based on the gauge function, we can construct a bijective mapping between any pair of compact convex sets as:
Definition B.3 (Gauge Mapping). Let Z,X' C R be compact convex sets with interior points = int(2) 2) and x E int(()), respectively.
The gauge mapping t 2 + X' is defined as:
(xz) = 72 20)
(2 20) + xx" 2
(19)
7x 7x x0)
The inverse mapping +-
:X + 2 is given by:
7x (x - x² x,, x²)
Φ (x)
(x)
(x - x0)
20, XEZ X
(20)
7 x 20)
In essence, the gauge mapping scales the boundary of a convex set from an interior point to another convex set and translates it to its interior point. When Z is a unit p-norm ball, the gauge mapping is simplified in Def.
as:
|=1|p
7c(1 x-, x
((2)
70(2, x0) + VZ E,
x-1(x)
2²), Vx E C,
(21)
x - 1
Further, the gauge mapping between B2 and C can be simplified as:
(<2) de(x,, , z + x., Vz E B, Φ - de(x x x
x
Vx C (22)
21

---

Published as a conference paper at ICLR 2026
Proposition B.1 (Properties of gauge mapping). The gauge mapping between any pair of compact convex sets satisfies the following properties:
it is invertible.
it is continuous everywhere and twice differentiable almost everywhere in both directions.
it is bi-Lipschitz homeomorphism.
Proof. First, the invertibility can be easily verified by:
② (x)) = ② (x - x² x²)
(xr - x²) =0)
(23)
72 (r x0 20)
2z 72 x- ,2
(0-x" xx)
( - x. )=,0) 20) 7x(1
1 x0)
(x- xx))
(24)
(0 20 _00
x ) 72 x² 2-)
7x 72( (00--" ,") )2)
2x(2-2 12 ((x - -0) 22) 1x x0)
(= x2, ,20)
1/x ((x x²)) x>) 72(x x x²,20) 20) (xr - x²) 4
(25)
2x(x=x"
Yz (a x² 22
= (r x + x =
(26)
where Eq. (25 is derived by the positive homogeneity of the gauge function.
Second, the continuity and differentiability are derived from the properties of the gauge function and elemental compositions in the gauge mapping.
Third, the homeomorphism derived from the continuity in both directions and the invertibility.
Proof of Bi-Lipschitz Properties of Gauge Mapping in Prop.
To derive the bi- bi-Lipschitz properties, we first consider the gauge mapping between B2 and C and derive its Lipschitz constants shown in Prop. 4.1aas as follows,
| 20(2, =
((z)
+ x,, Vz E B2
(27)
Differentiating (=) with respect to (using the product and quotient rules) yields
0$ ||
|| 72(2, 20 2 (Venc(z,2"))
+
(28)
dz
vc (z,
x))
(2, x 1
Taking the operator norm and applying the triangle inequality gives
04
|
||=||2
|| 20)2
V=ro(2,2))
<
+
(0(2,20)) +
0=
7C
(2,
(29)
< ro +
+
ni
where in the last inequality we have used the facts that (i) for E B2 one has ||||| < 1, (ii) the gauge function satisfies (0(2, [|l=|//r: ||=]] and (iii) (V, xI bounded by 1/ri
In summary, we obtain
< 2r.
(30)
0z
which proves that the forward Lipschitz constant of t satisfies
2
Forward Lipschitz:
Lg < 2r. +
(31)
22

---

Published as a conference paper at ICLR 2026
Next, consider the inverse gauge mapping from C to the 2-norm ball
"Xc (ar
x=; x²)
1 (z)
|| x (xr x²), Vo E
(32)
Differentiating with respect to x gives
(x [
04-1
Vrc (z - z², x²))
x
+ 7 x², x0)x))
1
(33)
dz
x
|
x 2
Taking norms and again using the triangle inequality leads
( r Using the bound vc (x - x², x², x0) E |||r x most 1/|| - x² (12), we have
② 0x
Vre(x x2, x0)09)
| x² 112 )
)'
+ 7c - =x)200)
I
1w-
a x" 13
11 x
(34)
2 x
|| [1/|
x 1/|l/ and the projection term has norm at
04-
1 1
(35)
dx
r3
11
n
Thus, the inverse Lipschitz constant is bounded by
2
Inverse Lipschitz:
Lo
(36)
This completes the proof of Prop.
To extend the Bi-Lipsehitz properties any pair compact convex sets, e.g., and X, we decompose the gauge mapping as
X ( = (B2) $; (
(37)
where 2 = Φ (B2) and X = 2 (B2).
Leveraging the inequality of Lipschitz for mapping compositions as Lfiofa S T, T2g2, we can conclude that the gauge mapping between any pair of convex sets is bi-Lipschitz.
B.2.3 COMPUTATION METHODS
IP computation
In practice, we seek such "central" interior point by solving the following residual minimization problem through convex optimization in the offline phase (Tordesillas et 2023):
min
(38)
s.l. (x Sn 1, 1......mm
(39)
We note that solving this convex optimization problem with linear objective incurs only polynomial time complexity. As this computation is performed offline prior to model training, it adds negligible overhead to the overall computational cost.
Gauge function computation
Computation of gauge mapping essentially involves calculating the gauge function or (inverse) distance function. Following the established closed-form distance function calculation for several common convex sets in (Tordesillas et al. 2023), we give a summary in Table It provides closed- form expressions for the inverse distance function across various constraint types. Most matrix calculations can be computed and stored offline before being applied for online inference.
When the inverse distance function lacks explicit expression, employ efficient bisection algorithm detailed in algorithm 3] This algorithm supports batch processing, enabling efficient parallel computation for multiple inputs simultaneously.
23

---

Published as a conference paper at ICLR 2026
Table 8: Closed-form Expressions for Inverse Distance Functions (Tordesillas et al. 2023
Constraints
Formulation
Inverse Distance Function
Intersections
(œ) 20, .. (m) 2 50}
Kg v) max 15
(Ka (x", ,1)
Linear
(x) a=x a 0
Kag (x ,2) b-aTzol }+
a a'v
Quadratic
[Q(r) = r ='Qz + -
Kag (x", 0) {1/root(Aq. Bq, Cq)}
Second Order Cone gs(a) |A'x p1a (a x b) S Kig (I v) {1/root (As, Bs, Cs)] }
Matrix Cone
9Mr (M (x) Σi=
F;
Kag (x .v) max(eig(1.
(-S))) (-S)L)}²
Notation: 2,
Q E
gux* .
DER", R* Fo:.. 1 Fn RExE
Aq = Qu, Bq 2x +v++avaaa
Tu e, = Q2" +a ta
As (A 6) (A v)-(a u)² B5 = 2(A +p) (A v)-2(a x +bila) Cs
(a b)²
H Fa Ei= 2 F. H = L Et F..
(+) max(-;
6 moot(21: 22. m3)
denotes quadratic equation solution
cig( X) Au.... An denotes the eigenvalues satisfying det Af) Note that the maximum eigenvalue needed.
Thus, power iteration methods can be applied to compute it efficiently.
Note that v-independent
be computed only and stored for
Algorithm 3 Bisection Algorithm for Point-to-Boundary Distance
Input: A compact convex set C, an interior point x E Eint(C), and unit vector v. 0 and = 2: while |o; - aa1 2 do
1: Initialize:
3: if' x:" + au v E C then
4: increase lower bound: a f an.
5: double upper bound: Oaa t- am
6: else
7: bisection: Om (an + o2)/2
8: if u + Om v E C then
9: increase lower bound: og 0'm
10: else
decrease upper bound: a₂ f- a'm
12: end if
end if
14:
end while
Output: de (xx, v) Orn
C GAUGE MAPPING FOR NON-CONVEX CONSTRAINTS
In this section, we extend the conventional gauge mapping over convex sets in Euclidean space to certain non-convex settings, including star-convex sets and geodesic-convex sets.
C.I STAR-CONVEX SET
Definition C.1 (Star-convex Set). A set S C I is said to be star-convex with respect to interior point d E (S) if for every & E S, the line segment connecting x" and æ lies entirely within S, i.., (z 2²) ) for all 0 [0, 1].
The gauge mapping framework can be extended beyond convex sets to certain non-convex domains, particularly star-convex sets. This extension significantly broadens the applicability of our GFM approach to more complex geometric constraints encountered in practical applications. A star-convex set S is characterized by the existence an interior point x such that any line segment connecting xx to any point in the set remains entirely within the set. While star-conver sets lack the full convexity property, they retain a critical radial structure that allows gauge mappings to be constructed similar manner.
24

---

Published as a conference paper at ICLR 2026
Z
x x= (Z)
X
Z2
X2
x
B
S
Figure 9: Gauge mapping construction for star-convex set.
For a star-convex set S with interior point x², we can define the gauge function as:
(s(x, ()) inf (> 2 0 x E X - x²)}
(40)
This formulation captures the minimum scaling factor needed to reach point when scaling along the ray from x²_ The corresponding gauge mapping + B - S between a unit p-norm ball and the star-convex set follows the same construction as in Definition 4.] As illustrated in Figure the gauge mapping preserves the radial structure of the star-convex set while establishing a homeomorphism with the unit ball. This property particularly valuable for handling constraints with non-convex geometries that still maintain visibility from an interior point, such as fp-norm balls with < 1 or certain non-convex polytopes arising in practical applications.
The computation of gauge functions for star-convex sets follows similar principles as for convex sets, though additional care may be needed:
For simple star-convex sets (e.g., f--norm balls with p < 1), closed-form expressions for the gauge function can often be derived as Def. 4.
For star-convex sets defined by the union of convex components by non-convex inequalities, bisection methods remain applicable for computing the gauge function.
Selection of the interior point x becomes more eritical for star-convex sets, as it determines the visibility region and thus the quality of the mapping. When multiple interior points are viable, selecting one that maximizes the minimum distance to the boundary often yields better numerical properties.
C.2 GEODESIC-CONVER SET
When the underlying space is Riemannian manifold rather than Euclidean space, the notion of straight lines is replaced by geodesics. In this context, one speaks of geodesic convexity.
Definition C.2 (Geodesic-Convex Set). Let (M, .) be a Riemannian manifold, and let U A be subset. V is geodesically convex if, for every pair of points x, E U, there exists a unique minimizing geodesic 7) [o. 1] + M with respect to g that satisfies 170) x, (7(1 = y. n(t) E U, Vt E [0,
To extend the gauge mapping with the exponential map. Since 2/ C M geodesic-conver the exponential map at any interior point is a diffeomorphism onto its image (Lee] 2006). Given this property, any point x E U with r * x can be uniquely represented in geodesic polar coordinates as x = exp (r, Vz), where r" E l/ is an interior point U4 Ur E Sn-1 C T. M is a unit vector over the tangent space T2 M, and rg = (œ, x²) is the geodesic distance from x to X.
geodesic-conee convex set over manifold, we replace the linear structure
Similarly, we can establish the geodesic gauge mapping that transforms unit ball from the tangent space to the geodesic-comee convex set on the manifold. As illustrated in Figure I0 this mapping establishes correspondence between the unit ball in the tangent space and the geodesic-convex set the manifold. The construction respects the intrinsic geometry of the manifold while preserving the radial structure that is characteristic of gauge mappings.
25

---

Published as a conference paper at ICLR 2026
= (())
14
0
U
B
M
Tx² M
Figure 10: Gauge mapping construction for geodesic-convex set.
Definition C.3 (Geodesic Gauge Mapping). Without loss of generality, let B - {= E Tyo M
||=||2 < }} denote the unit ball in the tangent space x² The geodesic gauge mapping $ : and its inverse are defined as
((2) ,xp, (du(z 2/ 2/|1=/11 2), (z) Tx T dx/" Ux) Uns Vr expx (F=Ux) El
Va E B
(41)
(42)
where the point-to-boundary geodesic distance du(x" uz) along U, is defined as du(v))) sup 0 expx (A Ux) U}.
Note that the woint-to-boundee geodesic distance may not admit a closed-form expression; we can apply the bisection methods similar to Alg. 3 to locate the boundary points starting from an interior point.
This extension enables our GFM framework to handle constraints manifolds, expanding its application domains such as inverse problems with manifold constraints and atmospheric modeling over spherical geometries.
C.3 MORE GENERAL NON-CONVEX SET
When applying our methods to more general non-convex sets, a natural choice is the class of non- convex sets that homeomorphic generally admit closed-fore ffrm homeomorphism. One promising approach is to apply an invertible neural network to learn the target homeomorphism, which offers both feasibility and approximation guarantees (Liang et al. 2023 2024). For more general non-convex constraints, low cwwcooply constrained generation schemes remain largely unexplored, particularly for NP-hard problems where even identifying an interior point poses significant computational challenges, such as non-convex quadratic equations or mixed-integer formulations. The exploration of more general and theoretically guaranteed approaches constitutes a promising direction for future research.
unit ball. However, such Dall-homeomorphic sets do not
D THEORETICAL ANALYSIS OF GFM
D.1 ASSUMPTION AND PRELIMINARY
We made the following assumptions for the error analysis, which are common for error analysis for
the flowddifuusion-asee generative models (Kwon et al. 2022; Benton et al. 2023 Xie et al., (222).
Assumption (Regularity of Distribution (Wan et al. 2024 ) We assume the data distribution
satisfies the following regularity conditions:
It has finite 2-moment, i.e., fg- ((dx))((dx) < co;
Its support, Ω supp(p), has positive reach;
It has non-vanishing density, i.e., there exist constants & 0 and 0 such that for any radius >> 0, there exists a constant CR > 0 where for any small radius < r < and any E (0) n Ω we have B, (x)) 2 Cpr.
26

---

Published as a conference paper at ICLR 2026
The finite 2-moment requirement is standard in the literature for analyzing generative models. The positive reach condition is satisfied in our setting since we consider a convex support set f = C. The last requirement ensures the density does not vanish within the support, which prevents pathological cases where probability mass concentrates on lower-dimensional manifolds.
There also other (more restrictive) regularity conditions for the data distribution in the literature, including the covariance condition (Benton et al. (2023) and
(Gao et al.) 2024a), and convex support set (Gao al., 2024b). It remains largely open for the minimal necessary conditions to establish the well-posedness of flow matching models.
Assumption (Neural Network Approximation Properties). Given a distribution p with support Ω, let u(x, i) be the designed target vector field and up be the neural network approximated vector field the flow-matching model.
Up is Lg-Lipschitz for Ω and t [0, 1];
The l approximation error is bounded eg Ex,,://oo( (12,, u(ze; where pr is the probability density driven by the target vector field u.
The Lipschitz constant of trained neural network is bounded within a compact set in our setting ? = C. The training loss in Eq. ①I is equivalent to the vector field approximation error up to constant (Lipman et al. 2022). Therefore, we can minimize the approximation error through proper training of the loss function.
Lemma 2 (Error Bound for Flow Matching (Benton et al., 2023)). For the vanilla flow matching
model: -o So (ar, L) dt, with induced probability distribution po at t 1. The Wasserstein-2 distance between the data distribution Pdota (1) and the approximated distribution po po(x) is bounded by
W2(Pa Pa(x)) S cc co
(43)
Lemma 3 (Error Bound for Reflected Flow Matching (Xie et al. 2024)). For the reflected flow
matching model over the convex domain C: ø1 Io 0g(x, +)+RR((()) dt, with induced probability distribution pi at t = 1. The Wassserstein-2 distance between the data distribution pdata (x) and the approximated distribution po (²(x) is bounded by
Wa(Paata (z(2)) S e/24Lo
(44)
D.2 PROOF OF REGULARITY GFM IN PROP.
Proof. Given the original data distribution p over compact convex set Dpp = C satisfying the regularity conditions; We verify the three regularity conditions one by one for the transformed data distribution q 4 p over unit ball S2 B, where & is a bi bi-Lipschitz homeomorphism (c.g.. gauge mapping).
Finite second moment: since g has bounded support, it immediately implies finite second moment.
The support of q is unit ball 2g = B, which has infinite positive reach. We remark that even the original support 2p does not have a positive reach (e.g., star-shaped set), after transformation, the S2y is always unit ball and satisfies this regularity condition on the reach.
Positive reach:
Non-vanishing density: Given the bi-Lipschitz property, we have:
T =21 < (=1) +(22))l < T22|==1 22l1, Vz =2 E B
(45)
where L, 1/T.4 and L = L. are the (inverse) Lipschitz constants.
For any E S2 n (0) and any small radius r < f. where c is from the non-vanishing density property of we have:
al (=) = (((b
(46)
2 pp(rrrr (=)))
(47)
since ((B,(=)) (T)r(((=)) by the bi-Lipschitz property.
27

---

Published as a conference paper at ICLR 2026
Since & is bi-Lipschitz, for any R O, there exists R' such that (B 0 2ggg S BR (0) n.. By the non- vanishing density property of p. there exists a constant Cr > such that p(BL,r(((2)) 2 CR (Lir).
Therefore, we have
(B, 2 Lg
(48)
This establishes the non-vanishing density property for q-
D.3 PROOF OF ERROR BOUND OF GFM IN PROP. 5.2
Proof.
If || - x2212)d7}
(Panta: Po Po )==
inf T=I1(pacu pgT))
(49)
1f [10(z1) 2x22)) 113)d7}
inf
(50)
7=II(quata: og
1/ ||2 22(13)d7}
L3 7- (adammn:0)))
inf
(51)
< L3 W/ [gdatas 46)
(52)
(53)
where Eq. (50 is by the equivalence of the distributions under homeomorphic push-forward mapping: Pdata = 4 and p Qy Eq. (5T) is by the Lipschitz property of gauge mapping & shown in Prop. 4.1 Eq. (53) is the error bound of the regular reflected generation over the unit ball under Lemmal 3
The Lipschitz constant of gauge mappings between compact sets remains bounded inher- ently. This property stands in contrast to mirror mapping-based generative models (Liu et al. 2024b), which map open convex as points near the boundary are mapped to infinity, significantly complicating approximation error analysis. Our Gauge Flow Matching circumvents this limitation, providing theoretical guarantees on the Wasserstein-22 distance between the learned and data distributions. To optimize the model's 's performance, can further reduce the Lipschitz constant the gauge mapping by identifying an interior point x" that serves as the "center" of the constraint set.
Remark.
to IR... In the latter case, the Lipschitz constant can grow unbounded
D.4 PROOF OF GENERATION COMPLEXITY OF GFM IN PROP. 5.3
Proof. Closed-form computation: Based on Table 8 the algorithmic complexity of naturally derived under different constraint functions.
Bisection based computation: Based on algorithm 3 the algorithmic complexity for the bisection algorithm needs the feasibility check E C. which essentially involves calculating the constraint function 9e for = 1,..., m. The number of bisection steps needed to achieve a target error is derived as (Ciam((C 1og(1/(tiz)).
E EXPERIMENTAL. SETTINGS
In this section, we describe the experimental setup used to generate the results reported in Section 7 Our proposed model is implemented in PyTorch (Paszke et al. 2019), with all models trained using Adam optimizer (Kingma & Bal 20114) with hyperparameters 3 0.99, 83 = 0.999, and learning rate of 10-3 For sample generation, solve the forward ODE using the method proposed by (Chen
28

---

Published as a conference paper at ICLR 2026
al. 2018). and follow the approach in (Xie et al. 2024) solve the reflected ODE in Equation
E.
CONSTRAINED GENERATIVE MODELS BASELINES
We compare approach against the following state-of-the-an constrained generative models, selected for their ability to handle various constraint types:
FM: Vanilla flow matching that transforms a Gaussian distribution the target distribution using
linear conditional flow (Lipman et al. 2022 Liu et al. 2022b).
DM: Vanilla diffusion model with pariance-preservinng diffusion process (Ho et al. 2020 Song
et al. 2020).
Reflection: A method that applies reflection terms when generated samples encounter constraint boundaries (Xie al. 2024).
Metropolis: Metropolis Sampling approach for approximating Inrection-based generation (Fish man al. 2024).
Projection: An approach utilizing orthogonal projection when generated samples violate con- straints (Christopher et al. 2024).
GFM: Our proposed framework as detailed in Section
We remark that all models share the same training settings (dataset, optimizers, hyperparameters) and generation settings (ODE methods, step size). For Reflection, Metropolis, and Projection models, the same velocity models are used. Therefore, they share the same training time reported in the Table.
S
The additional constraint-handling mechanisms (e.g., reflection and projection) are implemented based on references. We emphasize that closed-form projection operators generally do not exist for the complex constraints in experiments, such intersections of polytopes and ellipsoids, and semidefinite cone constraints. Therefore, we employ state-of-the-ard optimization solvers: (i) MOSEK for convex projection and (ii) IPOPT for non-convex cases. While these highly optimized solvers, iteratively solving constrained optimization problem at each sampling step for each sample (whenever the ODE trajectory violates constraints) remains computationally expensive. This cost particularly significant for nigh-dimensional problems or tight constraint sets where violations occur frequently throughout the sampling process.
We also remark that, in the geodesic-convex set generation tasks, the FM and DDPM models are trained generate samples in the tangent space of prescribed point, which is the same across all methods.
E.2 REFLECTION COMPUTATION
We follow the notations Xie et al. 2024) and (Lou Ermon 2023 Let L, be the reflection term that reflect the outward velocity the boundary dC. Given an initial point Zinits the reflected ordinary differential equation is
dr; = va t) + dLe,
(54)
xog Tinit-
(55)
Intuitively, the reflection term L, equation (54) pushes the trajectory back to the domain once the trajectory hits the boundary. Under mild conditions, the solution to the reflected ODE exists and is unique (Xie et al. 2024).
Empirically, the reflected ODE can be solved by numerical solvers, for example, using the Buler method. The Euler method iteratively compute the trajectory points by
Ee+61 = 2, + öt u(r,, t),
(56)
where t E [0, 1) is an intermediate time and öt > 0 is small time step. Euler step with reflection can be computed as
6t xt + 56 (u(x,, 1) - - 2(1 (u)(ee 1))² nacnax
(57)
29

---

Published as a conference paper at ICLR 2026
Algorithm 4 Euler method with reflection
Input: velocity field v(2,, f), domain C, an initial point x E C, number of steps N, a function that gives outward normal vector (a)( and a function that computes the distance to boundary sac Uz),
1: Let 01 01:001111 = z0 = I
2: for i = 00, 1, ..... .N do
3:
let f = D2 u(z; t) update J,-1 I; + öt
4:
5:
if 1++1 & then
6:
compute distance to boundary 8; = min sae (0c öt v.)) sōl
7:
compute boundary point 24
8:
compute normal vector n, = nac(s1))
9:
compute reflection term LG 20t(1 8;) n,74
10: update Ii+1 = I++1 + Li
end if
12:
end for
Output: EN
Table 9: Closed-form Expressions for Distance and Normal Vector Functions
Constraints
Formulation
Distance Function
Normal Vector
Unit Ball
| 1 - E
s(z, (root(e" u, 2n x 1)}+
n(x) x/||a>||
Unit Cube
s(z, v) 1,,. {[1 sigu(vs)a))/u/
n(x) earg max, 1.2...... [zy
Linear
a' x <
s(x, = (b (b-a" 22'=2/aa"v
n(x) a/|[all
Quadratic
Qu + a" x x
s(x,)) (root(Aq. Bq, Cq)
n(x) -2Qx
Notation: 2,0 R" R...ERR,DERRRAA Re.."e..REEEEE
R PER" Fo.... F. E Rkak
2 e1 (1.0,0..... 0) = (0, 1,000 O)..... en- (0.0. 1,0), = (0, 0).1..
Ag Qu, BQ = 2x Qu v, Co T Qa" +
(+)* = max(:,)) 0).
-22= 1124 denotes the quadratic equation solution
(oot(#1,22,24)
where min {1 inf( & >
| ae + sôt o(x,,)) () & C}} is the distance from He to the nearest boundary aC along direction u( f), and nac is the outward normal vector at 2t + - E aC. Algorithm 4 summarizes the Fuler method with reflection and table 9 lists reflections for several common convex sets.
E.3 PERFORMANCE EVALUATION METRICS
Feasibility ratio computation Suppose that 31, .2....... ....,. UN are generated samples, gı(x), 1.2.... are constraints of interest. One sample x is feasible if
g (x) S 0, 1, 2, ...... m.
(58)
The feasibility ratio of the batch samples is defined as
( X2, 22.,,...., # 1, 2. ..... N. (ax) S 0, = 1,2, ..., m)
(59)
N
In our experiments, we set the tolerance as 10-6
E.4 CONSTRAINT FORMULATIONS IN ILLUSTRATIVE EXAMPLES
We first give the specific formulations of those constraints:
Convex Set:
{sr E R.2 Ax b, | c|12 s 2.5,
+Qu+t"r+dd00
(60)
where A, b, Q, p. d are randomly sampled.
30

---

Published as a conference paper at ICLR 2026
Star-convex Set:
{x RR² |||||| I <Ta, (x) := 1 ais arccanaacctn(22//11
(61)
where (0, and n E 24 determine the size and number of stars, respectively. Geodesic-convex Set:
{ E S Ar - b)
(62)
where S is the sphere in 3-dimensional Euclidean space and is constrained by set of linear inequalities.
The data are sampled from a mixed Gaussian distribution, whose location parameter and covariance are generated randomly.
Model settings: We model the velocity field with 2 hidden layers with exponential linear unit (ELU) activation functions. We train the models for 10000 epochs with batch size of 256, and prior distribution as the uniform distribution over the constrained/ball domain. Samples are generated by the Euler algorithm 1,000 steps. For the reflected/projected methods, additional reflection/projection performed after each step.
E.5 MIRROR MAP GAUGE MAP ON SIMPLEX DOMAIN
We compare our GFM model against several baseline methods: Vanilla Flow Matching, DDPM, and Mirror Diffusion. We consider the standard simplex in d dimensions:
Sa = (x E RRd x4 2 O,
Σ rs S 1).
(63)
Data preparation: We generate training data from a mixture of multivariate normal distributions, with each component centered near distinct geometric feature of the simplex. The mixture includes vertex-centered modes, one origin-centered mode, and one face-centered mode. The location parameters defined as:
. 0.1/d.0.9.0.1.. 14 001/440..../00001144....00 0.1/d), i= 1,,,2....,4
d-i-1
11 (0.1/d. .../d......111... 0.1/d):
My (0.9/d.... 0.9/d).
Here, μ4 places most of its mass at the i-th vertex, Ht, is centered near the origin, and μ distributes mass in the simplex face. We generate total of 10,000 training samples from this mixture.
Model settings: For illustrative 2-dimensional cases, we modeled the time-variant velocity field using layers with 256 units each and ELU activation functions. For c0-dimensional cases, we modeled the time-variant velocity field with 4 hidden layers of 512 hidden units, incorporating residual connections and a bottleneck structure (Liu et al. 2024a).
As shown in Table TO in the synthetic simplex domain example, our GFM method achieves compa- rable performance to the vanilla flow matching baseline and other constrained generation baselines in low-dimensional settings (d 2). while achieving significantly faster inference speed. In high- dimensional settings (d = 50), GFM attains the best fidelity to the data distribution while maintaining faster inference speed. It is worth noting that the Mirror Map model (Liu et al.) 2024b) exhibits fast performance in both training and inference phases. This is because a low-complexity, closed-form mirror map exists for this simple simplex domain. However, the mirror map approach fails for more complex constraints, such as the PSD cone constraint.
E.6 CONFIGURATION GENERATION IN ROBOTIC CONTROL BENCHMARK
We follow the procedure reported by (Jaquerr et al. 202T) to learn the trajectories of the manipulability ellipse. In planar letter drawing problems, the manipulability ellipses are modeled by SPD matrices
31

---

Published as a conference paper at ICLR 2026
Table 10: Results for synthetic generation tasks the simplex domain.
Method
FM
DDPM
Mirror Map
GFM
Feasibility (%)
94.5
93.8
100
100
MMD (x 10
6.26
14.02
8.34
6.53
Training
0.09
0.09
0.13
0.13
Inference (s)
0.94
0.46
0.97
Feasibility (%)
4.35
100
MMD (x10-
9.08
1.82
6.00
2.56
d 50
Training (s)
0.14
0.13
0.14
0.15
Inference (s)
0.93
1.20
0.93
1.38
(a) Data
(b) FM
(c) DDPM
(d) Mirror Map
(e) GFM
= [X E g2x2 X > 0, tr(X) C}. In our formulation, we transform it into the equivalent
linear matrix inequality formulation:
M = f 3
21
0
+ 22
0
1 + 23
(
0 0 = 0. tr(M) = x1 + xy I C 1
(64)
We additionally learn to sample the 2-dimensional trajectories. Therefore, our models are parameter- ized to generate samples in R², x M
Model settings: We model the variant velocity field with hidden layers with 256 units each and ELU activation functions. We train the models for 10,000 epochs with batch size of 256, and prior distribution as the uniform distribution over the constrained/ball domain.
Remark:
The feasibility ratios reported in Table 11 also high for methods that do not force feasibility, e.g., vanilla flow-matching and diffusion models. The possible reason for that phenomenon that the training data is highly concentrated in the interior of the domain, and only few data points are close to the boundary. Therefore, few generated samples are infeasible.
E.7 SOLUTION GENERATION FOR RELAXED COMBINATORIAL PROBLEMS
We further evaluate our approach on high-dimensional solution generation for relaxed combinatorial problems (Kook & Vempala 2024). The distribution we consider is particularly interesting recovers, as special cases, the sampling problems associated with Max-Cut SDP relaxations and minimum volume ellipsoid problems, defined as:
X exp (-((A,X) + ||X B
||X
C|l
log det X))
(65)
s.t. X z 0, XD 24, We
(66)
32

---

Published as a conference paper at ICLR 2026
Table 11: Results for robotic manipulability ellipse generation task.
Methods
FM
DDPM
GFM
Feasibility (%)
99.92
99.94
100
MMD (x10-)
5.39
5.58
9.85
Training
3.43
3,64
3.55
Inference (s)
0.50
0.86
0.83
(f) Data
(g) FM
(h) DDPM
(i) GFM
Figure 11: Visual comparison of 1000 samples between data distribution, DDPM model, Flow
Matching model, and our GFM model on robotic manipulability ellipse generation task.
The constraints ensure that X is positive semidefinite n x n matrix satisfying m linear inequalities. Each term in the energy function captures different structural aspects of the optimization land- scape, including linear objectives, quadratic penalties, proximity constraints, and determinant-basse regularization.
By developing generative models for this class of problems, we aim to provide tools that not only sample feasible solutions but also help explore the solution space in ways that could reveal insights about the underlying combinatorial structure. This capability is particularly valuable for understanding the landscape of near-optimal solutions and for generating diverse candidate solutions that might be refined by downstream processes.
Data preparation: We sample from the target distribution using the Hit-and-Run (Bélisle al. 1993) algorithm. The Hit-and-Run sampling algorithm a Markov Chain Monte Carlo (MCMC) method that generates random samples from a high-dimensional distribution by iteratively selecting random direction and moving to new point along that direction according to the target distribution. To generate training samples, use the Hit-and-Run sampler with 1000 burn-in, and sample new points using I-dimensional Metropolis sampling with 100 burn-in.
Model settings: We model the time-variant velocity field with hidden layers with n²/2 /2 hidden units with residual connections and bottleneck structure (Liu et al.) (202a)). We train the models for
1 1000 epochs with batch size of 256, and prior distribution as the uniform distribution over the constrained/ball domain.
E.8 CONSTRAINED CONDITIONAL TIME SERIES GENERATION
We follow the experimental setup of (Narasimhan al. 2025 to generate time-series data subject to physical and statistical constraints:
Emin
Fmax
Mmin S (wee(11.)) S Mmax
var S max
(67)
capacity bounds
volume bounds
variance bounds
These constraint bounds derived from empirical statistics of the historical dataset. The time series generated via conditional forecasting, using past observations as model input.
Data preparation: We the PEMS-BAY traffic dataset (Yoon et al. 2019; Li et al.) 2018). partitioned by days. The models learn to generate time series for one day conditioned on data from the previous day. Days with incomplete data excluded from the dataset.
33

---

Published as a conference paper at ICLR 2026
Model settings: We model the conditional time-variant velocity field using 4-layer neural network with 64 hidden units per layer, incorporating residual connections and a bottleneck structure (Liu et al. 2024a). All models are trained for 4,000 epochs with batch size of 180. For the prior distribution, we use a standard Gaussian for vanilla Flow Matching and projection-base methods, and uniform distribution over the ball domain for GFM.
Part of PEMS-BAY time scrics
2017-01-13
2017401-17
201 7-0--11
201 7-01-5
7-00-20
2017-02-42
2017-00-06
Vasilla
Profect
CFPM
2017-01-05
201 7-01-11
201 3011-01-28
2011-42-27
2017-04-12
201 7-05-66
201 7-06-122
Figure 12:
Visual comparison of 24 time series strides between data distribution, Vanilla Flow Matching model, Projection Flow Matching model, and our GFM model constrained traffic time series generation task.
E.9 WATERMARKED IMAGE GENERATION TASKS
We evaluate GFM constrained image generation with embedded watermarks using CIFAR-10, which consists of 32 32 RGB images. Following the watermarking methods in (Liu et al.) 2024b) we partition pixels x E [0, 1]d into unconstrained public pixels X1 and watermark pixels X2 C. where C is randomly sampled polytope for watermarking.
Model settings: All baselines are distilled from a publicly available checkpoint flow matching models (Tong et al.) 2024)), which share the same U-Net architecture (34M parameters), All models use identical distilled hyperparameters (100 epochs) and sampling settings (100-step Euler method) on CIFAR-10 10 for fair comparison. We follow the method in (Salimans & Ho 2022 to train the distilled models by minimizing the MSE from the teacher model:
Edistin ( = =B (lo (200 t) -o (x)) I)
(68)
where v, is obtained by integrating teacher ODE from Zo with step length 0.001.
Data preparation: We select 32 pixels as watermarks constrained via random polytope C = {xg Ax < b) with 0. The watermarked training dataset is created by projecting the selected pixels in the original CIFAR-10 dataset onto the polytope C, while other public pixels remain unchanged.
34

---

Published as a conference paper at ICLR 2026
(a) FM
(b) Projection
(c) Reflection
(d) GFM
Figure 13: CIFAR-10 samples with embedded watermarks generated by various methods using 100-step Fuler sampling.
35

---

Published as a conference paper at ICLR 2026
E.10 SETTING OF SCALABILITY TESTS OF GAUGE MAPPING
To evaluate the scalability of gauge mapping, we conducted experiments with randomly generated inequality constraints. We tested scalability across three dimensions; problem dimension, batch size, and number of constraints.
We established baseline configuration with 500-dimensional variables (n 500), constraint dimension & = 500, 8 constraints, and batch size 256. For each scalability test, we varied one parameter while keeping the others fixed at their baseline values:
Problem dimension scalability: We varied n E [10, 3000] while fixing k 500, m = 8, and batch size 256. The specific dimensions tested follow Table 8
Batch size scalability: We varied the batch size from 10 to 100,000 while fixing n 500, k = 500, and = 8.
Constraint scalability: We varied the number of constraints m E [8, 512] while fixing n 500, k = 500, and batch size 256.
For each configuration, we measured the computation time of the gauge function evaluation. The results are illustrated in Fig. g
To demonstrate the efficiency of the bisection algorithm (Algorithm 3). tested gauge computation on convex high-dimensional polynomials generated using the sum- sum-of con-oof-quurees (SOS) method.
-um-of-qquares polynomial of n variables and degree d defined as
p( 22,
Fn)) = (mono,(I), ,,, x,))) Q: mono I2...., Fr,,
(69)
where mono,( 22, Fnn) a vector of all monomials up to degree d, and Q positive semi- definite matrix. Each monomial of degree k takes the form xt 2 ... x.f where C1. ... Cn. and C1 + C2 + ..... Cn k.
NUO
We performed the bisection-bais -based gauge function, and the results are summarized in Table 12
Table 12: Gauge function computation time of high-dimensional polynomials using bisection
No. of polynomials No. of variables
Degree
No. of monomials
Time (s)
10
4
4,356
0.025
10
50
4
7,8, 276
276
0.492
10
100
4
26, 532, 801
5.592
E. ABLATION STUDY
We consider the following ablation study to examine key components of GFM, including
Impacts of the interior point selection: Central interior point vs Near-boundary interior point
Impacts of different generation strategies: eeffection-bssdd sampling vs Projection- based sampling
Impacts of the initial set: fy-oorm ball loc -norm cube
We follow similar setting in Sec. 7.1
36

---

Published as a conference paper at ICLR 2026
Table 13: Performances on joint convex under different settings.
L
Loe
Central
Boundary
Central
Boundary
MMD 10 3)
0.34
43.50
0.35
43.59
Reflection
Training (s)
0.22
0.22
0.21
0.23
Inference (s)
0.91
0.86
0.86
0.88
MMD (x (x10-))
3.14
37.82
3.34
38.13
Projection
Training (s)
0.22
0.22
0.21
0.23
Inference (8)
0.48
0.50
0.47
0.50
LetCentrl
L Boundary
CanCaal
L. Boundryy
Figure 14: Visual comparison in the transformed space over L2-norm ball and Lec-norn cube: We compare gauge mappings using central versus near-houndary interior points, and evaluate samples generated via projection versus reflection-basee central interior points reduces distortion in the transformed space, facilitating more effective GFM training. Reflection based sampling achieves superior generation quality compared to projection methods, while the choice of initial set shows negligible impact generation quality.
strategies. Results show that gauge mapping with
37


Reviews:
[{"review_id": "1HH2Hi34wu", "reviewer": "Reviewer_FHKh", "summary": "GFM enforces hard constraints by bijective mapping any compact convex support set  \"C\"   to the unit ball B via a gauge map,training and generating inside B (with cheap boundary reflection) and mapping back, so samples are strictly feasible by construction. The gauge map is bi-Lipschitz, keeping distortion bounded and preserving distributional regularity when transferring training/sampling between \"C\"  and B. The framework extends to star-convex and geodesic-convex sets via appropriate generalized gauge maps. Experiments show 0% violations, strong distributional fit (e.g., low MMD), and faster generation than prior constrained methods.", "questions": "1. How do the authors envision extending GFM to more general non-convex domains beyond the star-convex and geodesic-convex cases?\n\n2. The paper mentions that extending GFM to discrete generation via relaxation or embedding is a promising direction. Could the authors elaborate on how this might be achieved?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The author proposes a generalized gauge mapping for constrained generation on any compact convex.The map Φ:B^2↔C is bijective and bi-Lipschitz with explicit bounds, which yields a clean distributional error transfer for any p-Wasserstein distance(Proposition 4.1),so training in B_2 preserves accuracy up to a bounded factor; feasibility is strict by construction.Extensions to star-convex and geodesic-convex sets (Appendix A) further broaden scope.\n\n2. Computational efficiency:overall cost ≈ unconstrained models + small mapping overhead.Gauge values are closed-form for many sets (linear/quadratic/Second Order Cone/...); for general convex sets a 1D bisection along rays computes boundary intersections rapidly. The approach scales to high-D and complex constraints.\n\n3. The empirical evaluation is convincing: 0% violations, low MMD (often outperforming projection/reflection/Metropolis baselines), and faster inference —showing constraints do not degrade sample quality.\n\n4. The paper is clearly written:clear positioning (comparative table), intuitive figures for the gauge map, and precise propositions.", "weaknesses": "1. Sensitivity to interior point & data assumptions:Robustness still hinges on a good interior point, which can be non-trivial in high dimensions; moreover, requirement of the data distribution may limit real-world performance.\n\n2. Although many cases admit closed-form or 1D bisection, per-sample gauge evaluation could dominate for very high-D or costly oracles (large LPs/implicit constraints),as shown in Figure 6.\n\n3. Reproducibility:Many evaluations use programmatically generated target distributions tied to the constraints. While standard for this setting, the paper would be stronger with public tasks + released seeds/code, plus tests on real constrained domains (e.g., PSD cones, Birkhoff polytope, trajectory constraints).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5qKNmWOXuD", "reviewer": "Reviewer_B79x", "summary": "The paper introduces Gauge Flow Matching, a constrained sampling approach that generalizes to arbitrary convex sets. The methodology builds on prior work for mirror mappings and reflected diffusion, originally proposed within diffusion frameworks, providing a generalization for flow matching with broader coverage of convex constraints. The experimental analysis on synthetic settings reports promising results for improving fidelity to the training distribution while improving runtime over existing constrained sampling approaches.", "questions": "- In Tables 3 and 5, the training time and inference times are lower than the \"vanilla\" models. Can the authors explain why this would be?\n\n- Along similar lines as the previous question, the runtime for the reflection and projection algorithms (Tables 3 and 4) seem unexpectedly high. Can the authors speak to how these were implemented? Given the simple nature of these constraints, it would seem these operations could be more efficiently computed (e.g., an efficient closed-form projection operator).", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "- **Motivation and Articulation:** The paper's motivation is well established and articulated. The writing is fairly easy to follow, and the methodology is communicated well.\n\n- **Theoretical Support:** The theoretical grounding presented in the paper strengthens the claims made by the authors. The analysis appears to be sound, and the provided bounds contribute to the soundness of the paper.\n\n- **Run-time Efficiency:** The efficiency arguments are an interesting contribution. While I believe the authors could make a better case for this by motivating it with time-sensitive, real-world experimental settings, I believe there are some relevant applications where this could be applied.", "weaknesses": "- **Limited Contribution:** The novelty of the method seems to be directly tied to the application of the existing gauge mapping [1] within flow matching models. The actual integration of this into the flow matching model is very closely tied to prior work on mirror maps and reflection based techniques. While the authors do point out that mirror maps do not generalize to arbitrary convex sets, it seems that beyond this mapping \"substitution\", the contribution is fairly limited.\n\n- **Limited Empirical Evaluation:** While the experimental evaluation provides strong results, it relies solely on synthetic test cases with little to no real-world meaning. Furthermore, for many of these test cases, rejection sampling seems to be the better approach, as feasibility rates are higher than 90%, and samples have higher fidelity to the training set. Hence, the results would be much more compelling if they were applied to real-world settings where constrained sampling would be practically applied; this seems to be a missed opportunity, especially if such settings could provide a compelling case for the introduced efficiency. \n\n- **Selected Baselines:** It is surprising that the authors have not chosen to compare to mirror diffusion models [2] in any of their analysis. Given how closely the method overlaps with mirror mappings, the omission of this comparison makes it difficult to assess whether the gauge map is effective -- which is the core contribution of the work. Furthermore, this would be a valuable comparison to assess how different bijective mappings influence the generation quality. \n\n---\n\n[1] Tabas, Daniel, and Baosen Zhang. \"Computationally efficient safe reinforcement learning for power systems.\" 2022 American Control Conference (ACC). IEEE, 2022.\n\n[2] Liu, Guan-Horng, et al. \"Mirror diffusion models for constrained and watermarked generation.\" Advances in Neural Information Processing Systems 36 (2023): 42898-42917.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tz7tPPJ73j", "reviewer": "Reviewer_rKAZ", "summary": "This paper addresses a critically important and timely challenge of constrained generative modeling by proposing a novel idea called gauge flow matching (GFM). The underlying idea of GFM is to introduce a bijective gauge mapping to transform a complex constrained generative modeling problem into an equivalent modeling problem over a simple unit ball domain. The paper presents rigorous theoretical validation/analysis on the strict constraint guarantee and a lower computational complexity on not only convex feasibility regions but also certain types of non-convex feasibility regions. It also presents empirical validation of the GFM idea on various benchmarks, in comparison to other state-of-the-art constrained generation methods. The proposed idea is technically sound, and the scientific breakthrough is clear. The presentation and the style of writing are also effective, maintaining a good balance between high-level intuitive explanation and low-level mathematical details. As such, I believe this paper deserves presentation at ICLR.", "questions": "I don't have questions.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Although the idea of bijective mappings to transform constrained domains into a new domain that has more favorable analytical properties existed before this paper, it was still understudied, and the previous work was largely limited to some limited problems and/or theoretical guarantees. Hence, I find the idea in this paper sufficiently original and of significance, presenting creative combinations of existing ideas (e.g., guage mapping + bijective mapping to simplify constrained generation), new scientific knowledge (e.g., theoretical guarantees, complexity analysis, benchmark test results), and such.\n\nI also appreciate the style and quality of writing. While engaging a broader group of audience effectively with a clear and intuitive explanation of the method and the rationale behind the method, the paper still presents enough technical details with mathematical rigor.", "weaknesses": "First and foremost, I would've appreciated more experiments on challenging generative tasks. I understand that the authors did their best to cover the spectrum of discussions from theoretical to empirical validations, the real-world generative tasks (i.e., robotic manipulability and relaxed combinatorial problems), but I still feel underwhelmed when it comes to the empirical validation portion of this work. Image generation, for example, would've been a great benchmark test, even if the scope of images was limited to a small, limited variety (e.g., microscopic images as opposed to the whole set of natural images). I understand such a task might be \"too non-convex\" and might be beyond the scope of this work, but I still remain curious to know if and how far the proposed GFM idea could push the limit. (Especially because the other parts of the works are very well done!)\n\nOther than that, I only have some minor suggestions regarding mostly presentation, which the authors may or may not want to consider:\n- I think Figure 1 captures the core idea of your work very well. Why don't you move it to an earlier page?\n- Also, figure captions in general could be more elaborate and contain more details.\n- Excessive use of bold faces and color highlights hurts readability in my opinion. As an example (out of other such cases), I don't see the need to emphasize the acronym \"GFM\" and the word \"limitation\" in the conclusion section by making them boldface. This tendency is found across the whole paper, which the authors are strongly encouraged to revisit and think through.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)