Which is the better way to use AdaLayerNormZero?

[JohnHerry]

I had read the code of class AdaLayerNormZero, it seems there are two ways to use this module on timestep_embedding, One way is the emb will be a value from outside this module, Another way, is to set the num_embeddings initial param to make this AdaLayerNormZero module to keep its time_embedding layer weights.
The first way will save time_embedding weights (in the self.emb=CombinedTimestepLableEmbeddings), to let all DiT AdaLayerNormZero blocks share a single set of time_embedding weights; While the second way makes each AdaLayerNormZero keep independent time_embedding weights.
The second way is more costly, but will it make better performance at the whole model?

[demirciberk]

Confirmed from source (src/diffusers/models/normalization.py, https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py):

python
class AdaLayerNormZero(nn.Module):
def init(self, embedding_dim: int, num_embeddings: int | None = None, norm_type="layer_norm", bias=True):
super().init()
if num_embeddings is not None:
self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim)
else:
self.emb = None
...
def forward(self, x, timestep=None, class_labels=None, hidden_dtype=None, emb=None):
if self.emb is not None:
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
emb = self.linear(self.silu(emb))
...

So num_embeddings is not "better vs worse", it toggles whether the module owns its own CombinedTimestepLabelEmbeddings (timestep sinusoidal embed + class-label embed + MLP) or expects a precomputed emb tensor from outside.

Where each pattern is actually used:

• Shared-embedding pattern (num_embeddings=None, emb passed in): FluxTransformerBlock (transformer_flux.py) — self.norm1 = AdaLayerNormZero(dim) then both norm1 and norm1_context receive the same externally-computed emb each forward: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py. SD3, PixArt-style DiT-derivative models follow the same approach.
• Per-block-embedding pattern (num_embeddings set): DiTTransformer2DModel (dit_transformer_2d.py), which is the original Peebles & Xie "Scalable Diffusion Models with Transformers" (DiT) architecture, https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/dit_transformer_2d.py — self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) for every block.

Facts vs hypothesis on "which is better":

Confirmed:

• The shared pattern computes the timestep/label embedding MLP once per forward pass and reuses it across all blocks; only the per-block self.linear (the AdaLN modulation projection producing shift/scale/gate) differs per block. This matches the original DiT paper's design, where a single TimestepEmbedder/LabelEmbedder feeds all blocks and each block only has its own small adaLN-Zero linear layer (see the original DiT repo: https://github.com/facebookresearch/DiT/blob/main/models.py).
• The per-block pattern in DiTTransformer2DModel instantiates a full separate CombinedTimestepLabelEmbeddings (sinusoidal timestep embed + label embed + 2-layer MLP) per block, which is a strictly larger parameter count and more compute than the shared version, and diverges from the original DiT paper's architecture as published.
• diffusers uses num_embeds_ada_norm in DiTTransformer2DModel primarily to match a specific checkpoint format/config it loads (it is a legacy/compat path for that particular model implementation), not a documented performance optimization.

Hypothesis (no evidence found):

• There is no ablation, benchmark, or comment in the diffusers repo, its PRs, or the original DiT paper claiming that giving every block its own timestep/label embedding weights improves output quality. The DiT paper's own experiments used the shared-embedding design that later became the "pass emb from outside" pattern in diffusers' newer models (Flux, SD3, PixArt).
• If you want to test whether per-block independent embeddings genuinely help, you'd need your own ablation (train small DiT variants with num_embeddings=None+shared emb vs. one embedding module per block, same total parameter budget) — nothing in this repo or its linked papers answers that empirically.

Practical recommendation: unless you specifically need to reproduce/finetune from DiTTransformer2DModel's existing checkpoint format, prefer the shared-embedding pattern (num_embeddings=None, compute emb once, pass it into every AdaLayerNormZero.forward(..., emb=emb)), as used by Flux/SD3 — it is both the architecturally faithful choice and cheaper. There is no published evidence in this repo that per-block embedding weights improve model quality; if you find or run such an ablation, that would need to be validated separately.