Improving the Diffusability of Autoencoders

We deeply appreciate the reviewer’s valuable feedback and constructive recommendations. Below, we systematically respond to each issue highlighted. We will ensure comprehensive incorporation of all suggestions into our manuscript.

It is not clear why high-frequency components have higher dimensionality.

Reviewer JqrQ has raised the same concern, and due to the limited space, we politely refer to our argumentation in that other response.

It is not clear why high-frequency components are more susceptible to error accumulation.

Similarly to the previous question, we kindly refer the reviewer to our response to Reviewer JqrQ on the same matter.

Can SE be helpful just because it improves the reconstruction of low-frequency components?

Not quite. PSNR/SSIM scores are extremely sensitive to low-frequency reconstructions, and all the autoencoders (vanilla, fine-tuned and fine-tuned+SE) perform on par in terms of these metrics.

From-scratch training.

Due to the space limit, we are again forced to refer the reviewer to our response on the same question to Reviewer MRKn. We apologize for this inconvenience.

The improved performance can be due to a different reason: high-frequency components of images have a lower entropy and should be easier to model.

In this [plot], we visualize the entropy of the latents for regularized/non-regularized FluxAE at various frequencies, and also RGB. Entropy is computed by obtaining a histogram of 100 bins for each frequency and its corresponding density. While for images space, high-frequencies indeed have lower entropy, for the latents, HF entropy exhibits a much flatter or an even slightly increasing profile indicating that they might be harder to model. We note that while higher entropy, intuitively, should lead to harder modeling of a distribution, it does not necessarily affect the quality of the final samples obtained from passing the latents through the decoder. Higher frequencies with smaller scale can still exhibit high entropy while our SE reduces the dependence on these frequencies.

Figure 5 does not show that SE preserves more content compared to the baseline.

We try being gentle about enforcing SE regularization not to affect the reconstruction quality. Figure 5 shows that our regularized AE does not introduce spurious high frequencies when they are chopped off in the latents. This also results in meaningful improvements in the corresponding reconstruction metrics as confirmed in Figure 8. Other examples (e.g., [this one]) can show this effect more noticeably.

[3 writing issues]

We fully agree with the remarks and will incorporate them in the manuscript.

Why are some amplitudes greater than 1 after normalization by D_{0,0} in Figures 2 and 3?

D_{0,0} corresponds to the lowest frequency, which is basically the mean of all the values. It might not necessarily have the largest amplitude compared to other frequencies (though this rarely happens in natural signals). For example, [this figure] shows an example of an image which has 0 amplitudes everywhere, except for the “[[0,1]]” spatial frequency.

Why do high-frequency components become more pronounced as bottleneck channels increase?

This is an interesting question. We hypothesize that increasing the autoencoder's bottleneck channels enables the model to better capture finer, high-frequency details. Initially, with limited capacity, the encoder prioritizes smoother, low-frequency information. As capacity grows, it encodes additional high-frequency details, which are unique and informative. However, without explicit regularization promoting frequency-based disentanglement, these high-frequency components distribute across channels in an unstructured manner. Thus, higher-dimensional bottlenecks enhance high-frequency representations but do not yield systematic frequency-specific disentanglement per channel. We will clarify this point in the revised manuscript.

Results for various numbers of steps.

We generated the Table 1 results using 256 steps (L#318). Additional plots for FluxAE/CogVideoX-AE/LTX-AE across step counts (16,32,64,128,256) for FID/DinoFID and FVD (images: 50K samples, videos: 10K samples) are provided [here (see the neighboring folders as well)]. Regularized autoencoders consistently improve diffusability. We will include these results (+ DiT-XL/2 for CogVideoX-AE) in the final paper.

Extended KL influence discussion.

We provided an expanded KL discussion (noise injection, relevant literature, RGB+noise spectrum analysis) in response to Reviewer tijt.

We welcome any additional suggestions from the reviewer.

We thank the reviewer for their valuable remarks. Below, we address each raised concern.

Why high-frequency components have higher dimensionality?

We agree this could be clearer. By "low-frequency components," we mean DCT coefficients required to reconstruct a feature map downsampled by factor k per spatial dimension. High-frequency components are remaining coefficients needed to fully reconstruct the original map. Since DCT coefficients scale quadratically with k, we have:

Number of high-frequency components = (k² − 1) × Number of low-frequency components

For k=2 (our experiments), high-frequency components thus have three times higher dimensionality. We'll clearly state this in the final version.

It is unclear why higher frequencies are only generated in the final steps.

This claim originates from prior literature (Rissanen et al., "Generative Modelling With Inverse Heat Dissipation"; also "Diffusion is spectral autoregression" and DCTDiff). Intuitively, high-frequency components are too easily erased by noise in the early denoising steps, which is why the model can only pick them up for smaller noise levels, which relate to the final denoising steps. We will extend the discussion of this topic in the earliest revision.

It is unclear why higher frequencies are more prone to error accumulation.

This observation is also an already explored phenomenon in the broader diffusion literature. A solid reference we can point to is Li et al. “On Error Propagation of Diffusion Models” (ICLR 2024) with Figure 2 of their paper serving as a clear illustration of such behaviour. Intuitively, the denoising process of a diffusion model can be represented as an ODE trajectory. And ODEs are known to accumulate errors very quickly even for simple processes (e.g., even for the simplest y’ = y equation, forward Euler solver would accumulate the error proportional to h² with each step of size h, and this error compounds exponentially as steps progress). Then, since diffusion models generate high frequencies later in the trajectory (per our previous response), the error accumulation is amplified specifically for them.

Could SE interact with other AE properties, rather than spectral properties alone?

In Table 6 of the appendix, we provide the ablation for direct high frequency chop-off: it only has influence on spectral properties (and this influence is the same as for SE), eliminating the interaction of other factors. It noticeably improves the diffusability, but we opt for the equivalent SE regularization since it’s much simpler and less error-prone to implement and should be easier to adopt by the community.

Computational cost and potentially unfair AE training budget comparison.

We measured the FLOPs of FluxAE (for the batch size of 1 and resolution of 256²) using [fvcore]. The entire encoder-decoder pass has 447 GFLOPS, and is split between the encoder/decoder as 136 vs 311 GFLOPS. Our regularization reuses the encoder pass and only runs the decoder with x2 or x4 reduced resolution (the scale sampled randomly during training). This results in 77.6 or 19.4 extra GFLOPs of the decoder, which is almost exactly 1/4 or 1/16 of the decoder compute or +17% or 4.5% of the total forward pass. Since we sample x2 or x4 downsampling factor equally randomly, this results in ~10.75% of total FLOPs overhead for our regularization.

To strengthen our point even further, we ran an experiment where the baseline FluxAE was fine-tuned strictly for 2 times longer (for 20K instead of 10K iterations). The resulted DiT-B/2 model achieved FID@5k and DinoFID@5k of only 33.99 and 642.7 vs the corresponding metrics of 25.87 and 551.27 of our FluxAE+FT-SE, fine-tuned for only 10K steps.

Balancing reconstruction and SE regularization.

We appreciate the raised concern and we provide the ablation over the SE strength in the table [here]. We remain cautious about reducing the influence of the main reconstruction term not to lose in terms of the reconstruction quality.

Training AE and LDM on the same dataset.

We conducted from-scratch FluxAE and CogVideoAE training on ImageNet/Kinetics datasets, followed by DiT-B/2 training, as detailed in our response to Reviewer MRKn. In both cases, SE regularization consistently improved results.

Unclear if suppressing high frequencies motivates improved performance.

We emphasize that merely suppressing high frequencies is insufficient; it is also crucial to prevent the decoder from arbitrarily amplifying them (we provide a detailed discussion on this in Sec. 3.3). This ensures alignment between the diffusion process’s strength—its ability to generate low-frequency components—and human perception, which is more tolerant to errors in low frequencies. Our experiments with direct frequency chop-off (described above) articulates this more explicitly.

We thank the reviewer for their insightful comments, which have greatly improved our work. Below, we address each concern.

Ambiguous caption in Figure 2

Figure 2 shows spectra of latent codes from real images encoded with from-scratch trained FluxAE with varying bottleneck sizes. We clarified the caption accordingly.

Exploring influence of bottleneck dimensionality on high-frequency components with more architectures

Upon closer inspection, we discovered a broadcasting issue in our spectrum computation pipeline, which led to distorted plots in the original submission. Updated FluxAE results are [here]. We further added analyses for two additional architectures: [WanAE] (which has the same bottleneck dimensionalities as CogVideoX-AE, but considerably faster to train) and [LTX-AE], trained with increasing bottleneck channel sizes. As one can see from these figures, autoencoders with higher channel sizes tend to possess high-frequencies of larger relative magnitude.

Fig. 3 does not show clear KL-high-frequency trend

The FluxAE KL plot previously had the same broadcasting issue; corrected results are [here]. We included similar analyses for [WanAE] and [LTX-AE]. These demonstrate a clearer relationship between KL and high-frequency energy.

The influence of KL on the high-frequency spectrum can be attributed to its role in injecting random noise during the encoding stage—an effect present both during training and inference. As KL regularization increases, so does the level of injected noise. Since random Gaussian noise has a uniform power spectrum, this flattens the frequency distribution, disproportionately inflating the high-frequency tail. We illustrate this in [this plot], where progressively increasing Gaussian noise (added to normalized RGB signals in [[–1, 1]]) results in a visible elevation of the high-frequency content. A similar effect was also independently observed in concurrent work on [SwD distillation (Figure 1)].

This reveals a nuanced trade-off: KL regularization aligns latents to the standard normal prior, facilitating downstream diffusion (as noted by LSGM), yet increases high-frequency energy, potentially hindering diffusability. We believe this trade-off warrants further attention.

Eq. 2 missing SE regularization loss weight

We thank the reviewer for pointing this out and corrected the equation.

Table 3 unclear relationship between KL strength and high-frequency components

There may be a misunderstanding: Table 3 illustrates KL's negative effect on diffusability, whereas its impact on high-frequency energy appears in Figure 3. To strengthen this analysis, we added results from DiT-L [here].

This table shows increased KL generally boosts small-model LDM performance, but at the expense of poorer reconstruction and stability, limiting scalability (consistent with SD3 findings). In contrast, our regularization improves LDM performance without harming reconstruction, scaling well to larger models.

Clarify “diffusability”: does it include efficiency and quality, and are there metrics?

We use "diffusability" strictly for generation quality, independent from efficiency. All else equal, increased diffusability should enhance LDM generation quality. Unfortunately, we found no reliable quantitative metrics correlating consistently with diffusability, including diffusion loss magnitude or decoder Lipshitz constants (as theoretically connected via Eq. 28 in [LFM]). We believe it is still important to introduce such a term to the community to bring attention to this property, since it remains largely ignored these days.

Clarify how spectral analysis suggested high-frequency components impact diffusability

We initially sought high-compression autoencoders but quickly found that increasing bottleneck size significantly reduces diffusability, keeping other factors constant. Then, we concurrently worked on cascaded latent diffusion pipelines, which required autoencoders to support downsampling, motivating us to explore their spectral characteristics. This revealed how increased bottleneck dimensionality relates to diminished diffusability, informing our central hypothesis.

We would greatly welcome any further comments the reviewer might have.

We sincerely thank the reviewer for their thorough feedback. In what follows, we carefully respond to each of the points raised. All comments and suggestions will be fully reflected in the revised manuscript.

From-scratch training

Our main motivation of fine-tuning instead of training from scratch is three-fold: Be able to compare against strong established benchmarks; Eliminate the possibility that we might have some “bug” in the training pipeline, which our regularization is rectifying (please, note that none of the explored SotA AEs release their training code). Carry more value to the community of improving popular autoencoders.

That being said, we launched multiple experiments for from-scratch training for FluxAE and CogVideoAE. Moreover, we opted for using public datasets to have fully self-sufficient experiments. Namely, we trained FluxAE from-scratch on ImageNet for 200K steps and CogVideoAE from-scratch on Kinetics-700 for 60K steps (it’s a heavyweight AE and slower to train). Due to the limited rebuttal period time, we only have the DiT-B/2 results till 300K steps on ImageNet, and till 250K steps on Kinetics. For both datasets, the regularized AEs lead to better LDM convergence. DiT-B/2 for non-regularized vs regularized AEs perform as follows:

• For ImageNet:
  • DinoFID@5k: 569.93 vs 561.4
  • DFID@5k: 27.94 vs 28.79
• For Kinetics:
  • DDinoFID@5k: 652.5 vs 561.21
  • DFID@5k: 21.54 vs 19.66
  • DFVD@5k: 265.94 vs 379.88

In all cases, our regularization improves the diffusability. We’ll include the full training results (400K) in the revised version of the paper.

​​Training on public data

We thank the reviewer for pointing this out, we include the results on public data together with the from-scratch training experiments in the previous message.

Some training details are under-explained (e.g. self-cond)

Our self-conditioning mechanism follows prior work without any modifications (i.e., RIN, FIT, WALT). Namely, during training with a 90% probability, we run an auxiliary forward pass with the DiT model, take its activations from the last block (i.e., right before the “unpatchify” projection), project them with a linear layer and add as residuals to the input tokens after “patchification” of the main training forward pass. For that auxiliary forward pass, following RIN, we use the same noise level $\sigma$ and “no-grad” context (i.e., we do not backpropagate through the auxiliary forward pass)

We added the self-conditioning discussion into the “Implementation details” appendix. We would be grateful to the reviewer if they point out any further missing implementation details and we would happily include them into the submission.

The DiT training details section in appendix A ends on a trailing sentence.

We thank the reviewer for pointing out that writing mistake. The sentence was intended to convey: “In essence, this reduces the total dataset size, but since we do the same procedure for the entire CogVideoX-AE family, the models are comparable between each other.” We fixed the error.

Explanation of non-monotonic KL influence on reconstruction quality (Table 3).

Yes, FluxAE was indeed unstable when fine-tuned with KL regularization for some of the KL β weights: it was stable for β of 0, 1e-7, 1e-4, and 1e-3, but not for 1e-6, 1e-5, 1e-2, and 1e-1. For from-scratch training (which we were doing for Figure 3), it was stable for the entire range (from 0.0 to 0.1), but after some threshold of β <= 1e-4, there was almost no difference in PSNR or FID. Our intuition is that high-capacity autoencoders (like FluxAE) can accommodate high KL penalty for their latents (for from-scratch training, we started noticing degradation only for β >= 1e-3), and their reconstruction quality is governed by other factors, making KL interference less predictable (up to a certain factor).

Should the reviewer have further comments, we would be glad to incorporate them fully into our manuscript.