Diffusion Mental Averages

1. Problem Statement

The paper addresses the challenge of synthesizing a single, sharp, and realistic "mental average" (prototype) of a concept directly from pre-trained text-to-image diffusion models.

• The Gap: While diffusion models can generate countless variations of a concept (e.g., "dog," "firefighter," or abstract concepts like "freedom"), they lack an explicit mechanism to output a representative prototype that captures the model's internal semantic consensus.
• Limitations of Existing Methods:
  • Pixel-space averaging: Requires spatial alignment (e.g., GANgealing), which is infeasible for complex or abstract concepts and results in blurry images due to the averaging of fine details.
  • Data-centric distillation: Methods like Dataset Distillation (D4M, MGD3) aim to compress datasets for downstream tasks but often produce prototypes that are either inconsistent (varying wildly between runs) or lack visual realism.
  • Latent-space averaging: Standard autoencoders (VAEs, GANs) have explicit bottlenecks allowing for embedding averaging and decoding. Diffusion models, however, lack a single semantic bottleneck; semantic information is distributed across timesteps and layers, making direct averaging impossible without a decoder.

2. Methodology: Diffusion Mental Averages (DMA)

The authors propose DMA, a model-centric approach that treats averaging as a trajectory alignment problem within the diffusion model's semantic space.

Core Concept: Trajectory Alignment

Instead of generating samples and averaging them post-hoc, DMA jointly optimizes multiple noise latents so their denoising trajectories converge toward a shared semantic consensus.

1. Semantic Space (h-space): The method leverages the h-space (the bottleneck layer of the U-Net), which recent studies show encodes interpretable, approximately linear semantic information.
2. Progressive Alignment: The process operates across diffusion timesteps ($t$):
   • Initialization: Start with $K$ noisy latents ($z^{(0)}_k$) drawn from a Gaussian distribution.
   • Iterative Optimization: At each timestep $t$:
     • Extract the h-space activation $H(z^{(t)}_k)$ for each latent.
     • Compute the mean activation $\bar{A}^{(t)} = \frac{1}{K}\sum H(z^{(t)}_k)$.
     • Noise Optimization: Optimize each latent $z^{(t)}_k$ to minimize the L2 distance between its activation and the mean $\bar{A}^{(t)}$. This forces all trajectories to align on the shared semantic features at that specific stage of denoising.
     • Denoising: Perform one step of DDIM denoising to move to $t-1$.
   • Termination: This alignment continues until a cutoff timestep $t_{stop}$, after which standard denoising completes the generation.
3. Result: The final output is a single, sharp prototype that represents the "mental average" of the concept as encoded by the model.

Handling Multimodal Concepts (Mode Discovery)

For concepts with distinct modes (e.g., "crane" as a bird vs. a machine, or "dog" by breed), a single average creates a blur. DMA addresses this via:

• Clustering: Samples are clustered in a stable semantic space (e.g., CLIP or BLIP-VQA for grounded attributes like ethnicity).
• Cluster-Specific Conditioning: To ensure the average faithfully reflects the specific cluster's semantics (bridging the gap between the clustering space and the diffusion h-space), the authors apply lightweight conditioning:
  • Textual Inversion: Learns a specific token embedding for the cluster.
  • LoRA (Low-Rank Adaptation): Fine-tunes the conditioning pathway to capture richer intra-mode variations (found to be superior to Textual Inversion for complex modes).
  • DMA is then run independently within each cluster.

3. Key Contributions

1. First Model-Centric Averaging: Introduces the first method to generate consistent, realistic averages directly from pre-trained diffusion models without access to the original training dataset.
2. Trajectory Alignment Formulation: Recasts the averaging problem as aligning denoising trajectories in the h-space, solving the issue of distributed semantic information across timesteps.
3. Multimodal Generalization: Extends the method to handle ambiguous concepts by combining clustering with LoRA/Textual Inversion, enabling the discovery of distinct semantic modes.
4. Architecture Agnosticism: Demonstrates that the method generalizes across different Stable Diffusion variants (SD1.5, Realistic Vision, PixelArt) and architectures (DiT).

4. Experimental Results

The authors evaluated DMA on 12 concepts across four categories: Animals, Persons, Objects, and Abstract concepts.

• Baselines: Compared against GANgealing (spatial alignment), Avg. VAE (latent averaging), D4M, and MGD3 (dataset distillation).
• Metrics:
  • Consistency: How similar are results across different random seeds? (DMA is significantly more consistent than D4M/MGD3).
  • Representativeness: How close is the prototype to the semantic center of the distribution? (DMA achieves the lowest distance scores).
  • ImageReward: A human preference metric for quality and alignment. (DMA achieves the highest scores).
• Qualitative Findings:
  • Realism: Unlike baselines which produce blurry or distorted artifacts, DMA generates sharp, structured images (e.g., a clear astronaut with a visor, rather than a vague helmet shape).
  • Abstract Concepts: DMA successfully visualizes abstract concepts like "freedom" (Statue of Liberty) or "poverty," whereas baselines fail or produce inconsistent results.
  • Bias Detection: The method reveals model biases, such as gender bias in "soldier" (SD1.5 generates males, while PixelArt is gender-neutral) and cultural biases in "Italy" (converging on Venice canals).

5. Significance and Impact

• Interpretability Tool: DMA provides a concrete visual summary of what a diffusion model "thinks" a concept looks like, serving as a powerful lens for auditing model biases and understanding learned representations.
• Data Efficiency: It offers a way to compress the knowledge of a massive dataset into a single, high-quality prototype, potentially useful for privacy-preserving data sharing or efficient downstream training.
• Methodological Advancement: By proving that semantic alignment can be achieved through trajectory optimization in the h-space, the paper opens new avenues for manipulating and analyzing the internal mechanics of diffusion models beyond simple generation.

In summary, Diffusion Mental Averages bridges the gap between the generative power of diffusion models and the need for interpretable, representative prototypes, offering a robust tool for both analyzing model behavior and generating high-quality concept summaries.