data
DiTConfig
extends
DiffusionModelConfigDiTConfig(sample_size: int | tuple[int, int] = 32, in_channels: int = 4, out_channels: int = 8, act_fn: GenerativeActivation = 'silu', num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: BetaSchedule = 'linear', prediction_type: Literal['epsilon', 'sample', 'v_prediction'] = 'epsilon', patch_size: int = 2, hidden_size: int = 1152, depth: int = 28, num_heads: int = 16, mlp_ratio: float = 4.0, frequency_embedding_size: int = 256, num_classes: int = 1000, class_dropout: float = 0.1, conditioning: DiTConditioning = 'adaln_zero', learn_sigma: bool = True)Configuration for the DiT family.
Parameters
sample_sizeint= 32Side of the latent the transformer sees —
32 for the
256-pixel models and 64 for the 512-pixel ones, both after
the VAE's factor-of-eight downsample.in_channelsint= 4Latent channels.
out_channelsint= 8What the decoder emits:
2 * in_channels, a noise and a
diagonal covariance. Set it equal to in_channels for a model
that predicts noise alone and fixes the variance.patch_sizeint= 2Side of the square patch each token covers. The number after the
slash in the paper's names —
XL/2 is the XLarge backbone at
patch 2. Halving it quadruples the token count and at least
quadruples Gflops, which is the axis the scaling study rides.hidden_sizeint= 1152Transformer width.
depthint= 28Number of blocks.
num_headsint= 16Attention heads per block.
mlp_ratiofloat= 4.0Feed-forward expansion.
frequency_embedding_sizeint= 256Width of the sinusoid the timestep is expanded to before the
two-layer MLP that maps it to
hidden_size. The reference
implementation fixes this at 256 independently of the model's
width, so it is a field rather than a reuse of hidden_size —
tying the two would add (hidden_size - 256) * hidden_size
parameters to every variant and make published checkpoints
unloadable.num_classesint= 1000Label vocabulary. Index
num_classes is the null embedding
classifier-free guidance drops to.class_dropoutfloat= 0.1Probability of replacing the label with that null embedding
during training.
conditioning(adaln_zero, adaln, cross_attention, in_context)= "adaln_zero"Which of the paper's four designs to build.
learn_sigmabool= TrueWhether the decoder predicts the covariance alongside the noise.
Kept explicit because it is what makes
out_channels twice
in_channels, and a reader who changes one without the other
gets a shape error rather than an explanation.Notes
Reference: Peebles and Xie, "Scalable Diffusion Models with
Transformers", ICCV, 2023 (arXiv:2212.09748). Backbone
configurations are Table 1; the diffusion hyperparameters are ADM's,
which is a linear variance schedule over 1000 steps from 1e-4 to
2e-2 — the defaults DiffusionModelConfig already carries.
Examples
>>> from lucid.models.generative.dit import DiTConfig
>>> config = DiTConfig()
>>> config.depth, config.hidden_size, config.num_heads
(28, 1152, 16)
The token count is what the scaling study varies, and it comes from
the patch size rather than the width:
>>> DiTConfig(patch_size=2).num_patches, DiTConfig(patch_size=8).num_patches
(256, 16)Used by 3
Constructors
1dunder
__init__
→None__init__(sample_size: int | tuple[int, int] = 32, in_channels: int = 4, out_channels: int = 8, act_fn: GenerativeActivation = 'silu', num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: BetaSchedule = 'linear', prediction_type: Literal['epsilon', 'sample', 'v_prediction'] = 'epsilon', patch_size: int = 2, hidden_size: int = 1152, depth: int = 28, num_heads: int = 16, mlp_ratio: float = 4.0, frequency_embedding_size: int = 256, num_classes: int = 1000, class_dropout: float = 0.1, conditioning: DiTConditioning = 'adaln_zero', learn_sigma: bool = True)Initialise the block. See the class docstring for parameters.