Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers

The paper introduces the Hourglass Diffusion Transformer (HDiT), a scalable image generative model that enables efficient high-resolution pixel-space training without relying on latent autoencoders or multiscale architectures, achieving state-of-the-art performance on FFHQ-1024.

The Gist

Imagine you want to paint a masterpiece, but you have two choices for how to do it:

1. The "Blurry Sketch" Method (Current Standard): You first shrink your canvas down to a tiny, low-resolution sketch. You paint the sketch, and then you use a magic enlarger to blow it back up to a huge size. The problem? The magic enlarger is imperfect. It often smears the fine details, making hair look like wool, eyes look like smudges, and textures look soft and muddy. This is how most current AI image generators (like the original Stable Diffusion) work. They work in a "compressed" space to save computing power.

2. The "Direct Painting" Method (This Paper): You paint directly on the giant canvas, pixel by pixel, from the very first stroke. You don't shrink it down; you don't enlarge it later. You just paint the high-resolution masterpiece directly.

The Problem: Painting directly on a giant canvas is incredibly hard for a computer. If you try to use a standard "Transformer" (a type of AI brain known for being smart but computationally hungry) to look at every single pixel on a 1024x1024 image, the computer's brain explodes. The math required grows quadratically (like a square). If you double the size of the image, the work required goes up by four times. It's like trying to read every single letter in a library of books just to write one sentence.

The Solution: The Hourglass Diffusion Transformer (HDiT)

The authors of this paper built a new AI brain called HDiT. Think of it as a hierarchical "Hourglass" painter.

Here is how it works, using a simple analogy:

1. The Hourglass Shape

Imagine an hourglass.

• The Top (Wide): You start with the full, high-resolution image.
• The Middle (Narrow): The AI quickly shrinks the image down to a tiny, manageable core (like a 16x16 grid). Here, it figures out the "big picture" relationships (e.g., "This is a face," "The eyes go above the nose"). Because the image is small here, the computer can look at everything at once without getting overwhelmed.
• The Bottom (Wide): The AI expands the image back out to full size, adding the fine details (like the texture of skin or the strands of hair) as it goes.

2. The Secret Sauce: "Local" vs. "Global" Vision

The genius of HDiT is how it handles the different parts of the hourglass:

• At the Narrow Middle (Global Vision): When the image is tiny, the AI uses Global Attention. It looks at the whole picture at once to ensure the composition makes sense. This is expensive, but since the image is tiny, it's cheap to do.
• At the Wide Ends (Local Vision): When the image is huge (high resolution), the AI switches to Local Attention. Instead of trying to compare every pixel to every other pixel (which is impossible), it only looks at its immediate neighbors.
  • Analogy: Imagine you are painting a massive mural. To get the overall shape right, you step back and look at the whole wall (Global). But when you are painting the details of a flower petal, you don't need to look at the mountain in the background; you just need to look at the petals right next to the one you are painting (Local).

By doing this, the computer's workload grows linearly (like a straight line) instead of quadratically (like a steep hill). If you double the image size, the work only doubles, not quadruples.

3. Why This Matters

• No More Blurry Magic: Because HDiT paints directly on the high-resolution canvas (Pixel-Space), it doesn't need that imperfect "magic enlarger" (VAE) that current models use. The result is incredibly sharp, crisp images with fine details that other models miss.
• Scalability: Because the math is so efficient, they can train this model on massive images (1024x1024) without needing a supercomputer the size of a city.
• Better Editing: Since the AI understands the actual pixels, not a compressed code, it's much better at editing images. If you want to change the color of a shirt or fix a face, the AI knows exactly where the pixels are, rather than guessing based on a blurry sketch.

The Results

The paper shows that HDiT creates faces (on the FFHQ dataset) and objects (on ImageNet) that are sharper and more realistic than previous state-of-the-art models. It beats the competition in quality while being much more efficient, effectively bridging the gap between the "smart but slow" Transformers and the "fast but simple" older models.

In short: They built a new AI painter that knows when to step back to see the whole picture and when to zoom in to paint the details, allowing it to create stunning, high-definition art without burning out the computer.

Technical Summary

1. Problem Statement

Diffusion models have become the state-of-the-art for image generation, but scaling them to high resolutions (e.g., $1024 \times 1024$) directly in pixel space remains a significant challenge.

• The Latent Bottleneck: Most high-resolution models (like Stable Diffusion) use Latent Diffusion Models (LDMs), which compress images into a lower-dimensional latent space via a Variational Autoencoder (VAE) before generation. While efficient, this introduces quality limitations, such as the loss of fine details, artifacts, and poor reconstruction for downstream tasks like image editing.
• The Transformer Scaling Issue: Pure Transformer-based diffusion models (like DiT) offer superior scalability in parameter count but suffer from quadratic computational complexity ($O(n^2)$) regarding the number of image tokens ($n$). This makes training on high-resolution pixel data prohibitively expensive.
• Current Workarounds: Existing high-resolution pixel-space approaches often rely on complex multi-stage cascades, self-conditioning, or multi-scale losses, which add training complexity and architectural overhead.

2. Methodology: Hourglass Diffusion Transformers (HDiT)

The authors propose HDiT, a hierarchical, pure Transformer architecture designed to achieve linear computational complexity ($O(n)$) with respect to the number of pixels, enabling efficient high-resolution pixel-space synthesis.

Core Architectural Innovations

• Hierarchical Hourglass Structure: Inspired by Hourglass Transformers (Nawrot et al., 2022) and U-Nets, the model processes images at multiple resolutions.
  • Encoder: Spatially merges $2 \times 2$ tokens into one using Pixel-UnShuffle, reducing resolution as the network goes deeper.
  • Bottleneck: The innermost layer operates at a low resolution (e.g., $16 \times 16$ tokens) using Global Self-Attention to capture long-range coherence.
  • Decoder: Re-expands the resolution using Pixel-Shuffle and Skip Connections.
• Local vs. Global Attention:
  • Global Attention is restricted to the lowest resolution levels to ensure global coherence.
  • Local Attention (specifically Neighborhood Attention) is used at all higher resolution levels. This limits the quadratic cost of attention to a small, fixed core, allowing the rest of the network to scale linearly ($O(n)$) with image size.
• Skip Connection Mechanism: Unlike standard concatenation in U-Nets, HDiT uses Learnable Linear Interpolation (Lerp) for skip connections. The model learns a coefficient $f$ to interpolate between the skip connection and the upsampled branch ($x_{merged} = f \cdot x_{skip} + (1-f) \cdot x_{upsampled}$), which proved superior in deep hierarchies.
• Block Design Improvements:
  • Positional Encoding: Replaces standard additive positional embeddings with 2D Axial Rotary Positional Embeddings (RoPE), which improves generalization and reduces patch artifacts.
  • Feedforward Network: Uses GEGLU activation instead of standard GELU, with the modulation signal derived from the data itself rather than conditioning.
  • Normalization: Uses Adaptive RMSNorm (similar to LLaMA) conditioned on noise level and class, initialized with zero-output projections for stability.
• Training Strategy: The model is trained directly on RGB pixels without VAEs. It utilizes Soft-Min-SNR loss weighting (a smoothed version of Min-SNR) to improve convergence, particularly at high resolutions where standard noise schedules fail.

3. Key Contributions

1. HDiT Architecture: Introduction of a hierarchical Transformer backbone that scales to megapixel resolutions with $O(n)$ computational complexity, bridging the gap between the efficiency of CNNs (U-Nets) and the scalability of Transformers.
2. Pixel-Space High-Resolution Synthesis: Demonstrates the ability to train and generate high-quality images at $1024 \times 1024$ directly in pixel space, eliminating the need for latent compression (VAEs) and associated quality loss.
3. State-of-the-Art Performance: Sets new records for diffusion models on the FFHQ-1024 dataset and achieves competitive results on ImageNet-256, outperforming standard DiT baselines even at lower resolutions ($128 \times 128$).
4. Simplicity: Achieves these results without complex training tricks like progressive growing, multi-scale losses, or self-conditioning.

4. Experimental Results

• FFHQ-1024 ($1024 \times 1024$):
  • The HDiT model (85M parameters) achieves a FID of 5.23 (50k samples), significantly outperforming previous diffusion models like NCSN++ (FID 53.52) and latent-based DiT.
  • It sets a new state-of-the-art for diffusion models on this dataset, producing sharp, detailed faces with symmetric features, whereas latent models often suffer from blurriness and asymmetry.
  • It also achieves the best DINOv2-based metrics (FDD and KDD), which correlate better with human preference than FID for face generation.
• ImageNet-256 ($256 \times 256$):
  • A 557M parameter class-conditional model achieves a FID of 6.92 (without classifier-free guidance).
  • While slightly behind massive models using self-conditioning (like RIN or VDM++), it outperforms standard U-Net based pixel-space models (ADM) and latent models (DiT-XL/2) despite operating at a higher effective resolution (pixels vs. latents).
• Efficiency:
  • At $256 \times 256$, HDiT is 10x more efficient (in FLOPs) than a parameter-matched DiT.
  • At $1024 \times 1024$, the efficiency gap widens to >100x compared to DiT, making high-resolution training feasible.

5. Significance and Impact

• Eliminating the Latent Bottleneck: By proving that Transformers can scale efficiently in pixel space, HDiT removes the reliance on VAEs. This unlocks higher fidelity for image editing, inpainting, and controllable generation, as there is no information loss from compression/decompression.
• Scalability: The $O(n)$ complexity scaling makes it possible to train massive diffusion models on high-resolution data without the prohibitive costs associated with standard Transformers.
• Future Directions: The architecture is well-suited for other modalities (video, audio) and tasks (super-resolution, text-to-image). The authors suggest that combining HDiT with latent spaces could further push resolutions to multi-megapixel scales.

In summary, HDiT represents a paradigm shift in diffusion model architecture, demonstrating that hierarchical Transformers can efficiently handle high-resolution pixel data directly, offering a path to higher quality and more versatile generative AI without the compromises of latent space models.