Beyond Explicit Mirror Maps: Constrained Generative Modeling via Brenier Maps

Mirror-lifting paradigms for constrained generation are concise and efficient, but they suffer from fundamental limitations. Analytical mirror functions are only val-id for specific convex constraint domains, lacking a unified construction method and failing to handle non-convex constraints. Neural approximate mirror functions broaden applicable scenarios yet brings uncontrollable geometric distortion and heavy training overhead. To address these issues, we propose a unified framework for constructing constrained generative models based on Brenier maps. Theoretically, we prove that mirror functions on convex domains constitute a special case of dual Brenier functions. Leveraging singular set theory, we break the convexity constraint of mirror maps and enable unified adaptation to arbitrary convex and non-convex domains. Algorithmically, we develop a universal multi-stage solving strategy for Brenier maps. Anchor sampling and geometric variational principles are adopted to capture precise geometric structures, and a controllable smoothing mechanism is introduced to preserve critical geometric properties and improve generalization. Experimental results show that on three classical convex constraint domains, our method achieves better fitting accuracy for mirror functions than neural approximation baselines. For the watermark image generation task and five other complex real-world constrained generation tasks involving both convex and non-convex scenarios, our method surpasses existing SOTA baselines in generation quality and constraint satisfaction with lower computational cost, substantially improving the accuracy and efficiency of constrained generation.