AdaLN自适应层归一化：从原理到实现，深度拆解生成模型的核心组件

AdaLN 将条件向量动态生成 LayerNorm 的缩放和平移参数，是 DiT、文生图/文生视频大模型的核心组件。

核心原理

标准 LayerNorm 固定参数无法适配动态条件信号，AdaLN 将 $\gamma$ 和 $\beta$ 替换为条件向量动态生成的函数：

$\text{AdaLN}(x,c)=\gamma (c)\cdot \frac{x-\mathbb{E}[x]}{\sqrt{\text{Var}[x]+ϵ}}+\beta (c)$

• $x$：输入特征 $[B,N,D]$
• $c$：条件向量 $[B,C]$
• $\gamma (c),\beta (c)$：动态生成的缩放和平移参数

从 LN 到 AdaLN

归一化	特点
BatchNorm	对批次大小敏感，不适配 Transformer
LayerNorm	固定 $\gamma ,\beta$，成为 Transformer 标准
AdaLN	动态 $\gamma (c),\beta (c)$，适配条件生成

AdaLN-Zero

DiT 论文提出的工业级变体，核心设计：

1. 生成 6 个动态参数：${\gamma }_1,{\beta }_1,{\alpha }_1$（注意力分支）和 ${\gamma }_2,{\beta }_2,{\alpha }_2$（FFN 分支）
2. $\alpha$ 零初始化，实现恒等映射

计算流程：

$x_{\text{attn}}=\text{MHSA}({\gamma }_1\cdot \text{LN}(x)+{\beta }_1)$ $x=x+{\alpha }_1\cdot x_{\text{attn}}$

零初始化让初始阶段 Transformer Block 输出等于输入，解决深层模型训练不稳定问题。

PyTorch 实现

class AdaLN(nn.Module):
    def __init__(self, feature_dim, condition_dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.gamma_mlp = nn.Sequential(
            nn.Linear(condition_dim, feature_dim), nn.SiLU())
        self.beta_mlp = nn.Sequential(
            nn.Linear(condition_dim, feature_dim), nn.SiLU())
        nn.init.ones_(self.gamma_mlp[0].bias)
        nn.init.zeros_(self.beta_mlp[0].bias)
 
    def forward(self, x, condition):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        gamma = self.gamma_mlp(condition).unsqueeze(1)
        beta = self.beta_mlp(condition).unsqueeze(1)
        return gamma * x_norm + beta

适用场景

场景	条件类型
文生图/文生视频	时间步、文本嵌入
风格迁移	风格特征向量
多模态大模型	模态类型、模态特征

工程实践

• 条件向量先经 LayerNorm 归一化
• 条件映射网络 1-2 层即可，无需过深
• 多条件融合：先融合后映射，避免参数爆炸
• FP16 训练时 $ϵ$ 调整为 $10^{-4}$