Saining Xie
New York University

Joint work with Sanghyun Woo, Bill Peebles, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon

ConvNeXt
🚀ezinredoM

LinearGELU
LinearLN

ResNetDepthwise
7x7

A ConvNet for the 2020s, Liu et al., CVPR 2022

ConvNeXt🚀ezinredoM

LinearGELU
LinearLN

ResNetDepthwise
7x7

Attention is not essential for scalability.

ConvNet vs Transformer:
Supervised LearningConvNet vs Transformer:
CLIP TrainingConvNeXt is more efficient & more friendly to higher resolution input.

ConvNet vs Transformer:
Self-Supervised Learning

input

targetreconstruct

Masked Autoencoders (MAE)[He et al., CVPR 2022]

Q: Can we use ConvNeXt
as a backbone network for
self-supervised learning?

Masked Autoencoders (MAE)

inputtargetreconstruct- Asymmetric encoder-decoder architecture.
- Only encoding visible patches.
- How can ConvNet also achieve this?

Dense ConvNets:

Image as 2D pixel grid

[MinkowskiEngine, Choy et.al. 2019]

Sparse ConvNets:
Image as point cloudFCMAE: Fully Convolutional Masked Autoencoder
Input

ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders, Woo et al., CVPR 2023.

FCMAE: Fully Convolutional Masked Autoencoder
Encode visible patches using sparse convolutions
Input

rsespavcon

Sparse ConvNeXt

FCMAE: Fully Convolutional Masked Autoencoder
Add mask tokens
Input

rsespavcon
SparseConvNeXt

FCMAE: Fully Convolutional Masked Autoencoder
Reconstruct masked patches
Input

rsespav
conSparse
ConvNeXt

ConvNeXt BlockFCMAE: Fully Convolutional Masked Autoencoder
After pre-training, the weights are converted back to standard layers without requiring specialhandling.
Inputrsespav
conSparse
ConvNeXt

ConvNeXt Block

The impact of FCMAE
Train & Test Protocol• 800ep Pre-training on ImageNet-1K train• 100ep Fine-tuning on ImageNet-1K train• Test on ImageNet-1K val

FCMAE pre-trained ConvNeXt surpasses random initialization
but still lags the original 300ep supervised setup.

Feature Collapse
Feature collapse detected in the dimension-expansion MLP layers inFCMAE pre-trained ConvNeXtV1 blocks.

Dead Neurons

Saturated Neurons

GRN: Global Response Normalization

𝐻𝐶𝑊GRN: Global Response Normalization
1) Global feature aggregation

𝒢(𝛸)

𝛸!
𝐻𝐶𝑊

1×1×𝐶

GRN: Global Response Normalization
1) Global feature aggregation
2) Feature normalization

𝒢(𝛸)

𝛸!
𝐻

𝒩(||X! ||) =||X! ||∑ ||X! ||𝐶𝑊

1×1×𝐶

1×1×𝐶

GRN: Global Response Normalization
Three lines of code

1) Global feature aggregation
2) Feature normalization
3) Feature calibration

𝒢(𝛸)

𝛸!
𝐻𝒩(||X! ||) =||X! ||∑ ||X! ||𝐶𝑊

1×1×𝐶

1×1×𝐶

ConvNeXt V2
V1 Block

V2 Block

96-d96-dd7x7, 96d7x7, 96LN

1x1, 384
GELU

1x1, 96-LayerScale

LN

1x1, 384
GELU+GRN

1x1, 96

Impact of GRN
ConvNeXt v1

ConvNeXt v2

ConvNeXt v1

ConvNeXt v2

Results

Works for tiny tinymodels as well!

ConvNet vs Transformer:
Diffusion Models

Diffusion models use highly-nonstandard architectures…

Linear Decoder

Transformer Block
Transformer Block
Transformer Block
Transformer Block
Tokenize“A dog walked over to the barn…”

GPT-3 / Chinchilla / …

Diffusion models use highly-nonstandard architectures…

Linear Decoder

Transformer Block
Transformer Block
Transformer Block
Transformer Block
Tokenize“A dog walked over to the barn…”

GPT-3 / Chinchilla / …

Imagen U-Nets

Diffusion models use highly-nonstandard architectures…

Linear Decoder

Transformer Block
Transformer Block
Transformer Block
Transformer Block
Tokenize“A dog walked over to the barn…”

GPT-3 / Chinchilla / …

Imagen U-Nets

Diffusion models use highly-nonstandard architectures…

Linear Decoder

Transformer Block
Transformer Block
Transformer Block
Transformer Block
Tokenize“A dog walked over to the barn…”

GPT-3 / Chinchilla / …

Imagen U-Nets

Diffusion models use highly-nonstandard architectures…

Linear Decoder

Transformer Block
Transformer Block
Transformer Block
Transformer Block
Tokenize“A dog walked over to the barn…”

GPT-3 / Chinchilla / …

Imagen U-Nets

Language models are having all the fun

Scaling of the GPT Family for NLG

Scaling of the U-Net Family for Diffusion
Imagen. Google 2022.
Scaling Laws for Neural Language Models. OpenAI 2020.

“Diffusion models are about denoising images andso a convolutional inductive bias that encodes
locality is critical!! Furthermore, the U-Net
architecture contains many components that are
fundamentally important like long-range skip
connections and an encoder-decoder structure!”

Can diffusion models benefit from simple, scalable architectures?

Scalable Diffusion Models with Transformers, Peebles and Xie, arXiv 2023.

DiT: A New Class of Diffusion Models

DiT: A New Class of Diffusion Models

DiT: A New Class of Diffusion Models

Diffusion Basics𝑥!

• Sample a real image 𝑥!

z

Diffusion Basics𝑥!

• Sample a real image 𝑥!
• Sample noise z ~ N(0,I)

z

Diffusion Basics𝑥!

• Sample a real image 𝑥!
• Sample noise z ~ N(0,I)• Corrupt 𝑥! with z: call it 𝑥"zDiffusion Basics𝑥!

• Sample a real image 𝑥!
• Sample noise z ~ N(0,I)• Corrupt 𝑥! with z: call it 𝑥"• Train a neural net to predict z from 𝑥"zLatent Diffusion Models

High-Resolution Image Synthesis with Latent Diffusion Models. Rombach et al., 202

Important to Scaling: Measuring Model Complexity

Measuring Model Complexity

Parameter counts = badFLOPs = much better (*but not perfect)

DiT DesignLatent DiTs: Training

Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

Latent DiTs: Training

VAE
Encoder(8x ↓)
Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

SpatialLatentLatent DiTs: Training

VAE
Encoder(8x ↓)
Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

SpatialLatentForwardProcess

NoisedLatentLatent DiTs: Training

VAE
Encoder(8x ↓)
Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

SpatialLatentForwardProcess

NoisedLatent

DiTLatent DiTs: Training

VAE
Encoder(8x ↓)
Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

SpatialLatentForwardProcess

NoisedLatentDiT

NoisePrediction

CovarianceLatent DiTs: Training

VAE
Encoder(8x ↓)
Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

SpatialLatentForwardProcess

NoisedLatentDiT

NoisePrediction

Training losses
CovarianceLatent DiTs: Training

VAE
Encoder(8x ↓)
Input imaget
Forward process
timestepc
Conditional info(label, text, etc.)

SpatialLatentForwardProcess

NoisedLatentDiT

NoisePrediction

Training losses
CovarianceDiT Design Space: Transformer Block

Nx

NoiseΣ

32 x 32 x 432 x 32 x 4

DiT+
MLPt

Embed

c

NoisedLatentDiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4

LinearLayer Norm
Nx

Patchify+
MLPt

DiT Block

Embed

c

NoisedLatentxN

DiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4

LinearLayer Norm
Nx

Patchify+MLPt

DiT Block

Embed

c

NoisedLatentxN

DiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4+

LinearPointwise
Feedforward

Layer Norm
Nx

Layer Norm+

DiT BlockxN

Multi-Head
Self-Attention

Layer Norm

Patchify+
MLPt

Concatenate
on Sequence
Dimension

EmbedcNoisedLatentImage TokensConditioning(a) In-Context Conditioning

DiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4+

LinearPointwise
Feedforward

Layer Norm
Nx

Layer Norm+

DiT BlockxN

Multi-Head
Self-Attention

Layer Norm

Patchify+
MLPt

Concatenate
on SequenceDimension

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

DiT Design Space: Transformer Block
NoiseΣ

32 x 32 x 432 x 32 x 4++Linear

Pointwise
Feedforward

Pointwise
Feedforward

Layer Norm
Nx

Layer Norm

Layer Norm++

DiT BlockxN

Multi-Head
Cross-Attention

Multi-Head
Self-Attention

Layer Norm

Layer Norm

Patchify+
MLPt+
Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Layer Norm

EmbedcNoisedLatentImage TokensConditioning(a) In-Context Conditioning

Image Tokens

Conditioning(b) Cross-Attention

DiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4++Linear

Pointwise
Feedforward

Pointwise
Feedforward

Layer Norm
Nx

Layer Norm

Layer Norm++DiT BlockxN

Multi-Head
Cross-Attention

Multi-Head
Self-Attention

Layer Norm

Layer Norm

Patchify+
MLPt+
Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Layer Norm

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

Image Tokens

Conditioning(b) Cross-Attention

DiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4++LinearPointwise
Feedforward

Pointwise
Feedforward

Layer Norm
Nx

Layer NormxNMLPtScale, Shift+

Layer Norm

Multi-Head
Cross-Attention

Multi-Head
Self-Attention+

Layer Norm

Layer Norm+

Patchify+

Pointwise
Feedforward

Layer Norm+

DiT Block+Multi-Head
Self-Attention

Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Scale, Shift

Layer Norm

Layer Norm

MLP

Image Tokens

Conditioning

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

Image Tokens

Conditioning(b) Cross-Attention(c) Adaptive Layer Norm

DiT Design Space: Transformer Block
NoiseΣ32 x 32 x 432 x 32 x 4++Linear

Pointwise
Feedforward

Pointwise
Feedforward

Layer Norm
Nx

Layer Norm

xN

MLPt

Scale, Shift+

Layer Norm

Multi-Head
Cross-Attention

Multi-Head
Self-Attention+

Layer Norm

Layer Norm+

Patchify+

Pointwise
Feedforward

Layer Norm+

DiT Block+Multi-Head
Self-Attention

Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Scale, Shift

Layer Norm

Layer Norm

MLP

Image Tokens

Conditioning

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

Image TokensConditioning(b) Cross-Attention(c) Adaptive Layer NormAside: Identity Initialization in ResNets

Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. Goyal et. al 2018.

Aside: Identity Initialization in ResNets

Original U-Net code from Ho et. alAside: Identity Initialization in ResNets

Not mentioned inpaper, but superimportant in practice!

Original U-Net code from Ho et. al

DiT Design Space: Transformer Block
Noise

Σ

32 x 32 x 432 x 32 x 4++Linear

Pointwise
Feedforward

Pointwise
Feedforward

Layer Norm
NxLayer NormxNMLPtScale, Shift+

Layer Norm

Multi-Head
Cross-Attention

Multi-Head
Self-Attention+

Layer Norm

Layer Norm+

Patchify+

Pointwise
Feedforward

Layer Norm+

DiT Block+Multi-Head
Self-Attention

Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Scale, Shift

Layer Norm

Layer Norm

MLP

Image Tokens

Conditioning

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

Image Tokens

Conditioning(b) Cross-Attention(c) Adaptive Layer Norm

DiT Design Space: Transformer BlockNoiseΣ

32 x 32 x 432 x 32 x 4+

LinearPointwise
FeedforwardLayer Norm
Nx

Layer Norm

xN

MLPt

Pointwise
Feedforward

Scale
Pointwise
Feedforward

Scale, Shift+

Layer Norm

Multi-Head
Cross-Attention

Multi-Head
Self-Attention+

Layer Norm

Layer Norm

Scale+Patchify++

Layer Norm+

DiT Block+Multi-Head
Self-Attention

Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Scale, Shift

Layer Norm

Layer Norm

MLP

Image Tokens

Conditioning

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

Image Tokens

Conditioning(b) Cross-Attention(d) Adaptive Layer Norm ZeroDiT Design Space: Transformer Block
NoiseΣ32 x 32 x 432 x 32 x 4+

LinearPointwise
Feedforward

Layer Norm
Nx

Layer Norm

xN

MLPt

Pointwise
Feedforward

Scale
Pointwise
Feedforward

Scale, Shift+

Layer Norm

Multi-Head
Cross-Attention

Multi-Head
Self-Attention+

Layer Norm

Layer Norm

Scale+Patchify++

Layer Norm+

DiT Block+Multi-Head
Self-Attention

Multi-Head
Self-Attention

Concatenate
on Sequence
Dimension

Scale, Shift

Layer Norm

Layer Norm

MLP

Image Tokens

Conditioning

EmbedcNoisedLatentImage Tokens

Conditioning(a) In-Context Conditioning

Image TokensConditioning(b) Cross-Attention(d) Adaptive Layer Norm Zero

DiT Block

DiT Block

ADM/LDM U-Net BlockScaling DiT

DiT Design Space: Patch Size

𝑝𝐼𝑝

Noised Latent
I×I×C𝐼

DiT Design Space: Patch Size

DiT Block𝑝𝑝

Noised Latent
I×I×C

𝐼

Input Tokens L × C

𝐿 = 𝐼/𝑝 (𝐼DiT Design Space: Patch Size
Design space: p = 2, 4, 8
DiT Block𝑝𝑝

Noised Latent
I×I×C

𝐼

Input Tokens L × C

𝐿 = 𝐼/𝑝 (𝐼DiT Design Space: Model Size

DiT Design Space: Model Size
Design space: S, B, L, XL model configs

DiT Design Space

Block design: adaLN, cross-attention, in-context blocks
Patch size:

2, 4, 8Model size:

S, B, L, XL configs

Experiments

DiT Training Recipe
Training hyperparameters are almost entirely retained fromADM, and identical across all model sizes and patch sizes.
Did not tune learning rates, decay/warm-up schedules, Adam
β1/β2 or weight decays.
Stability issues with U-Net.

More compute = better generative models

More tokens

More compute = better generative models

More tokens

More compute = better generative models

More tokens

More compute = better generative models

More tokens

More compute = better generative models

Bigger DiTs

Bigger model

Smaller patch size

Visual ScalingBigger model

Smaller patch size

Visual ScalingBigger model

Smaller patch size

Visual ScalingBigger model

Smaller patch size

Visual ScalingBigger model

Smaller patch size

Visual ScalingBigger model

Smaller patch size

Visual Scaling

Compute is all that matters

Compute is all that matters

Compute is all that matters

Sampling compute cannot compensate for model computeDiT-XL/2 512x512 Samples

DiT-XL/2 512x512: High GuidanceDiT-XL/2 512x512: High GuidanceDiT-XL/2 512x512: High GuidanceDiT-XL/2 512x512: High GuidanceDiT-XL/2 512x512: High GuidanceDiT-XL/2 256x256: Effects of Guidance Scales = 1.5

s = 2.0s = 4.0s = 6.0

DiT Latent Walk (512x512)

DiT Class Embedding Walk (512x512)

DiT Class Embedding Walk (512x512)

DiT Dogball

BigGAN DogballComparison to SOTA generative models

Comparison to SOTA generative models

Comparison to SOTA generative models

Code: github.com/facebookresearch/DiTAlso available in the Hugging Face diffusers libraryOkay, But why I see people arestill using UNet?
DiT with COCO captions

a giraffe standing in front of agrassy plain and blue sky.

a big black bear laying in thegrass surrounded by trees.

two images of open suitecases
full of toiletries.