Neural Approximate Mirror Maps for Constrained Diffusion Models

Authors: Berthy Feng, Ricardo Baptista, Katherine Bouman
Primary area: generative models
Keywords: generative models, diffusion models, mirror maps, constrained generation, inverse problems

Abstract:
Diffusion models excel at creating visually-convincing images, but they often struggle to meet subtle constraints inherent in the training data. Such constraints could be physics-based (e.g., satisfying a PDE), geometric (e.g., respecting symmetry), or semantic (e.g., including a particular number of objects). When the training data all satisfy a certain constraint, enforcing this constraint on a diffusion model makes it more reliable for generating valid synthetic data and solving constrained inverse problems. However, existing methods for constrained diffusion models are restricted in the constraints they can handle. For instance, recent work proposed to learn mirror diffusion models (MDMs), but analytical mirror maps only exist for convex constraints and can be challenging to derive. We propose *neural approximate mirror maps* (NAMMs) for general, possibly non-convex constraints. Our approach only requires a differentiable distance function from the constraint set. We learn an approximate mirror map that transforms data into an unconstrained space and a corresponding approximate inverse that maps data back to the constraint set. A generative model, such as an MDM, can then be trained in the learned mirror space and its samples restored to the constraint set by the inverse map. We validate our approach on a variety of constraints, showing that compared to an unconstrained diffusion model, a NAMM-based MDM substantially improves constraint satisfaction. We also demonstrate how existing diffusion-based inverse-problem solvers can be easily applied in the learned mirror space to solve constrained inverse problems.

Decision: accept (poster)
Reviewer scores: 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730137595368

thread_type:
official_review

<thread_object>
creator: reviewer_ftxm

<content>
summary:
This paper proposes a technique to learn a mirror map for general constrained diffusion generation. The resulting neural approximate mirror map transforms the constrained problem domain to an unconstrained space, where the diffusion model is trained. The training loss encourages that the inverse of the mirror map lies within the constraint set. Numerical experiments on various problems demonstrate the efficacy of the proposed technique in enforcing the constraints.

strengths:
- The method is well motivated. Proper constraint satisfaction is challenging in diffusion generation, limiting many important applications in engineering, physics and computer vision. 

- The proposed approach is sensible, and to the best of my knowledge novel. It improves the flexibility of mirror diffusion models by obviating the need for analytical mirror maps.

- The experimental results are promising on a wide array of problems. 

- The paper overall is well-written and easy to follow.

weaknesses:
- It is unclear how complex of a constraint the proposed method can handle, as the gap between NAMM and the baseline is less significant for the semantic problem. 
- It is unclear if there is a systematic way to tune the introduced hyperparameters, and how sensitive the performance is in higher dimensions to $\sigma_{max}$.

questions:
- It appears in Fig. 6 that constraint distance is fairly sensitive to $\sigma_{max}$ and $\lambda_{constr}$. Can authors propose systematic ways to tune the hyperparameters? How does the sensitivity depend on the dimensionality of the problem or the complexity of the constraint?
- Minor comment: the citations in Section 2.1 take up a large portion of the paragraph somewhat reducing readability (top of page 3).
- Minor comment 2: Fig. 2 is never directly referenced in the text.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730557287311

thread_type:
official_review

<thread_object>
creator: reviewer_isd2

<content>
summary:
This paper presents neural approximate mirror maps (NAMMs) to enforce soft constraints for diffusion models (DMs). NAMMs employ two neural networks to learn a mirror map that maps constrained points into the mirror space, and its inverse that transforms data back to the constraint set. A mirror diffusion model (MDM) can be trained in the learned mirror space, and its generated samples can be mapped to the constraint set via the inverse map. This method is tested on five benchmark problems, ranging from physics-based, geometric to semantic constraints, and the results show the proposed method improve constraint satisfaction compared to a vanilla unconstrained DM. And, this paper also demonstrates NAMMs leads to less constraint violation when solving constrained inverse problems.

strengths:
1. NAMMs generalize the concept of true mirror maps to learn approximate mirror maps to handle non-convex constraints.
2. NAMMs can handle physics-based, geometric and semantic constraints, while existing methods are restricted in the types of constraints they can handle.
3. NAMMs not only help diffusion models, but also help VAEs to improve constraint satisfaction, showing the potential to be compatible for other generative models. And NAMMs are also helpful to diffusion-based inverse-problem solvers for solving constrained inverse problems.

weaknesses:
1. Theoretically, NAMMs lack the guarantee of the existence and uniqueness of the mirror maps when applied to non-convex problems.
2. The proposed method is validated on five benchmark problems in the main experiments to show the superiority of applying NAMMs to diffusion models in generating constrained data. However, in the experiments to solve inverse problems, ablation experiments, and the experiments applied to VAE, this method is only carried out on partial problems and does not fully demonstrate its performance on the three types of constraints mentioned.
3. Finetuning is an important part of the proposed method introduced in section 3. But, it is mentioned in subsection 4.1 that “We show results from a finetuned NAMM, but as shown in Section 4.3, finetuning is often not necessary”. Moreover, in the ablation studies of constraint loss and mirror map parameterization, and experiments about the VAE, fine tuning is not used.
4. The basic unconstrained model is used for comparison, but the comparison with another existing methods dealing with constraints is lacking.

questions:
1.	In this paper, the robustness of the model is enhanced by introducing noise into the mirror space. What does robustness mean here, and are there any experimental results that support the robustness of the method?
2.	In order to fully demonstrate the performance of the method, it is necessary to supplement its experiments on five benchmark problems and an additional baseline model.
3.	If fine-tuning is considered to be one of the important components of the method and one of the contributions of this paper, more experimental support is needed.
4.	On line 1097, there is a clerical error, “a la”.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707824891

thread_type:
official_review

<thread_object>
creator: reviewer_y9gb

<content>
summary:
The paper proposes neural approximate mirror maps for constraint data generation with diffusion models. Compared to typical mirror diffusion models, the mirror map is parameterized by the gradient of ICNN and learned via penalizing the differentiable constraint distance, thereby being applicable to general non-convex constraints. The forward and inverse mirror maps are learned by a combination of cycle consistency loss, constraint loss and regularization loss. Experiments in several settings ranging from physics-based to semantic demonstrate the effectiveness on constraint satisfaction, training efficiency and constrained inverse problem solving.

strengths:
- The method is applicable to more general constraints than previous works.
- The cycle-consistency loss tailored for mirror maps and diffusion models is sound.
- The experiments are conducted on diverse settings, including constrained DPS.
- The ablation studies are comprehensive.

weaknesses:
- The experiments are primarily toy. It is not clear whether the proposed method can scale to high dimensions and apply to domains such as images.

questions:
- The regularization loss is introduced to ensure a unique solution according to the paper. What is the meaning of unique? As the mirror map is parameterized as the gradient of ICNN, the reversibility is already ensured.
- Is the method scalable? For example, can it be applied to the image settings in Reflected Diffusion Models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731550678966

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We thank the reviewers for their feedback and suggestions. NAMMs offer a way to flexibly incorporate more general constraints than possible with previous methods (y9Gb, iSD2), and we demonstrate this with promising results “on a wide array of problems” (FTXm). All reviewers appreciated the broader impact of our approach, as it addresses the challenge of constrained generative modeling for “important applications in engineering, physics and computer vision” (FTXm). They also recognized the broad applicability of our approach to inverse problems (y9Gb, iSD2) and different generative models (iSD2). Overall, our proposed approach is “sound” (y9Gb) and “novel” (FTXm), and it is backed by a “well-written” paper (FTXm) and “comprehensive” ablation studies (y9Gb). We will address individual reviewers’ questions in the individual responses.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper suggests incorporating constraints into diffusion generative models by mapping samples to an unconstrained space, training a generative diffusion process in the unconstrained space, and then sampling by computing the inverse map of samples in the unconstrained space. The mapping into the unconstrained space (mirror map) is defined as gradient of a convex function and is jointly optimized with the inverse (mirror) map to cover the likely sample space in the unconstrained space as well as other desired properties. 

The method is applicable to a wider set of constraints compared to previous works (e.g., non-convex sets in R^d), the optimization of the mirror map and its inverse are solid, and there is a rather extensive experimental section. In terms of weaknesses, it is not clear how scalable the method is to higher dimensions (e.g., using gradients of ICNN and the mirror/inverse-mirror computation), and theoretically the mirror and inverse-mirror maps are not guaranteed to be defined uniquely or actually satisfy forward-inverse relation.