Neural Approximate Mirror Maps for Constrained Diffusion Models

Thank you for the positive feedback.

A hyperparameter sweep for \sigma_\max and $\lambda_\text{constr}$ can be used to determine the values that would provide the best performance in terms of constraint distance and distribution-matching accuracy. As detailed in Appendix C.1, we did not extensively search the hyperparameter space and simply looked at \sigma_\max=0.1, 0.5 and $\lambda_\text{constr}=0.01, 1$.

Intuitively, it would make sense that \sigma_\max should be increased with a higher dimensionality: because data points are more separated in higher dimensional spaces, adding more noise would ensure the maps are trained accurately at regions between the data points. We did not find a great deal of hyperparameter sensitivity with the dimensionalities considered in our experiments. We believe that a related theoretical direction to pursue is to understand how the best possible constraint satisfaction depends on the complexity of the constraint and how that might inform the choices of $\lambda_\text{constr}$ and \sigma_\max.

We believe that the performance gap for the semantic constraint has to do with the gradient-based nature of our approach. One avenue we are exploring is the use of gradient-free methods to optimize the NAMM, which could better handle constraints with irregular or undefined derivatives.

Thank you for the suggestions for improving the text; we will revise the text accordingly.

Thank you for your feedback and suggestions.

Robustness of our method is enhanced by introducing noise in the mirror space. By robustness, we mean that the inverse mirror map is able to restore a wider region of $\mathbb{R}^d$ to the constraint set. A robust inverse map works not just for points from the data distribution but also for points somewhat off the data manifold. Figure 6 demonstrates the robustness provided by introducing noise in the mirror space. When the noise level \sigma_\max is too low (i.e., \sigma_\max=0.001), the constraint distances are higher. This indicates that training the NAMM with too little noise in the mirror space makes it difficult for the inverse mirror map to handle any errors introduced by the mirror diffusion model.

We performed ablation and inverse problem experiments on a subset of the considered constraints for the sake of presentation clarity. We chose to focus on the more complex physics-based constraints to highlight the applicability of our method to physics-constrained inverse problems. We would be happy to include experiments on all five constraints in an appendix.

To our knowledge, there is no other approach that is applicable to all the constraints we consider. Since there was no relevant baseline, we instead considered a modification of the DPS method to act as a baseline we call “constraint-guided DPS.”

We do not consider finetuning to be an essential component of the method. As Figure 4 and Table 1 show, most of the performance is already achieved before finetuning. We suggest finetuning as an optional additional step to further boost the constraint satisfaction. We believe it is a strength of our method that the results are not so dependent on finetuning.

Thank you for the positive feedback and clarifying questions.

We would like to clarify that all of our experiments are conducted on image datasets. For example, the divergence-free constraint is demonstrated on 128x256 images. The NAMM architecture can easily scale to even higher dimensions, and we expect the bottleneck to be the diffusion model itself, which is known to be slow for generating very large images. To overcome this bottleneck, it is possible to train a latent diffusion model in the mirror space, where the latent space would be lower dimensional and thus cheaper to sample in.

Our method can be easily applied to the constraints considered in Reflected Diffusion Models [1], which are convex and have simpler analytical forms in comparison to the physics-based and semantic constraints we consider in the paper. For example, we can easily design a constraint distance for our method to enforce the constraint that pixel values are bounded between two values.

The regularization loss ensures uniqueness of the forward and inverse mirror maps. Just the fact that the forward map is invertible does not ensure uniqueness: as a simple example, the forward map $f(x) = x$ has the inverse $g(x) = x$, but it could be scaled to $f’(x) = 2x$ and have the inverse $g’(x) = 0.5x$. The regularization loss helps resolve this scale/shift degeneracy.

References:

[1] Aaron Lou and Stefano Ermon. “Reflected Diffusion Models.” ICML 2023.