Paper: Isometric Representation Learning for Disentangled Latent Space of Diffusion Models (ICLR-2024)

Reviewer comment (Evaluation):
Several key experiments are highly subjective. For example, the opening example of Fig 1 claims that the proposed approach is better than spherical interpolation. In my subjective view, the opposite seems true. Which method actually is better is not obvious to me.

Author response:
__[W1,6. Subjective experimentation and interpretation]__

We agree that evaluation of qualitative results is in nature subjective, and each person may think differently unless the quality gap is extremely obvious. 

To support quantitatively that our model actually generates a smoother interpolation, we added LPIPS plot between each adjacent frames at Fig. 6 (Blue: Baseline, Orange: Ours). Also, we revised Fig. 1 to present results with more pronounced differences and added Fig. 6 and VII for more examples. We hope these can provide more intuitive demonstration that smoother interpolation can be achieved for two paths sharing very similiar boundary points.

Understanding the visual inspections cannot be a concrete support for our conclusion, we added several more quantitative metric comparisons supporting for how much the encoder has become close to isometry (MCN, VoR) and how much the latentspace is disentangled (LS), in Tab. 1 and 2.

We hope these additional results could resolve your concerns, but we welcome further questions if you have any.

Reviewer comment (Theory):
The paper several times claims that by being isometric all sorts of benefits are gained, e.g. the representation becomes disentangled and reflects semantics perceived by humans. Given that we are not talking about being isometric to a human perception space, then I struggle with the mathematical validity of such statements. I really wish the authors would tone down such claims, as the resulting paper would be much more convincing.

Author response:
__[W3. Isometry in human perception space]__

We appreciate the reviewer for this constructive feedback. Our arguement that we make the mapping from human perception to latent space isometry is based on the implicit assumption that the mapping from H space to human perception is isometric, same as in [1]. However, as the reviewer pointed out, human perception space is just an abstractly defined space which can be highly subjective, and we totally agree with that the argument is not mathematically rigorous. Reflecting the reviewer's comment, we toned down our statement in the revised manuscript (marked blue). Specifically for toning down, we revised the following:

- Abastrct
  - perfectly align --> align close
  - semantically clear and geometrically sound latent space --> semantically and geometrically better latent space
  - marking a significant advance in aligning latent spaces with perceptual semantics. --> suggesting that our method helps to align latent space with perceptual semantics.
- 1. Introduction
  - each dimension in the latent space is aligned --> we could find a linear subspace of the latent space that is aligned
  - which is known to preserve --> which is known to _better_ preserve
- 2.2 Intermediate latent space $H$ as a semantic space
  - Therefore, if we can make the mapping from $\mathcal{X}$ to $\mathcal{H}$ isometric, then $\mathcal{X}$ will be isometric to the data manifold $\mathcal{X}_0$ as well. --> Therefore, we claim that as the mapping from $X$ to $H$ becomes closer to isometric, the mapping of the data manifold from $X$ can also become more isometric.

We have toned down the overly strong or exaggerated climas, but if you still find any parts to be excessively assertive or exaggerated, please feel free to let us know. We will gladly incorporate your feedback.

[1] Kwon, Mingi, Jaeseok Jeong, and Youngjung Uh. "Diffusion models already have a semantic latent space." arXiv preprint arXiv:2210.10960 (2022).

Reviewer comment (Evaluation):
Many conclusions of the paper rely on visual inspection of interpolation plots. I honestly cannot tell which is better (e.g. I think DDPM looks better than the proposed method in Fig 8, but how am I supposed to determine which actually is better).

Author response:
__[W1,6. Subjective experimentation and interpretation]__

We agree that evaluation of qualitative results is in nature subjective, and each person may think differently unless the quality gap is extremely obvious. 

To support quantitatively that our model actually generates a smoother interpolation, we added LPIPS plot between each adjacent frames at Fig. 6 (Blue: Baseline, Orange: Ours). Also, we revised Fig. 1 to present results with more pronounced differences and added Fig. 6 and VII for more examples. We hope these can provide more intuitive demonstration that smoother interpolation can be achieved for two paths sharing very similiar boundary points.

Understanding the visual inspections cannot be a concrete support for our conclusion, we added several more quantitative metric comparisons supporting for how much the encoder has become close to isometry (MCN, VoR) and how much the latentspace is disentangled (LS), in Tab. 1 and 2.

We hope these additional results could resolve your concerns, but we welcome further questions if you have any.

Reviewer comment (Theory):
To me, it is not straightforward how “ISOMETRIC REPRESENTATION LEARNING” is related to “DISENTANGLED LATENT SPACE“.

Author response:
__[W1. Relation between Isometric Representation Learning and Disentangled Latet Space]__

As mentioned in [B], there are various definitions for disentanglement, but a common goal in generative modeling is a latent space that consists of linear subspaces, each of which controls one factor of variation. A latent space is considered ideal if it maximally achieves this.

A key assertion in [1] is that the H-space of the diffusion models are closely aligned with the semantic space. Mathematically, this can be interpreted as the mapping from the H space to the semantic space being close to isometric. Furthermore, it implies that the more isometric the encoder mapping from X-space to H-space is, the more disentangled X-space becomes. In this respect, we aim to achieve the disentanglement of the latent space by regularizing the encoder to be as isometric as possible through optimizing the isometric loss.

To measure how isometric the encoder becomes, we measured mRTL in the original manuscript. Recognizing that this might not be sufficient, we additionally measure mean condition number (MCN) [A] and variance of Riemannian metric (VoR) [A] for further verification. In Tab. 1 in the revised manuscript, we see that the use of isometry loss results in lower mRTL, MCN, and VoR across all datasets, indicating that the encoder has successfully become closer to isometric.

Although we have conceptually explained that achieving an isometric encoder leads to a disentangled latent space, we would like to more intuitively measure disentanglement of latent space. To this end, we adopt perceptual path length (PPL) and linear separability (LS), metrics used by StyleGAN [B] to claim the latent space is disentangled. The results in Tab. 1 and 2 show a significant reduction in these metrics, signifying that our latent space is actually disentangled.

Reviewer comment (Theory):
It seems like the $H$ space in [1] (considered a reliable semantic space by the authors) is obtained from the pre-trained Diffusion Models. If Diffusion Models are trained from scratch with an additional objective, how do the authors ensure that the $H$ space in [1] and the $H$ space in this paper have similar properties?

Author response:
__[W3. Preservation of H space with our method]__

[1] argues that H-space exists in the bottleneck layer of the U-Net, for all pretrained diffusion models, and it is reasonable to deduce that H space exists given that the denoising score matching (DSM) loss has converged. Therefore, it can be inferred that H-space exists if the diffusion loss converges to a similar point even when the isometric loss is added.

Additionally, we can check with the FID scores to see if the diffusion loss has converged well. We observed that the addition of the isometry loss to diffusion loss does not significantly alter the convergence point of the diffusion loss and shows comparable FID scores, which leads to natural inference that H-space will also exist in our diffusion model.

However, as the reviewer suggested, the introduction of an additional objective could potentially alter the property of H-space. Therefore, we have conducted an additional experiment using our model to show that our H-space has similar properties. Please see Appendix E and Fig. VI in the revision.

Reviewer comment (Novelty):
The introduction is not persuasive enough about why the proposed method is needed.

Author response:
__[What is this for?]__

This work is about learning a 'better' or geometrically sound latent space of diffusion models, better aligned with human perception. In contrast to [A], which proposes an indirect manipulation in semantic space, we aim to analyze and improve the latent space by our geometric regularizer, ultimately desirable to the following aspects:

1) __Addressing entanglement in existing diffusion models__: In the discussion section of [C], they directly mention that "_it does not guarantee the perfect disentanglement and some directions are entangled_", and "_We have observed that some of the discovered latent vector occasionally leads to abrupt changes during the editing process_" as their limitations. Our work aims to address these issues. In fact, even a smooth latent traversal has not been possible with existing diffusion models. As seen in [E, F], latent traversal of diffusion models exhibits abrupt changes, compared to that of GANs [G]. Our method has significantly improved this, as illustrated in Fig. 1 and 6 in our revision.

2) __Direct image manipulation and interpolation__ without requiring access to the encapsulated diffusion pipeline. This naturally __reduces computational cost__ in image editing and interpolation to almost half, by dropping asymmetric reverse process which requires two times of entire reverse process. 

3) Geometrically sound latent space improves image __stability and consistency__, as well as __invertibility to its latent code__. StyleGAN [D] applied the path length regularizer to achieve a similar impact on GAN, achieving a near-isometric mapping from latents to images, improving image stability and invertibility, which advances in generation quality, image manipulation. Similarly, our work regularize the mapping from X space to H space to be more isometric, and to the best of our knowledge, this is the first work demonstrating similar impacts on the diffusion model, which was previously challenging due to its iterative sampling nature.


For these reasons, we believe our work introduces a first approach to analyze and learn a 'better' latent space to the diffusion community using tools of Riemannian geometry and can also be generalized into large-scale conditional models such as StableDiffusion, supported by the demonstration given at [H].

Reviewer comment (Writing):
The reviewer understands $H$ as a pre-trained feature space of UNet (i.e., the bottleneck layer). There are several questions regarding the $H$ space that are not clear to me. The authors have mentioned in the Introduction that “However, manipulating attributes indirectly through $H$ is not fully desirable. One reason is additional computations accompanied with this indirect manipulation, as it requires two times of the score model evaluations and training an extra neural network to find local editing directions at every point of $H$.” What does it mean by “as it requires two times of the score model evaluations”? Can it be understood without knowing the paper by Kwon et al. [1]? I also read Section 2.2, but it is still not understandable. Why does the method in [1] require sampling twice? and why does not it happen in the proposed method in this paper?

Author response:
__[Q1-1. Sampling twice]__

According to Theorem 1 in [1], asymmetric reverse process is required for image change, and this requires two independent inference of score model with different inputs: $\epsilon_t(x_t)$ and $\tilde{\epsilon_t}(x_t) = \epsilon_t(x_t|f_t)$, where $f_t$ is the additional module to find the editing direction. With our improved latent space $X$, sampling procedure is identical with that of ordinary diffusion models, e.g., DDIM reverse process, which needs just a single inference of score model. We revised our manuscript with this explanation in Introduction to be self-contained.

Reviewer comment (Theory):
From the part “training an extra neural network to find local editing directions at every point of $H$”, the reviewer assumes it means the method in [1] requires finding a semantically meaningful direction by training an additional module. Why does it not happen in $X$ space while it is needed in $H$ space?

Author response:
__[Q1-2. Why we need additional NN only with H space?]__

If the latent space is disentangled, we can manipulate an attribute on an image by obtaining a corresponding "consistent" direction in the latent space and moving along the direction, without any additional module inference. What we mean by consistent here is that for every latent vector $x$, we call a direction $\Delta x$ "consistent", if $x + \Delta x$ yields an expected image. However, since $X$ in the existing diffusion models is not disentangled, there exists no consistent direction, and thus the editing trajectory must change its direction every time after stepping forward at $X_t$ or $H_t$. This is the method implemented with the additional training of a neural network. In this sense, manipulating indirectly through $H$ is not fully desirable, and alternatively we aim to disentangle the $X$ space, so that we can manipulate through the consistent directions.

Reviewer comment (Writing):
What are the $z$ and $z’$?

Author response:
__[Q2. $z$ and $z'$?]__

$z$ and $z'$ are the local coordinates of Riemannian manifold $M_1$ and $M_2$, respectively.
We added the description of $z$ and $z'$ in the caption of Fig. 2 (Currently Fig. 3) and in the revised main text.

Reviewer comment (Writing):
What is the $\it{f}$ on the bottom of Fig. 2? The reviewer saw the definition “A mapping between two Riemannian manifolds”, but it is not clear how it is implemented. Also, using a different letter with the encoder of $s_\theta$ is recommended to avoid any confusion.

Author response:
__[Q3. $f$ in Fig. 2?]__

First of all, we apologize for the confusion between $f$ and $\mathbf{f}$. $f$ in Fig. 2 (Currently Fig. 3) is a map between local coordinate spaces, while $\mathbf{f}$ is a map between Riemannian manifolds, following the notation from [E].

To address the confusion between notations, we changed the notation $\mathbf{f}$ to $e_\theta$ in the revision.

In our notation, $f = \Phi \circ e_\theta \circ \Phi_{n-1}^{-1}$, and this is implemented by direct composition of functions. We only need $\mathbf{J}_f(z) =  \frac{\partial f}{\partial z}(z)$, so we only compute $\mathbf{J}_f(z)$, which is indirectly done utilizing jacobian vector product (JVP). Note that in our work, we use $\Phi = id.$, based on the linearity of $H$ space shown in [1].

[E] Jang et al. "Geometrically regularized autoencoders for non-Euclidean data." ICLR 2022.

Reviewer comment (Writing):
What is $\phi$? in Fig. 2?

Author response:
__[Q4. What is $\Phi$?]__

$\Phi$ in Fig. 2 (Currently Fig. 3) is a chart, mapping from Riemannian manifold to local coordinate space. [F] This could be any homeomorphism, and we choose $\Phi = \text{id}$ following the linearity argument of $\mathcal{H}$. [1]
We added a description to the caption of Fig. 2 (Currently Fig. 3) and the main text.

Reviewer comment (Writing):
What is the input of $\mathcal{L}_{iso}$?

Author response:
__[Q5. Input to $\mathcal{L}_{iso}$]__

$L_{iso}$ gets the encoder of score model $e_{\theta}$, and timestep $t$. We changed the notation $L_{iso}$ to $L_{iso}(e_{\theta}, t)$ and its descriptions for clarity in Sec. 3.2.

Reviewer comment (Presentation):
I recommend moving the illustration part (Fig. 4 and its description) to the beginning of Section 3.

Author response:
__[Q6. Moving Fig. 4]__

We appreciate for this suggestion. Reflecting the reviewer's comment, we moved them to the beginning of Sec. 3 in the revised manuscript.

Reviewer comment (Experiments):
What are the interesting applications that the proposed method shows better performance than the baseline?

Author response:
__[Q7. Interesting applications]__

First, we can perform a smooth interpolation between two images, semantically mixing them with a user-defined weight. Baseline models exhibit rough or sub-optimal interpolating results (such as appearance of woman between two men), while our work resolves this issue. In the revised manuscript, we added more interpolation results that shows our method's advantages; please see Fig. 6 and Fig. VII in the revision.

Second, image editing directly in the $X$ space will be possible. As shown in Fig. 7, one may find a meaningful editing direction in the $X$ space and edit the desirable attribute by moving the $x_T$ along the direction. The $X$ space of the base diffusion models is not quite aligned with the semantic space. When $x_T$ is moved naively along a consistent direction, unexpected attributes often change together. Direct manipulation in the $X$ space is important, since it does not require access to the encapsulated pipeline (=the entire reverse process) of diffusion model.

Third, Fig. V shows multiple attributes can be manipulated without dependence to each other, thanks to our disentangled latent space. This will allow easier and more precise image editing.

Reviewer comment (Experiments):
How was the baseline interpolation done? Is it based on the naive slerp? If so, from my perspective, it is mandatory to compare the results with [1] since the authors have mentioned that [1] is an important related work.

Author response:
__[Q8. Comparison with [1]?]__

Yes, we use the naive slerp for both ours and the baseline interpolation.

Although [1] is an important related work proposing the $H$ space as the semantic space of diffusion models, the interpolations presented in the paper is based on the _change of editing strength while $x_T$ is fixed_, whereas our interpolation is based on the _change of noise $x_T$_  itself ($x(t) = Slerp(x_T^1, x_T^2; t)$). For this reason, the comparison between ours and [1] will be not apple-to-apple.

Reviewer comment (Evaluation):
The authors argue their proposed methods have 1. better interpolation and 2. better disentanglement. What is the measure supporting the second argument?

Author response:
__[Q10. Metric for disentanglement]__

As a response to the reviewer's question, we provide an additional metric, linear separability (LS), which has been used to measure disentanglement of the latent space by StyleGAN [B]. Please see Tab. 2 in the revised manuscript. With our method, LS significantly improves, indicating that the latent space is more disentangled.

Reviewer comment (Reproducibility):
I am not sure how heavy the computation of the Isometry Loss is. The authors used the stochastic trace estimator to substitute the trace of Jacobian to Jacobian-vector product (JVP).

Author response:
We sincerely thank the reviewer for valuable comments and support for our paper.

__[W1. About Isometry Loss]__

Given that $X \in R^{256 \times 256 \times 3}$ and $H \in R^{8 \times 8 \times 512}$, the encoder's Jacobian $J$ contains 6,442,450,944 elements. With a float32 data type, the Jacobian matrix uses approximately 24GB of memory. If we compute the Jacobian directly, it takes 202.7 seconds under our NVIDIA A100 environment.

In contrast, the JVP does not explicitly calculate the entire Jacobian matrix, but it directly computes the product of the Jacobian matrix with a specific vector, requiring only $(256 \times 256 \times 3 + 8 \times 8 \times 512) \times 4 = 91,750$ bytes, which is approximately 0.875MB of memory. In our isometry loss, we compute JVP three times to estimate a trace of Jacobian. On average, computing a single JVP takes 0.6 seconds.

We have added an additional experiment on the stochastic trace estimator in Appendix A based on your comments.


Please don't hesitate to share any additional comments or questions.

Reviewer comment (Evaluation):
I am curious about how accurate the stochastic trace estimator is.

Author response:
We sincerely thank the reviewer for valuable comments and support for our paper.

__[W1. About Isometry Loss]__

Given that $X \in R^{256 \times 256 \times 3}$ and $H \in R^{8 \times 8 \times 512}$, the encoder's Jacobian $J$ contains 6,442,450,944 elements. With a float32 data type, the Jacobian matrix uses approximately 24GB of memory. If we compute the Jacobian directly, it takes 202.7 seconds under our NVIDIA A100 environment.

In contrast, the JVP does not explicitly calculate the entire Jacobian matrix, but it directly computes the product of the Jacobian matrix with a specific vector, requiring only $(256 \times 256 \times 3 + 8 \times 8 \times 512) \times 4 = 91,750$ bytes, which is approximately 0.875MB of memory. In our isometry loss, we compute JVP three times to estimate a trace of Jacobian. On average, computing a single JVP takes 0.6 seconds.

We have added an additional experiment on the stochastic trace estimator in Appendix A based on your comments.


Please don't hesitate to share any additional comments or questions.