Diffusers 归一化层详解AdaLN、adaLN-Zero 与 RMSNorm 在扩散 Transformer 中的实现【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本文基于 Diffusers 官方 API 文档docs/source/en/api/normalization.md与源码src/diffusers/models/normalization.py系统讲解 Diffusers 为支持各类扩散模型而封装的 15 个归一化层从基础的AdaLayerNorm到adaLN-Zero、adaLN-Single再到面向特定模型PixArt、SD3.5、CogVideoX、CogView3、Lumina、Mochi的定制变体。读完本篇你将理解扩散 TransformerDiT中“条件调制归一化”的通用范式——即如何把时间步/条件嵌入投影为 shift、scale、gate 参数来控制归一化后的特征流并能读懂每个类的构造参数、forward 语义及返回值约定。一、为什么扩散模型需要专门的归一化层扩散模型的骨干UNet 或 DiT Transformer除了对输入特征做归一化外还需要一种机制把条件信息扩散时间步t、类别标签、分辨率/宽高比、文本上下文等注入到每一层。Diffusers 将这类“归一化 条件调制”的组合封装为统一的models.normalization模块位于 normalization.py。从源码结构看整个模块的类可以归为四大族基础自适应归一化AdaLayerNormshift/scale 调制、AdaGroupNormGroupNorm 版本、AdaLayerNormContinuous连续条件向量版本adaLN-Zero 家族在 shift/scale 之外额外输出 gate门控参数且权重常初始化为零用于 SD3/Flux 等 DiT 的残差门控——包括AdaLayerNormZero、AdaLayerNormZeroSingle、SD35AdaLayerNormZeroX、CogVideoXLayerNormZero、CogView3PlusAdaLayerNormZeroTextImage、LuminaRMSNormZero、MochiRMSNormZero单步条件嵌入族AdaLayerNormSinglePixArt-Alpha 风格把分辨率/宽高比并入条件、LuminaLayerNormContinuous纯归一化RMSNorm、MochiRMSNorm、GlobalResponseNormConvNeXt-v2 风格、LpNorm。此外模块还包含两个内部辅助组件FP32LayerNorm强制 float32 计算的 LayerNorm 子类供AdaLayerNormZero的norm_typefp32_layer_norm使用以及版本兼容的LayerNorm别名torch ≥ 2.1.0 时直接是nn.LayerNorm否则回退到一个带bias参数的自实现版本见 normalization.py。下面按官方 API 文档的类逐一展开并给出源码级的参数与 forward 行为。二、基础自适应归一化层AdaLayerNorm时间步嵌入驱动 shift/scaleAdaLayerNorm是最早引入的自适应 LayerNorm其结构为SiLU → Linear → chunk → (1scale)*norm(x)shift。构造参数见 normalization.py参数类型 / 默认值说明embedding_dimint必填条件嵌入向量维度num_embeddingsint, optional若给定则内部创建nn.Embedding(num_embeddings, embedding_dim)查表嵌入否则调用方直接传入temboutput_dimint, optional默认为embedding_dim * 2norm_elementwise_affinebool默认False归一化层是否带仿射参数norm_epsfloat默认1e-5归一化数值稳定项chunk_dimint默认0决定 chunk 切分维度见下文forward 接收x特征、timestep可选仅当有查表嵌入时和temb可选的条件向量核心逻辑为temb self.linear(self.silu(temb)) if self.chunk_dim 1: # 该分支目前仅 CogVideoX 与 OmniGen 使用顺序为 shift, scale shift, scale temb.chunk(2, dim1) shift shift[:, None, :] scale scale[:, None, :] else: scale, shift temb.chunk(2, dim0) x self.norm(x) * (1 scale) shift值得注意的细节是chunk_dim语义默认chunk_dim0时按 batch 维切分要求temb是 2D 的[B, 2D]向量chunk_dim1时按特征维切分并扩展为[B, 1, D]。源码注释明确说明这一分支“目前专用于 CogVideoX 与 OmniGen”且两分支的解包顺序相反shift, scalevsscale, shift——这是历史遗留的模型差异移植权重时需特别留意。AdaGroupNormGroupNorm 版的自适应调制AdaGroupNorm把 LayerNorm 换成F.group_norm适合卷积式特征见 normalization.py。构造参数embedding_dim、out_dim、num_groups可选act_fn归一化前的激活函数名经get_activation解析与eps1e-5。其 forward 假定输入为 4D[B, C, H, W]且条件emb形状为[B, T, embedding_dim]emb self.linear(self.act(emb)) # [B, T, out_dim*2] emb emb[:, :, None, None] # 展到 [B, T, out_dim*2, 1, 1] scale, shift emb.chunk(2, dim1) x F.group_norm(x, self.num_groups, epsself.eps) x x * (1 scale) shift即每个时间步T对应一组独立的 scale/shift这正是视频 DiT如 CogVideoX 早期实现对时空特征逐帧调用的方式。AdaLayerNormContinuous面向连续条件向量许多 DiT 模型把条件嵌入时间步经 MLP 后的向量直接传入归一化层而不是在层内部完成“时间步 → 嵌入”的编码。AdaLayerNormContinuous就是为此设计见 normalization.py构造参数embedding_dim输出投影维度、conditioning_embedding_dim输入条件维度、elementwise_affineTrue、eps1e-5、biasTrue、norm_typelayer_norm或rms_normforward 为scale, shift chunk(linear(silu(conditioning_embedding)))再norm(x) * (1 scale) shift源码中有一段显式 NOTE归一化层自身带有 scale/shift 参数、而输出紧接着又被条件向量缩放平移这种双重仿射在原始模型代码中如此实现通常建议将elementwise_affine设为False一个易被忽略的细节self.linear前会把条件向量.to(x.dtype)并随后转回原 dtype注释说明这是为兼容 hunyuanDiT 中条件被上采样到 float32 的情况。LuminaLayerNormContinuous只有 scale 的连续条件版LuminaLayerNormContinuous是 Lumina 图像模型专用变体见 normalization.py结构与AdaLayerNormContinuous类似但条件只投影出scale无 shift且可选地附加第二层投影linear_2当out_dim给定时把归一化后的特征再线性投影到out_dimemb self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) x self.norm(x) * (1 emb)[:, None, :] if self.linear_2 is not None: x self.linear_2(x)它支持同样的norm_typelayer_norm/rms_norm与eps1e-5、biasTrue默认值在 transformer_lumina2.py 等 Lumina 系 Transformer 中被引用。三、adaLN-Zero 家族shift/scale gate 的门控归一化adaLN-Zero 范式源自 DiT 论文在 shift/scale 之外再输出若干gate参数归一化特征进入注意力/MLP 后其输出会乘以 gate 再进入残差。Diffusers 中的AdaLayerNormZero是这一范式的主实现见 normalization.py。AdaLayerNormZero构造参数embedding_dim、num_embeddings可选给定则使用CombinedTimestepLabelEmbeddings同时编码时间步与类别标签、norm_typelayer_norm或fp32_layer_norm后者使用FP32LayerNorm保证数值稳定、biasTrue内部线性层把嵌入映射到6 * embedding_dim解包为shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlpforward 返回 5 元组(调制后的 x, gate_msa, shift_mlp, scale_mlp, gate_mlp)。也就是说只返回“注意力输入”的调制结果MLP 分支的 shift/scale/gate 留给调用方DiT 块在 MLP 前再应用。这种“半调制”的返回值约定是阅读 SD3、Flux 等 DiT 块代码时的关键forward里不会看到 MLP 的 scale/shift 被直接用于当前类而是被上层TransformerBlock消费。AdaLayerNormZeroSingle面向“单分支”注意力结构的精简版见 normalization.py线性层输出3 * embedding_dim只解包shift_msa, scale_msa, gate_msaforward 返回(调制后的 x, gate_msa)不含 MLP 调制参数。norm_type仅支持layer_normelementwise_affineFalse, eps1e-6。SD35AdaLayerNormZeroXSD3.5 的扩展变体见 normalization.py线性层输出9 * embedding_dim解包为三组 (shift, scale, gate)shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 emb.chunk(9, dim1)forward 对同一份norm_hidden_states施加两次不同调制msa与msa2返回 7 元组(调制后 x, gate_msa, shift_mlp, scale_mlp, gate_mlp, 第二次调制后的 x2, gate_msa2)。可以推断这是为 SD3.5 双注意力分支结构准备的“一份归一化、两种调制”设计norm_type目前仅支持layer_normeps 固定1e-6。CogVideoXLayerNormZero同时调制特征与文本上下文CogVideoXLayerNormZero见 normalization.py把“调制 latent 特征”与“调制文本 encoder 输出”合并进一个类构造参数为conditioning_dim条件向量维度独立于特征维、embedding_dim特征维度、elementwise_affineTrue、eps1e-5、biasTrue。forward 中线性层把temb映射到6 * embedding_dim并解包为shift, scale, gate, enc_shift, enc_scale, enc_gate用同一个 norm 层分别调制hidden_states与encoder_hidden_states返回(hidden_states, encoder_hidden_states, gate, enc_gate)。与AdaLayerNorm的chunk_dim1分支相配合构成 CogVideoX 的逐帧条件调制链路。CogView3PlusAdaLayerNormZeroTextImage文本/图像双流 12 路调制面向 CogView3-Plus 的双流结构见 normalization.py构造参数为embedding_dim条件维与dim特征维线性层输出12 * dim解包成两流各 6 路shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp emb.chunk(12, dim1)内部持有norm_x与norm_c两个无仿射 LayerNormeps1e-5分别调制x图像流与context文本流。forward 返回 10 元组两流各自的调制后特征 5 路调制参数供 Transformer 块分别消费。LuminaRMSNormZero 与 MochiRMSNormZeroRMSNorm 基座的 zero 调制这两个类把基座从 LayerNorm 换成 RMSNorm且只输出 scale无 shift各解包 4 路scale_msa, gate_msa, scale_mlp, gate_mlpLuminaRMSNormZeronormalization.py构造参数embedding_dim, norm_eps, norm_elementwise_affine注意其线性层输入维度是min(embedding_dim, 1024)——即允许用一个低维瓶颈最大 1024去生成4 * embedding_dim的调制参数再经RMSNorm(embedding_dim, eps)归一化后输出x * (1 scale_msa)返回(x, gate_msa, scale_mlp, gate_mlp)MochiRMSNormZero定义在 transformer_mochi.pyAPI 文档中同样收录构造参数embedding_dim, hidden_dim, eps1e-5其内部使用RMSNorm(0, eps, False)——dim0加上elementwise_affineFalse表示这是一个无参数的纯 RMS 归一化所有缩放职责交给1 scale_msaforward 全程在 float32 下计算再转回原 dtype返回(hidden_states, gate_msa, scale_mlp, gate_mlp)。四、AdaLayerNormSinglePixArt 风格的条件嵌入器AdaLayerNormSingle与前面几类语义不同它不做任何调制而是承担“条件编码 一次线性映射”的完整职责见 normalization.py。源码 docstring 明确引用了 PixArt-Alpha论文 Section 2.3构造参数embedding_dim、use_additional_conditionsFalse内部使用PixArtAlphaCombinedTimestepSizeEmbeddings其中分辨率嵌入维度固定为embedding_dim // 3forward 接受timestep、added_cond_kwargsresolution、aspect_ratio等附加条件、batch_size、hidden_dtype返回(linear(silu(embedded_timestep)), embedded_timestep)——前者是 6 倍维度的调制参数由 DiT 块再 chunk 成 6 路后者是原始条件向量供后续层复用例如投影进文本上下文。因此在 PixArt 类 DiT 中你会看到AdaLayerNormSingle的输出一头供给注意力前的 scale/shift/gate另一头与文本特征拼接这正是use_additional_conditions控制宽高比条件开关的用途。五、纯归一化层RMSNormRMSNorm是标准 RMS 归一化Zhang et al., 2019的实现见 normalization.py构造参数dim、eps必填、elementwise_affineTrue、biasFalse数值策略无论输入 dtype方差计算始终.to(torch.float32)进行再rsqrt(variance eps)缩放若权重是 fp16/bf16还会把特征先转到权重 dtype 再乘权重兼顾精度与显存硬件适配检测到 NPUis_torch_npu_available时走torch_npu.npu_rms_norm硬件算子否则走 PyTorch 纯张量实现。MochiRMSNormMochiRMSNormnormalization.py是 Mochi 模型专用的 RMSNorm 变体同样 float32 计算方差但总是把输出转回输入 dtype且不提供 bias 选项。源码中的 TODO 注释说明当前存在它是为了绕过分片 checkpoint 尚不支持_keep_in_fp32_modules的问题对应仓库 issue #10013未来可被标准RMSNorm取代。GlobalResponseNormGlobalResponseNormGRN来自 ConvNeXt-v2见 normalization.py对 4D 特征按空间维(1, 2)计算 L2 范数gx用其通道均值归一gx / (gx.mean 1e-6)得到nx最终输出gamma * (x * nx) beta x。gamma与beta均为(1, 1, 1, dim)且零初始化——零初始化意味着初始状态下 GRN 是恒等映射训练过程中再逐步学出全局响应的缩放。在 Diffusers 中它被用于 Sana 等高效 DiT 的通道特征调节。LpNormLpNorm是一个无参数薄封装normalization.pyp2, dim-1, eps1e-12默认值下等价于F.normalize即对指定维做单位化可用于方向特征/参考特征的长度约束。六、配套工具LayerNorm 兼容层与 get_normalization 工厂normalization.py底部还有两个实用组件虽然不在 API 文档的类清单中却是理解各模型配置的关键LayerNorm兼容别名normalization.pytorch ≥ 2.1.0 时就是nn.LayerNorm旧版本回退到自实现唯一差异是支持独立的bias开关torch 2.1 的nn.LayerNorm亦如此。AdaLayerNormContinuous与LuminaLayerNormContinuous中的norm用的就是它get_normalization工厂normalization.py按norm_type字符串batch_norm/layer_norm/rms_norm实例化对应模块统一了eps、elementwise_affine、bias的语义。例如 ResNet 块配置rms_norm时即获得上文所述的RMSNorm配置非法值会直接抛出ValueError。七、这些归一化层在哪些模型中生效normalization模块是 Diffusers 模型层models包内部的基础设施全库有数十个 Transformer、UNet 块文件直接from ...normalization import ...例如 transformer_sd3.py、transformer_flux.py、pixart_transformer_2d.py、cogvideox_transformer_3d.py、transformer_cogview3plus.py、transformer_mochi.py、transformer_lumina2.py、sana_transformer.py 等此外 UNet 侧的 unet_2d_blocks.py、resnet.py、attention_processor.py 也引用其中的归一化工具。结合前文的类定义可以归纳出选型对应关系模型/场景典型归一化类关键特征SD3 / Flux 系 DiTAdaLayerNormZero6 路 shift/scale/gatefp32_layer_norm可选SD3.5SD35AdaLayerNormZeroX9 路输出双注意力分支PixArt-αAdaLayerNormSingle分辨率/宽高比并入条件无内部调制CogVideoXAdaLayerNormchunk_dim1CogVideoXLayerNormZero逐帧时间条件文本流同 norm 调制CogView3-PlusCogView3PlusAdaLayerNormZeroTextImage12 路双流调制Lumina 图像模型LuminaLayerNormContinuous/LuminaRMSNormZero1024 瓶颈投影仅 scale 无 shiftMochiMochiRMSNormZero/MochiRMSNorm无参 RMS 基座float32 计算Sana 系 DiTGlobalResponseNorm零初始化 GRN 全局响应调制通用配置get_normalizationRMSNorm/LayerNorm按norm_type字符串统一构造八、实践要点与注意事项返回值约定决定调用方式adaLN-Zero 系列返回的是元组调制特征 若干调制参数直接打印或断言其输出形状前应先确认对应forward的解包顺序例如AdaLayerNormZero的 5 元组与SD35AdaLayerNormZeroX的 7 元组不可混用scale/shift 解包顺序存在分支差异AdaLayerNorm在chunk_dim0与chunk_dim1两个分支的解包顺序相反scale, shiftvsshift, scale源码注释指明这是 CogVideoX/OmniGen 的历史差异跨模型移植权重时不可想当然数值精度RMSNorm/MochiRMSNorm的方差计算、fp32_layer_norm选项与MochiRMSNormZero的 float32 前向都是针对 fp16/bf16 训练中归一化数值溢出问题的工程处理在低精度微调时保留这些默认行为可避免梯度异常elementwise_affine的双重仿射问题AdaLayerNormContinuous与LuminaLayerNormContinuous的源码注释都提示条件向量已经提供了 shift/scale 时内部归一化层的仿射参数是冗余的建议按原始模型配置决定是否关闭eps 取值不统一各类默认 eps 在1e-5多数类、1e-6adaLN-Zero 家族、1e-12LpNorm之间不同复现第三方权重时应以对应模型 config 中的实际取值为准而不是套用某一类的默认值。小结Diffusers 的models.normalization模块把扩散模型中“归一化 条件调制”的多种变体收敛为 15 个可复用的层AdaLayerNorm奠定 shift/scale 调制的通用范式adaLN-Zero家族在其上增加 gate 门控以适配 DiT 残差结构AdaLayerNormSingle、AdaLayerNormContinuous、Lumina*等则针对 PixArt、SD3.5、CogVideoX、CogView3、Lumina、Mochi 等具体模型的条件形态做了定制RMSNorm、GlobalResponseNorm、LpNorm提供纯归一化底座get_normalization工厂统一按norm_type字符串构造。理解这些类的构造参数、forward 返回值约定与精度处理策略是阅读、复现乃至为 Diffusers 新增 DiT 模型 Transformer 块的前置知识。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考