Technical Summary: Diffusion Domain Expansion (DDE)

Problem Statement

Pre-trained diffusion models have achieved state-of-the-art results in generating images, audio, and other data modalities. However, retraining these models from scratch to handle larger output domains (e.g., panoramic images, long audio tracks) or more complex conditioning inputs (e.g., multiple object locations) is computationally expensive. Existing approaches to expand generative domains often rely on fixed, algorithmic coordination mechanisms (e.g., MultiDiffusion) that do not require additional training but are limited to specific tasks and lack flexibility. Other methods, such as training dynamic diffuser modules, often require fixed modalities at training time or fail to generalize to domains larger than those observed during training.

The core challenge addressed by this paper is how to efficiently extend the generative domain of pre-trained diffusion models to generate larger objects and handle more complex conditioning without retraining the base models, while ensuring the solution generalizes to instances larger than those seen during the coordinator's training.

Methodology: Diffusion Domain Expansion (DDE)

The authors propose Diffusion Domain Expansion (DDE), a framework that employs a compact, trainable "coordinator" network to orchestrate the outputs of pre-trained diffusion models.

General Formulation

Given a pre-trained conditional diffusion model $p(x|y)$ with a denoising network $D(x, y, t)$, the goal is to generate an expanded object $X^{[L]}$ conditioned on expanded input $Y^{[L]}$, where $L$ represents the scale (e.g., number of patches or conditioning labels).

1. Decomposition: The expanded pair $(X^{[L]}, Y^{[L]})$ is decomposed into $L$ smaller parts $([x_1, \dots, x_L], [y_1, \dots, y_L])$ using a mapping function $F$. These parts fall within the base model's original domain.
2. Base Generation: Each part $x_i$ is processed by the pre-trained denoiser $D(x_i, y_i, t)$.
3. Coordination: A trainable coordinator network $C[\cdot]$ takes the sequence of denoised outputs $[D(x_i, y_i, t)]_{i=1}^L$ and the conditioning information $[y_i]_{i=1}^L$ as input. It learns to coordinate these outputs to produce a coherent estimate of the full expanded object $X^{[L]}(0)$.
4. Training Objective: The coordinator is trained by minimizing the denoising error of the composite denoiser $D^{[L]}$ on a dataset of size $L_{train}$. Crucially, the method is designed to generalize to test sizes $L_{test} \ge L_{train}$.

Coordinator Architecture

The authors utilize a Vision Transformer (ViT) architecture for the coordinator, adapted from Peebles & Xie (2023).

• Input Processing: The denoised patches from the pre-trained models are "patchified" and augmented with global positional encodings relative to the larger object. The authors employ Rotary Position Embeddings (RoPE) to handle larger spatial dimensions effectively.
• Reconciliation: During the forward process, the transformer processes all patches. To handle overlapping patches (used to ensure continuity), the final output is constructed by averaging the values in the overlapping regions, inspired by the MultiDiffusion sampling method.
• Efficiency: The coordinator is significantly smaller than the base model (e.g., 66M vs. 405M parameters in music generation) and requires fewer training epochs.

Key Contributions

1. Novel Framework: Introduction of DDE, a method that uses a small, trainable coordinator to extend the generative capabilities of pre-trained diffusion models to larger objects and more complex conditioning.
2. Generalization Capability: Demonstration that the coordinator can generalize to generate objects and handle conditioning inputs larger in scale than those present in its training data ($L_{test} \ge L_{train}$).
3. Architecture Design: Validation of a ViT-based architecture as a suitable coordinator, which learns long-range dependencies while enforcing local consistency through overlap averaging.
4. Efficiency: The approach achieves high-quality generation with significantly fewer parameters and training time compared to retraining base models or using larger coordinator architectures.

Experimental Results

The authors evaluated DDE across three domains:

1. Long Audio Track Generation (Slakh2100)

• Task: Generating music tracks of length $10l$ (120 seconds) using a base model trained on length $l$ (12 seconds).
• Baselines: Concatenation and MultiDiffusion.
• Results: DDE achieved a Fréchet Audio Distance (FAD) of 2.112 for 4l tracks and 2.142 for 10l tracks. This significantly outperformed MultiDiffusion (4.796) and Concatenation (4.596) for the 10l task. The coordinator (66M parameters) was much smaller than the base model (405M).

2. Conditional Image Generation (CLEVR)

• Task: Generating images conditioned on multiple object locations ($L$ labels), where the base model only handles one location.
• Baselines: RRR (Reduce, Reuse, Recycle) and MultiDiffusion.
• Results: DDE demonstrated superior accuracy in satisfying multiple conditioning constraints. For 5 conditionings, DDE achieved 44.5% accuracy (Heun sampler), compared to 33.6% for RRR and 36.3% for MultiDiffusion. The coordinator generalized from training on $L_{train}=2$ to testing on $L_{test}=5$.

3. Satellite Image Generation

• Task: Generating large satellite images conditioned on schematic maps.
• Baselines: MultiDiffusion.
• Results: DDE achieved a lower FID score than MultiDiffusion for both in-distribution ($N=96$, FID 31.75 vs 37.81) and out-of-distribution ($N=128$, FID 27.37 vs 35.01) sizes. Visual inspection showed DDE produced more coherent images without the visible seams observed in MultiDiffusion samples.

Significance and Claims

The paper claims that DDE offers a practical and efficient solution for scaling diffusion models without the prohibitive cost of retraining. By learning to coordinate pre-trained processes, the method enables:

• Domain Expansion: Generation of larger objects (longer audio, larger images) and handling of more complex conditions (multiple labels) beyond the original model's capabilities.
• Generalization: The ability to extrapolate to scales larger than those seen during the coordinator's training.
• Quality and Consistency: The ViT-based coordinator effectively learns long-range dependencies, resolving conflicts between local denoising steps to produce globally consistent outputs.

The authors note that the primary limitation is the requirement for additional data and time to train the coordinator model, though this cost is substantially lower than training a full diffusion model from scratch. Future work is proposed for applying DDE to protein sequence generation and video creation from short clips.