Alias-Free Latent Diffusion Models:Improving Fractional Shift Equivariance of Diffusion Latent Space (2503.09419v1)

Abstract: Latent Diffusion Models (LDMs) are known to have an unstable generation process, where even small perturbations or shifts in the input noise can lead to significantly different outputs. This hinders their applicability in applications requiring consistent results. In this work, we redesign LDMs to enhance consistency by making them shift-equivariant. While introducing anti-aliasing operations can partially improve shift-equivariance, significant aliasing and inconsistency persist due to the unique challenges in LDMs, including 1) aliasing amplification during VAE training and multiple U-Net inferences, and 2) self-attention modules that inherently lack shift-equivariance. To address these issues, we redesign the attention modules to be shift-equivariant and propose an equivariance loss that effectively suppresses the frequency bandwidth of the features in the continuous domain. The resulting alias-free LDM (AF-LDM) achieves strong shift-equivariance and is also robust to irregular warping. Extensive experiments demonstrate that AF-LDM produces significantly more consistent results than vanilla LDM across various applications, including video editing and image-to-image translation. Code is available at: https://github.com/SingleZombie/AFLDM

• The paper introduces AF-LDM, a latent diffusion model enhanced with an equivariance loss to maintain shift consistency even under fractional shifts.
• It combines ideal resampling techniques and alias-free modules in both VAE and U-Net architectures to suppress aliasing amplification during iterative denoising.
• The approach integrates Equivariant/Cross-frame attention to ensure global consistency, significantly improving performance in video editing and image-to-image translation.

Latent Diffusion Models (LDMs) have achieved impressive results in image synthesis but suffer from instability, where small shifts in input noise can lead to inconsistent outputs. This lack of shift-equivariance hinders their use in applications requiring temporal or spatial consistency, such as video editing and image-to-image translation. The paper "Alias-Free Latent Diffusion Models: Improving Fractional Shift Equivariance of Diffusion Latent Space" (2503.09419) investigates this issue and proposes Alias-Free LDM (AF-LDM) to enhance consistency by making the models more shift-equivariant.

The authors identify several reasons for the instability in standard LDMs (like Stable Diffusion, SD):

1. Aliasing amplification: The training process of the VAE encoder-decoder can amplify aliasing effects despite initial alias-free designs.
2. Accumulated aliasing: The iterative denoising process in the U-Net involves multiple steps, allowing aliasing to accumulate over time.
3. Non-equivariant self-attention: Standard self-attention modules, sensitive to global positions, are not inherently shift-equivariant under arbitrary (non-circular) shifts.

To address these challenges, AF-LDM introduces improvements in two main areas:

1. Continuous Latent Representation via Equivariance Loss:
   • The goal is to ensure that the latent space represents a band-limited continuous signal, making fractional shifts well-defined and predictable.
   • The model architecture incorporates anti-aliasing techniques inspired by StyleGAN3 (2106.11237) and Alias-Free Convnets (AFC) (2303.05253):
     • Ideal Downsampling: Replaces strided convolution with convolution + ideal low-pass filter + nearest downsampling. The low-pass filter is implemented in the Fourier domain.
     • Ideal Upsampling: Uses zero-padding + convolution with a sinc kernel, also implemented as multiplication in the Fourier domain.
     • Filtered Nonlinearity: Wraps nonlinear activation functions between ideal upsampling and downsampling to suppress high frequencies introduced by the nonlinearity.
   • These alias-free modules are applied to both the VAE and the U-Net.
   • Crucially, an equivariance loss is added to the training objective. This loss directly penalizes deviations from shift-equivariance:
     • For a network $f$ with scaling factor $k$, the loss is $$L = || f(T_{\Delta}(x)) - T_{k \cdot \Delta}(f(x)) ||_2^2$$, where $T_{\Delta}$ is a fractional shift operator (Fourier shift).
     • For the VAE: $L_{\text{eq}}^{\text{VAE}}$ is computed separately for the encoder $\mathcal{E}$ and decoder $\mathcal{D}$ using cropped integer pixel shifts $\Delta=(\Delta_x, \Delta_y)$ on the input image $x$ and corresponding fractional shifts $\Delta/k$ on the latent $z$. A valid mask $M_{\Delta}$ is used to exclude padded regions from the loss calculation.
     • For the U-Net: $L_{\text{eq}}^{\text{LDM}}$ is computed using fractional circular shifts $T_{\Delta}^{\text{cir}}$ on the noisy latent $z_t$. Circular shift is used for the input latent to avoid large padding effects in the iterative process, but a cropped mask $M_{\Delta}$ is still applied to the output difference.
   • This loss encourages the network to only utilize low-frequency information that is consistent across shifts, preventing aliasing amplification during training.
2. Equivariant Attention (EA) / Cross-Frame Attention (CFA):
   • Standard self-attention fails to be fully equivariant under non-circular shifts because the key and value pools change with the shifted input.
   • Equivariant Attention addresses this by fixing the keys and values from a reference frame ($x_r$) while allowing the queries to come from the shifted frame ($x_s$). The operation becomes $$\text{EA}(x_r, x_s) = \text{softmax}(x_sW^Q(x_rW^K)^\top)x_rW^V$$.
   • This mechanism is equivalent to Cross-Frame Attention (CFA), commonly used in video diffusion models during inference.
   • The authors highlight that using CFA is critical not just in inference but also during training (specifically, within the U-Net equivariance loss calculation) to prevent the loss from focusing on fixing the attention mechanism itself rather than suppressing aliasing.

Implementation Details:

• AF-LDM is built by modifying existing Stable Diffusion models. AF-VAE is initialized from SD VAE (kl-f8) and retrained on ImageNet. AF-LDM and AF-SD U-Nets are initialized from scratch or SD V1.5 and retrained on FFHQ and LAION Aesthetic 6.5+, respectively.
• Alias-free modules replace standard resampling (bilinear/nearest) and wrap nonlinearities in the VAE and U-Net.
• Training involves adding the proposed equivariance losses ($L_{\text{eq}}^{\text{VAE}}$ and $L_{\text{eq}}^{\text{LDM}}$) to the standard reconstruction/KL/GAN loss for VAE and diffusion loss for LDM/SD.
• Fourier shifts are used for implementing $T_{\Delta}$. Integer pixel shifts $(\Delta_x, \Delta_y)$ for VAE input are sampled from $[-\frac{3}{8}H, \frac{3}{8}W]$. Fractional shifts $(\Delta_x/k, \Delta_y/k)$ are used for latent spaces.

Practical Applications and Evaluation:

• Shift PSNR (SPSNR): A key metric to evaluate shift-equivariance. Higher SPSNR indicates better equivariance. It's calculated as PSNR between the output of a shifted input and the shifted output of the original input: $$\text{SPSNR}_f(x) = \text{PSNR}(f(T_{\Delta}(x)), T_{k \cdot \Delta}(f(x)))$$.
• Ablation Studies: Experiments show that while alias-free architecture helps, equivariance loss is crucial to maintain high SPSNR throughout training, especially for the U-Net across denoising steps (Fig. 4). Equivariant attention (CFA) combined with equivariance loss further improves results.
• Warping-Equivariant Video Editing: AF-SD's enhanced shift-equivariance translates to better robustness under irregular shifts like optical flow warping (measured by warping PSNR/MSE). A simple video editing method is proposed: invert video frames to latents (using CFA for consistency), then regenerate from inverted latents with a new prompt (again using CFA). This method achieves more consistent results than standard SD without explicit latent warping.
• Image-to-image Translation: AF-LDM improves consistency in tasks like $4 \times$ super-resolution (using latent I$^2$SB) and normal estimation (using AF-YOSO). Shifted inputs yield significantly more consistent outputs with alias-free models compared to baselines.

Limitations:

• The dependence on a reference frame for Equivariant Attention (CFA) means handling occlusions or new objects entering the scene is challenging, potentially leading to inconsistencies in these areas, similar to flow-based methods.

In summary, AF-LDM provides a practical approach to building diffusion models with enhanced shift-equivariance. By combining alias-free architectural design with a novel equivariance loss and strategically applying cross-frame attention during training and inference, it achieves significantly more consistent results for various generative and image-to-image tasks compared to standard LDMs. The code and models are available at https://github.com/SingleZombie/AFLDM.