Research map / Image generation / Fast diffusion sampling
Neural Approximate Mirror Maps for Constrained Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes neural approximate mirror maps (NAMMs), learned forward/inverse maps that transform constrained data into an unconstrained space so a mirror diffusion model can be trained there and its samples projected back to satisfy general, possibly non-convex constraints.
Problem
Diffusion models often fail to satisfy subtle constraints (physics, geometric, semantic) present in training data, and existing constrained diffusion methods only handle restricted constraint types such as convex sets or symmetry groups.
Approach
Learn an approximate mirror map g_phi (gradient of a strongly input-convex neural network) and an approximate inverse f_psi (ResNet CNN) using only a differentiable constraint distance, with cycle-consistency, constraint, and regularization losses over noisy mirror-space distributions; then train a standard diffusion model in the mirror space and map its samples back with f_psi.
Key results
- Finetuned NAMM-based MDM reduces mean constraint distance relative to the vanilla DM by 78% (total brightness), 91% (1D Burgers'), 77% (divergence-free), 96% (periodic), and 38% (count).
- In data assimilation, mirror DPS yields notably less PDE residual / divergence than DPS and constraint-guided DPS on Burgers' and Kolmogorov flow.
- Mirror VAE (MVAE) trained in the learned mirror space achieves lower divergence than a vanilla VAE on the divergence-free constraint while also improving MMD and KID.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train a NAMM-based MDM instead of a vanilla diffusion model | Diffusion model · Total Brightness | CD (↓): 2.20 ± 1.61 -> 0.49 ± 0.35 | Helped |
| train a NAMM-based MDM instead of a vanilla diffusion model | Diffusion model · 1D Burgers' | CD (↓): 0.45 ± 0.03 -> 0.04 ± 0.01 | Helped |
| train a NAMM-based MDM instead of a vanilla diffusion model | Diffusion model · Divergence-free | CD (↓): 6.96 ± 0.27 -> 1.57 ± 0.31 | Helped |
| train a NAMM-based MDM instead of a vanilla diffusion model | Diffusion model · Periodic | CD (↓): 1.04 ± 0.08 -> 0.04 ± 0.06 | Helped |
| train a NAMM-based MDM instead of a vanilla diffusion model | Diffusion model · Count | CD (↓): 0.56 ± 0.44 -> 0.35 ± 0.27 | Helped |
| add inverse-map finetuning (ablation) | NAMM-based MDM · Total Brightness | CD (↓): 0.57 ± 0.48 -> 0.49 ± 0.35 | Helped |
| add inverse-map finetuning (ablation) | NAMM-based MDM · 1D Burgers' | CD (↓): 0.09 ± 0.01 -> 0.04 ± 0.01 | Helped |
| add inverse-map finetuning (ablation) | NAMM-based MDM · Divergence-free | CD (↓): 2.51 ± 0.21 -> 1.57 ± 0.31 | Helped |
| add inverse-map finetuning (ablation) | NAMM-based MDM · Periodic | CD (↓): 0.04 ± 0.05 -> 0.04 ± 0.06 | No clear effect |
| add inverse-map finetuning (ablation) | NAMM-based MDM · Count | CD (↓): 0.71 ± 0.53 -> 0.35 ± 0.27 | Helped |
| replace the ICNN-gradient forward map with a ResNet-based forward map (ablation) | NAMM · Demonstrated constraints | constraint distance | Hurt |
| train the inverse map using noise levels integrated over [0, σ_max] instead of fixed σ=σ_max (ablation) | NAMM for 1D Burgers' · 1D Burgers' | constraint distance | Helped |
Surprising findings
- Generic diffusion models perform surprisingly poorly even on the simple total brightness constraint.
- Using a ResNet-based forward map instead of an ICNN gradient leads to worse constraint satisfaction, even though ICNN theory applies only to convex constraints.
Limitations
- Constraint satisfaction is approximate; the inverse map is not a true inverse and constraint distance is not exactly zero for count and other constraints.
- Less effective on highly non-convex constraints such as the learned count constraint, where the gap to the DM is smaller.
- FID and KID rely on Inception features tuned to natural images, which are imperfect for physics simulation data.
Experimental setup
Datasets: black-hole simulation images 1D Burgers' Crank-Nicolson solutions 2D Kolmogorov flow (jax-cfd) periodic unit-cell tilings simulated radio galaxy images
Benchmarks: total brightness 1D Burgers' divergence-free Kolmogorov flow periodic count data assimilation tasks
Metrics: constraint distance MMD KID FID
Baselines: vanilla diffusion model (DM) constraint-guided DPS (CG-DPS) DPS VAE
Builds on: mirror diffusion models diffusion posterior sampling input-convex neural networks CycleGAN NCSN++
Closest papers by meaning
- Composition and Alignment of Diffusion Models using Constrained Learning (NeurIPS 2025 · similarity 0.87)
- Constrained Diffusion Models via Dual Training (NeurIPS 2024 · similarity 0.86)
- Constrained Diffusion with Trust Sampling (NeurIPS 2024 · similarity 0.86)
- SITCOM: Step-wise Triple-Consistent Diffusion Sampling For Inverse Problems (ICML 2025 · similarity 0.86)
- A Variational Perspective on Solving Inverse Problems with Diffusion Models (ICLR 2024 · similarity 0.86)
- Fast constrained sampling in pre-trained diffusion models (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.