张芷铭的个人博客

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

核心原理

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

$$\text{AdaLN}(x, c) = \gamma(c) \cdot \frac{x - \mathbb{E}[x]}{\sqrt{\text{Var}[x] + \epsilon}} + \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 实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

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 训练时 $\epsilon$ 调整为 $10^{-4}$

Comments