Self-Attention Decomposition For Training Free Diffusion Editing

1. Problem Statement

Diffusion models have achieved state-of-the-art results in image synthesis, but controlling them for precise, targeted semantic editing (e.g., changing age, gender, or expression) remains difficult. Existing methods for finding interpretable editing directions face two primary limitations:

• Data Dependency: Most approaches rely on sampling large datasets to discover directions, which ties the results to specific data distributions and limits generalization to unseen images.
• Computational Cost: Many methods require training auxiliary networks, performing expensive geometric operations (e.g., Riemannian geometry), or extensive fine-tuning, leading to high latency.

The goal is to find a method that derives robust, interpretable editing directions directly from the pretrained model parameters without additional data sampling or training.

2. Methodology

The authors propose an analytical, training-free approach based on eigen-decomposition of self-attention weight matrices.

Core Hypothesis

The self-attention layers in diffusion models encode rich structural and semantic information learned during training. By analyzing the weight matrices of these layers, one can extract eigenvectors that correspond to meaningful semantic attributes.

Technical Derivation

1. Latent Whitening Assumption: The method operates under the observation that during early denoising timesteps ($0.5T \leq t \leq 0.8T$), the latent feature maps are approximately "whitened" (their covariance is close to the identity matrix, $E[Z^T Z] \approx I$). This allows the separation of self-attention effects from the specific data distribution.
2. Perturbation Analysis: The authors analyze the effect of perturbing the latent input $Z$ along a direction $n$ ($Z' = Z + \alpha n$). They aim to maximize the change in the attention output ($\Delta Attn$) to find directions that most significantly alter semantic attributes.
3. Weight Space Decomposition:
   • The self-attention block uses Query ($Q$), Key ($K$), and Value ($V$) projection matrices ($W_Q, W_K, W_V$).
   • Through first-order Taylor expansion and the whitening approximation, the authors derive that maximizing the impact on the attention output is equivalent to maximizing the quadratic form $n^T C n$, where $C$ is a combined matrix:
     $$C = W_Q^T W_Q + W_K^T W_K + W_V W_V^T$$
   • The Solution: The optimal editing direction $n^*$ is the principal eigenvector of matrix $C$. Subsequent eigenvectors provide orthogonal editing directions.
4. Application: These eigenvectors are applied as linear perturbations to the latent features during the specific early denoising window ($0.5T < t < 0.8T$).

3. Key Contributions

1. First Weight-Space Derivation: This is the first work to derive interpretable editing directions directly from the pretrained weights of diffusion models using eigen-analysis, eliminating the need for auxiliary training or large-scale sampling.
2. Theoretical Foundation: The paper provides a theoretical derivation showing that the principal components of the combined Query-Key-Value weight matrices yield semantically meaningful directions under the latent whitening assumption.
3. Training-Free Efficiency: The method requires zero additional data and zero fine-tuning, making it a "plug-and-play" solution for existing pretrained models.

4. Experimental Results

The method was evaluated on CelebA-HQ (faces), LSUN Cats, Cars, and Rooms, comparing against benchmarks like Asyrp, LocoEdit, Pullback, and NoiseCLR.

• Qualitative Performance:

  • Linearity: The method exhibits strong linear editing properties; varying the perturbation strength $\alpha$ results in proportional semantic changes (e.g., gradual aging or gender shifts).
  • Disentanglement: It achieves precise edits on specific attributes (e.g., "Chubby," "Hairstyle," "Eyes Closed") while preserving unedited regions better than baselines.
  • Generalization: The eigenvectors generalize across different datasets (faces, cars, rooms), capturing global attributes (age, gender) and localized features (nose shape, pose).

• Quantitative Performance:

  • Speed: The method reduces editing time by 60% compared to current benchmarks. While baselines take seconds to days (e.g., NoiseCLR takes 1 day), the proposed method takes approximately 3 seconds.
  • Image Quality: It achieves superior scores across all metrics:
    • SSIM: 0.92 (vs. 0.85–0.90 in baselines)
    • PSNR: 28.5 dB (vs. 23.9–26.8 dB)
    • Identity Preservation (FaceNet): 0.82 (vs. 0.68–0.78)
    • Semantic Alignment (Dir-CLIP): 0.75 (vs. 0.58–0.65)

• Ablation Studies: Confirmed that the editing strength is linearly controllable via $\alpha$ and that the method works effectively on models trained on diverse datasets (LSUN).

5. Significance and Future Work

• Efficiency & Accessibility: By removing the need for data sampling and retraining, this method makes high-quality semantic editing accessible and computationally efficient for any pretrained diffusion model.
• Interpretability: It offers a new lens for understanding diffusion models, suggesting that semantic structure is explicitly encoded in the self-attention weights.
• Future Directions: The authors plan to extend this analysis to cross-attention layers in Diffusion Transformers (DiTs) to enable joint text-visual control and explore robustness in higher-dimensional latent spaces.

In summary, this paper presents a breakthrough in diffusion editing by leveraging the mathematical properties of self-attention weights to create a fast, training-free, and highly effective method for semantic image manipulation.