Neural Approximate Mirror Maps for Constrained Diffusion Models

Authors: Berthy Feng, Ricardo Baptista, Katherine Bouman

ICLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental Our experiments show improved constraint satisfaction for various physics-based, geometric, and semantic constraints. We also discuss ablation studies and adapt a popular diffusion-based inverse solver to solve constrained inverse problems. Figure 3 highlights that MDM samples inverted through fψ are much closer to the constraint set than DM samples despite being visually indistinguishable. Table 1: Effect of finetuning. Constr. dist. (CD) = 100λconstrℓ. The improvements in mean CD are...
Researcher Affiliation Academia Berthy T. Feng California Institute of Technology EMAIL Ricardo Baptista California Institute of Technology EMAIL Katherine L. Bouman California Institute of Technology EMAIL
Pseudocode No The paper describes the methods through textual explanation and mathematical formulations, but it does not contain any clearly labeled pseudocode or algorithm blocks.
Open Source Code Yes Code can be found at https://github.com/berthyf96/namm.
Open Datasets Yes We demonstrate with a dataset of 64x64 images of black-hole simulations (Wong et al., 2022)... Using a Crank-Nicolson solver (Crank & Nicolson, 1947) implemented with Diffrax (Kidger, 2022), we numerically solved the 1D Burgers equation... We used jax-cfd (Kochkov et al., 2021) to randomly sample divergence-free, spectrally filtered initial conditions... Using the unit-cell generation code of Ogren et al. (2024), we created a dataset of 30000 training images... we simulated 128x128 images of radio galaxies with background noise Connor et al. (2022).
Dataset Splits Yes The dataset consists of 100000 training images and 100 validation images. The dataset consists of 10000 training images and 1000 validation images. In total, the dataset consists of 10000 training images and 1000 validation images. we created a dataset of 30000 training images and 300 validation images. The dataset consists of 10000 training images and 1000 validation images.
Hardware Specification Yes Each run was done on the same hardware (4 A100 GPUs).
Software Dependencies No The paper mentions specific tools and architectures like 'NCSN++', 'Adam optimizer', 'Cycle GAN', 'Res Net-based CNN', 'Diffrax (Kidger, 2022)', 'jax-cfd (Kochkov et al., 2021)', and 'simple MNIST classifier (8bitmp3, 2023)', but it does not provide specific version numbers for general software dependencies like Python or PyTorch.
Experiment Setup Yes We used the NCSN++ architecture... with 64 filters in the first layer and the VP SDE with βmin = 0.1 and βmax = 20. Training was done using the Adam optimizer with a learning rate of 0.0002 and gradient norm clipping with a threshold of 1. For gϕ, we followed the implementation of the gradient of an input-convex neural network (ICNN) following the implementation of Tan et al. (2023), configuring the ICNN to be 0.9-strongly convex. ... fψ was implemented as a Res Net-based generator with 6 residual blocks and 32 filters in the last convolutional layer. ... The NAMM was trained using Adam optimizer with a learning rate of 0.001 for the divergence-free constraint and a learning rate of 0.0002 for all other constraints. The regularization weight λreg in the NAMM objective (Equation 1) was fixed at 0.001. We used 3 ICNN layers for images 64x64 or smaller and 2 ICNN layers for images 128x128 or larger.