资讯详情

Diffusion Transformer数学推导与PyTorch最小可跑实现

📅 2026/10/27 21:05:52 | 华诺云谱 👁 阅读
Diffusion Transformer数学推导与PyTorch最小可跑实现
第一次把 diffusion transformer 的代码从公式抄到能跑我在单卡机器上盯着 loss 曲线看了三个小时——前 200 步掉得很漂亮第 300 步之后开始上下抖动最后停在某个明显偏高的平台上不再下降。那会儿我把问题归咎于学习率、数据、batch size折腾了一整圈最后发现是我把 adaLN 的调制分支初始化成了常规的正态分布导致每个 block 在训练一开始就把主干特征搅乱了。这个坑让我意识到一件事diffusion transformer 这种东西光会调库是不够的前向加噪那套数学、Transformer 里的条件注入方式、损失函数到底在优化什么这三块必须自己在纸上推一遍否则出问题时你连从哪儿查都不知道。这篇是我这个系列的第一部分只讲“数学地基”和“最小可跑实现”。具体来说会把扩散过程的前向链和闭式解拆开算一遍、把 Transformer 那半边注意力、patchify、时间步条件注入、adaLN-Zero 的零初始化的数学形式写清楚、然后把两边接起来的损失函数推导和 PyTorch 代码骨架给出来。适合两类人看一类是已经跑过 Stable Diffusion 那套 U-Net 架构、想搞明白为什么现在大家都在往 Transformer 骨架上迁的人另一类是被公式劝退过、想找一个“从符号到代码”的对照表的人。不需要你会推导变分下界但需要你知道什么是高斯分布、什么是矩阵乘法、能看懂 PyTorch 的nn.Module。系列第二部分我会讲采样器、CFG、以及把模型真正训到出图的那套工程细节这篇先把地基打完。1. 先把问题摆清楚为什么扩散模型要换掉 U-Net1.1 卷积骨架的天花板在哪U-Net 在扩散模型里统治了好几年它的优势是明确的卷积自带局部性和平移等变性这两个归纳偏置天然适合图像这种二维网格数据skip connection 让浅层的高频细节能直接送到深层重建质量好而且下采样-上采样的结构让感受野随深度指数级扩大计算量却是线性的。这套设计在 512×512 这个尺度上极其高效我在实际项目里用 U-Net 训 256×256 的模型单卡 24G 显存能把 batch 开到 64 左右收敛也很稳。但它的局限同样明显而且随着分辨率上去会越来越痛。第一感受野是靠堆层数堆出来的要建模“左上角的物体和右下角的物体之间的关系”这种长程依赖你得堆很深而且中间要经过多次下采样信息在瓶颈处已经被压得很扁了。第二U-Net 的超参结构非常“手工”——几层、每层多少个 block、attention 加在哪几个 resolution 上这些选择没有统一的最优解换个分辨率就得重新调结构。第三也是最关键的当你希望模型规模能像语言模型那样“加大参数量就变强”的时候U-Net 的 scaling 曲线不够平滑加大 depth 会带来梯度问题加大 width 的收益又很快饱和。Transformer 恰好补的就是这三块。自注意力是全局的任意两个 token 之间都是一步可达长程依赖不需要靠深度去换结构极其规整所有 block 长得一模一样扩参就是简单地把层数或宽度乘上去而且它已经在语言模型上被验证过 scaling law这套经验可以直接搬。这就是 diffusion transformer 这条路的核心动机——不是为了时髦是为了拿到一条更可预测的规模化曲线。1.2 把图像当成句子读patchify 的核心直觉Transformer 吃的是序列图像是二维网格中间得有个转换。这个转换就是 patchify把 H×W×C 的图切成 p×p 的小块每块拉平成一个向量再线性投影到 d 维就得到了一串长度为 N(H/p)×(W/p) 的 token 序列。这个操作你可以理解成“把图像当成一句话来读每个 patch 是一个词”。这里有个非常容易搞混的点我必须提前说清楚如果你在像素空间做 patchifytoken 数会爆炸。256×256 的图p2 的话是 128×12816384 个 token注意力复杂度是 O(N²)16384² 这个量级根本跑不动。所以实际做法几乎都是在潜空间latent space上做——先用一个自编码器把 256×256×3 压成 32×32×4 之类的潜表示再在这个 32×32 的网格上 patchifyp2 得到 256 个 token。256 个 token、768 维隐层这个规模就完全可控了注意力矩阵是 256×256一次前向的计算量在单卡上毫无压力。为什么 p 取 2 而不是 4 或 8这是计算量和表达力的折中。p 越小token 越多空间细节保留得越好但注意力成本按平方涨p 越大token 越少越省算力但每个 token 融合的空间范围变大细节会糊。在 32×32 的潜空间上p2 得到 256 个 tokenp4 只有 64 个——64 个 token 描述一张图实在太少了你会明显看到生成结果的结构崩塌。我自己实测过 p4 的配置FID 会差一大截而且小物体的边缘会糊成一片。所以业界的默认配置基本是潜空间 32×32 配 p2。1.3 三种骨架的横向对比在动手写代码之前先把选型的账算清楚。下面这张表是我自己在几个项目里对比下来总结的不是绝对结论但能帮你判断自己该走哪条路。维度纯 U-Net纯 Transformer (DiT)混合卷积下采样 Transformer 主干长程建模靠深度堆间接全局一步可达兼顾但结构复杂结构规整度低超参多高block 完全同构中扩参可预测性一般好接近线性可预测中高分辨率成本线性增长友好平方增长需要分块或窗口介于两者之间训练稳定性好成熟需要 adaLN、qk-norm 等技巧好上手难度低教程多中细节坑多高我的建议很直接如果你目标是 256×256 到 512×512 这一档、想验证一个新想法、或者需要可控的扩参实验走纯 Transformer 骨架如果你要做的是超大分辨率、或者是需要极强局部细节的任务比如某些图像修复场景卷积的局部性还是不能完全丢掉混合结构更务实。别一上来就追求“纯血”架构选择是为目标服务的不是为信仰服务的。2. 扩散过程那一半数学从加噪到闭式采样2.1 前向链每一步只加一点点高斯噪声扩散模型的前向过程定义成一个马尔可夫链从干净数据 x₀ 出发每一步加一点高斯噪声q(x_t | x_{t-1}) N(x_t; sqrt(1 - beta_t) * x_{t-1}, beta_t * I)这里 beta_t 是一个预先设定的方差调度variance schedule取值在 (0,1) 之间通常从 1e-4 线性或余弦地增长到 0.02总步数 T 一般取 1000。注意均值项里的 sqrt(1-beta_t) 这个系数不是随手写的——它是为了保证方差不爆炸。如果直接写成 x_t x_{t-1} noise每一代的方差都是上一代加 beta走到第 1000 步数值早就飘到天上去了。乘上 sqrt(1-beta_t) 之后每一步都做了一次“方差保留”的缩放整体信号能量被控制在合理范围。我习惯把这一步理解成“往一杯水里滴墨水每次滴一点点同时把杯子里的水倒掉一点点保持总量不变”。倒掉的那部分就是 sqrt(1-beta_t) 带来的收缩它保证了过程不会数值发散。2.2 重参数化为什么可以一步跳到任意 t马尔可夫链是一步一步走的但训练的时候我们显然不想真的循环 1000 次。这里就要用到重参数化reparameterization技巧。令 alpha_t 1 - beta_talpha_bar_t 连乘 alpha_1 到 alpha_t那么可以证明q(x_t | x_0) N(x_t; sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I)对应的采样写法就是最常用的那一行x_t sqrt(alpha_bar_t) * x_0 sqrt(1 - alpha_bar_t) * eps其中eps ~ N(0, I)这个式子的推导用的是高斯分布的可加性两个独立高斯的和还是高斯均值相加、方差相加。反复代入 T 次就能把连乘合并成 alpha_bar。它的意义极大——训练时任何 t 的样本都可以一步采样出来不需要真的跑 t 步链。这就是为什么扩散模型的训练可以完全并行。alpha_bar_t 这个量决定了整个训练的难度分布。我按 T1000、线性调度beta 从 1e-4 到 0.02手算了一组数你可以对照着感受一下talpha_bar_t约含义01.000干净图像1000.896保留约 90% 信号轻微噪声2600.497信噪比约 0 dB半信半疑5000.078只剩 7.8% 信号噪声占主导10004e-5几乎纯噪声算的过程不复杂log(alpha_bar_t) ≈ -sum(beta_i) - sum(beta_i²)/2。以 t500 为例sum(beta_i) ≈ 500×1e-4 (500²/2)×1.99e-5 ≈ 2.5375二阶修正约 0.008所以 log(alpha_bar) ≈ -2.546取指数得到 0.078。这个数字说明一个问题在线性调度下后一半步数基本是在纯噪声里空转。t500 时信号已经只剩不到 8%模型在这个区间能学到的东西非常有限但计算量一点都不少。这就是为什么后来大家改用余弦调度cosine schedule——它让 alpha_bar 的下降曲线更均匀中段的有效训练区间被拉长同样的 T 下采样质量会更好。2.3 反向过程与变分下界MSE 是怎么被“逼”出来的前向是加噪我们真正要学的是反向从纯噪声还原图像。反向过程也建模成高斯p_theta(x_{t-1} | x_t) N(x_{t-1}; mu_theta(x_t, t), sigma_t² * I)关键是均值 mu_theta 怎么参数化。这里有个很漂亮的结论给定 x_t 和 x_0真实的后验 q(x_{t-1} | x_t, x_0) 是解析可得的均值为mu_tilde [sqrt(alpha_{t-1}) * beta_t / (1 - alpha_bar_t)] * x_0 [sqrt(alpha_t) * (1 - alpha_bar_{t-1}) / (1 - alpha_bar_t)] * x_t方差为beta_tilde_t beta_t * (1 - alpha_bar_{t-1}) / (1 - alpha_bar_t)。既然后验是已知的那训练目标就很自然了——让预测分布去逼近真实后验用 KL 散度衡量。把 KL 展开、忽略与 theta 无关的常数项得到的是变分下界ELBO里的一个二次项L_t E[ (1 / (2*sigma_t²)) * || mu_tilde(x_t, x_0) - mu_theta(x_t, t) ||² ]到这里还是“预测均值”。再往下走一步既然 x_t sqrt(alpha_bar_t) x_0 sqrt(1-alpha_bar_t) eps把均值公式里的 x_0 用 x_t 和 eps 反解出来代进去你会发现 mu_theta 和 mu_tilde 的差本质上是“模型预测的噪声”和“真实加的噪声”之差。于是损失可以重写成L_simple E_{t, x_0, eps}[ || eps - eps_theta(x_t, t) ||² ]这就是所有扩散模型里最常见的那一行 MSE。值得强调的是这一步不只是等价变换还丢掉了前面的权重系数 1/(2*sigma_t²)。丢掉权重是个工程决定不是数学必然。理论上加权会更“正确”因为它反映了不同 t 下预测难度不同但实践发现去掉权重后模型反而更容易学因为高噪声区间的样本本来就更难预测如果按理论加权会把这些难样本的权重压得很低反而学不好。这个权衡我在第 4 节还会展开。2.4 采样器的选择DDPM 与 DDIM 的差别到底在哪训练完之后要采样采样器决定了从噪声到图像怎么走回去。最原始的 DDPM 采样是随机性的每一步都注入新的噪声x_{t-1} mu_theta(x_t, t) sigma_t * zz 是标准高斯t0 时t0 时 z0DDIM 的思路是把这个过程改成确定性的。当 eta0 时采样公式变成x_{t-1} sqrt(alpha_bar_{t-1}) * x0_hat sqrt(1 - alpha_bar_{t-1}) * eps_hat其中 x0_hat 是从当前预测的噪声反解出来的干净图像x0_hat (x_t - sqrt(1-alpha_bar_t) * eps_hat) / sqrt(alpha_bar_t)。这个式子有个巨大的好处——它不要求 t 和 t-1 相邻。你可以只取 50 个甚至 20 个时间步从 1000 步里等距跳着走速度提升 20 到 50 倍而图像质量下降很有限。我实测下来20 步 DDIM 和 1000 步 DDPM 在视觉上的差距远小于耗时上的差距。不过要注意一点DDIM 步数太少时会丢细节尤其是纹理丰富的地方容易出现“涂抹感”。我的经验是 50 步是个甜点20 步以下就开始明显掉质量了。这个跟调度方式也有关余弦调度在少步数下比线性调度更稳。3. Transformer 那一半数学注意力、切块与条件注入3.1 多头注意力写成矩阵是什么样子标准的缩放点积注意力Attention(Q, K, V) softmax(Q K^T / sqrt(d_k)) V其中 Q、K、V 都是 [N, d_k] 的矩阵。除以 sqrt(d_k) 的原因是如果 q 和 k 的每个分量都是独立的标准正态那它们的点积方差是 d_k维度一高点积的数值就会很大softmax 之后会退化成近似 one-hot梯度几乎为零。除以 sqrt(d_k) 把方差归一化回 1这是注意力能训起来的关键细节之一。多头就是把 d 维拆成 h 份每份 d/h 维各自做注意力再拼回去。以 d768、h12 为例每个头 64 维。多头的好处不是“算得更准”而是“能同时关注不同类型的关系”——我自己的观察是训完之后有些头明显在管局部纹理有些头在管全局布局虽然没有人显式指派。复杂度这块要算清楚注意力矩阵是 N×N所以时间和显存都是 O(N²d)。N256、d768 时一次注意力的 FLOPs 大约是 4N²d ≈ 2e812 层就是 2.4e9单张 A100 上跑一次前向是毫秒级。但如果你把 N 拉到 1024对应更高分辨率的潜空间成本就是 16 倍这时候就得考虑窗口注意力或者分块了。3.2 patchify 的尺寸账token 数怎么算这个账必须会算不然你连显存炸了都不知道为什么。公式很朴素N (H / p) * (W / p)输入patch 大小token 数注意力矩阵大小32×32 潜空间2256256×25632×32 潜空间46464×6464×64 潜空间210241024×1024256×256 像素空间810241024×1024256×256 像素空间21638416384×16384不可行最后一行就是前面说的那个坑像素空间配小 patch注意力矩阵有 2.68 亿个元素光存这个矩阵 fp16 就要 536MB再乘以层数和 batch直接原地爆炸。所以潜空间不是可选项是必需品。patch 嵌入的实现就是一个卷积nn.Conv2d(in_ch, d, kernel_sizep, stridep)。用卷积而不是真的做切块再拼是因为卷积核和步长都等于 p 的时候它干的事恰好就是“不重叠切块 线性投影”而且 GPU 上的实现效率更高。这个细节很多人第一次看代码会懵觉得怎么没看到 reshape。3.3 时间步和类别条件怎么塞进去扩散模型必须知道当前是第几步不然它没法判断该去掉多少噪声。条件注入方式主要有三种差异很大第一种是in-context conditioning把时间步和类别的嵌入当成额外的 token直接拼到序列前面让注意力自己去处理。这个最省事几乎不改结构但效果一般因为条件信息和图像 token 混在一起模型需要额外的容量去区分它们。第二种是cross-attention在 block 里加一个额外的交叉注意力层query 来自图像 tokenkey 和 value 来自条件嵌入。灵活度最高能处理变长条件比如文本但多了一层注意力和一堆参数训练也更容易不稳。第三种是adaLNadaptive layer norm不用额外 token也不用额外注意力直接把条件嵌入通过一个 MLP 映射成缩放和偏移参数去调制 LayerNorm 的输出。原话是modulate(x, shift, scale) x * (1 scale) shift。这种方式参数最少、计算最省而且效果在类别条件任务上是最好的。代价是它不太适合变长条件所以文本条件通常还是回到 cross-attention。我用下来最直接的感受是如果你的条件维度固定时间步、类别、分辨率、甚至一个固定长度的小向量闭眼选 adaLN如果条件是可变长的文本序列那没得选cross-attention 或者混合方案。3.4 adaLN-Zero为什么零初始化能救训练这是整个架构里最容易踩坑、也最值得讲的一个细节。adaLN 的调制分支一般是这样接的时间步嵌入 → SiLU → Linear(d, 6d) → 拆成 6 个向量shift、scale、gate 各两组一组给注意力一组给 MLP。关键在于这个 Linear 的输出层要做零初始化。为什么因为零初始化之后scale0、shift0x * (10) 0 x也就是每个 block 在训练开始时是恒等映射。整个网络在第一步就等价于“patch 嵌入 最后一层线性”主干是通的梯度能顺畅地从头传到尾。随着训练进行调制参数慢慢偏离零block 才开始真正起作用。如果按常规做正态初始化训练一开始每个 block 都会往特征里注入一个随机的缩放和偏移12 层叠起来特征的方差会指数级放大或缩小——这就是我开头说的那个坑loss 卡在平台上不降看起来像学习率问题其实是初始化问题。顺便说一句gate 那一路也是零初始化等于说残差分支一开始是被完全关掉的这也是 identity 性质的一部分。这个设计还有个附带好处它让训练早期的模型等价于一个浅层模型收敛特别快前三五百步 loss 掉得很陡。我对比过零初始化和普通初始化同样的配置下前者能早收敛大概 20% 的步数。4. 把两边接起来训练目标与损失函数落地4.1 从加权 ELBO 到“去掉权重”的工程取舍第 2 节推到了 L_simple这里把取舍讲透。理论上完整的 ELBO 是每一项带权重的和权重是 1/(2*sigma_t²)而 sigma_t² 在 t 大时接近 1在 t 小时接近 beta_11e-4。也就是说理论目标会给小 t低噪声的项分配高达 10⁴ 量级的权重给大 t 的项分配 1 量级的权重。这个加权看起来合理——低噪声时预测得准不准确实更影响最终画质。但工程上直接这么干会出问题不同 t 的损失量级差四个数量级梯度尺度极不平衡训练会非常不稳定。更麻烦的是高噪声区间的样本被过度压制之后模型学不会怎么从纯噪声里“起形”采样的时候前面几步就崩了后面再怎么修也救不回来。所以通行的做法就是直接扔掉权重所有 t 一视同仁。这个决定牺牲了理论上的最优性换来了训练稳定和采样质量。我自己的经验是如果你发现生成的图整体偏灰、结构感弱可以试试加一个温和的加权比如权重正比于 SNR 的某个幂次但千万别用理论权重那个基本一定崩。4.2 eps / x0 / v 三种预测目标的换算模型可以预测三样东西噪声 eps、干净图 x0、或者速度 v。三者数学上完全等价可以互相换算但训练动力学差别不小。由x_t sqrt(abar) x0 sqrt(1-abar) eps可以解出预测 eps 时x0_hat (x_t - sqrt(1-abar) * eps_hat) / sqrt(abar)v 的定义v sqrt(abar) * eps - sqrt(1-abar) * x0由 v 反解x0 sqrt(abar) * x_t - sqrt(1-abar) * veps sqrt(abar) * v sqrt(1-abar) * x_t。这两个式子代进去验证一下就能确认我当初推的时候还怀疑过符号实际展开就明白了。为什么要有 v 预测因为 eps 预测在 t 很小的时候有问题此时 abar 接近 1x0_hat 的式子里除以 sqrt(abar) 接近除以 1还好但 x_t 和 eps 项的量级关系导致信噪比极高模型稍微预测错一点点反解出来的 x0 就偏得很远损失被放大。用 v 预测相当于在 eps 和 x0 之间做了个平滑插值低噪声时它更接近预测 x0高噪声时更接近预测 eps整体梯度尺度更均衡。我的实测结论是T1000 的常规配置下 eps 预测够用但如果步数很少或者噪声调度很密v 预测会明显更稳。4.3 可跑的 PyTorch 代码骨架时间步嵌入用的是正弦位置编码和 Transformer 里那套一样只是把位置下标换成了时间步import math import torch import torch.nn as nn def modulate(x, shift, scale): return x * (1 scale) shift class TimestepEmbedder(nn.Module): def __init__(self, d, freq_dim256): super().__init__() self.freq_dim freq_dim self.mlp nn.Sequential( nn.Linear(freq_dim, d), nn.SiLU(), nn.Linear(d, d) ) def forward(self, t): half self.freq_dim // 2 freqs torch.exp( -math.log(10000.0) * torch.arange(half, devicet.device) / half ) args t[:, None].float() * freqs[None] emb torch.cat([torch.cos(args), torch.sin(args)], dim-1) return self.mlp(emb)block 的实现要注意两处LayerNorm 关掉仿射elementwise_affineFalse因为缩放偏移由 adaLN 提供调制层的最后一层零初始化。class DiTBlock(nn.Module): def __init__(self, d, heads): super().__init__() self.norm1 nn.LayerNorm(d, elementwise_affineFalse) self.attn nn.MultiheadAttention(d, heads, batch_firstTrue) self.norm2 nn.LayerNorm(d, elementwise_affineFalse) self.mlp nn.Sequential( nn.Linear(d, 4 * d), nn.GELU(approximatetanh), nn.Linear(4 * d, d) ) self.ada nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d)) nn.init.zeros_(self.ada[-1].weight) nn.init.zeros_(self.ada[-1].bias) def forward(self, x, c): s1, sc1, g1, s2, sc2, g2 self.ada(c).chunk(6, dim-1) h modulate(self.norm1(x), s1, sc1) h, _ self.attn(h, h, h, need_weightsFalse) x x g1.unsqueeze(1) * h h modulate(self.norm2(x), s2, sc2) x x g2.unsqueeze(1) * self.mlp(h) return x训练一步的核心就是采样 t、采样 eps、构造 x_t、算 MSEdef train_step(model, x0, opt, T1000): b x0.shape[0] t torch.randint(0, T, (b,), devicex0.device) eps torch.randn_like(x0) abar alpha_bar[t].view(b, 1, 1, 1) # 预先算好的表 x_t abar.sqrt() * x0 (1 - abar).sqrt() * eps pred model(x_t, t) loss ((pred - eps) ** 2).mean() opt.zero_grad() loss.backward() opt.step() return loss.item()注意alpha_bar一定要预先算成一张表放在显存里不要在每个 batch 里现场连乘那个开销白白浪费。另外 t 的采样分布也有讲究均匀采样是最常见的但如果你发现模型在某一段 t 上明显学不好可以对那段做过采样。5. 实操一张 24G 卡上的最小可行 DiT5.1 参数量和显存的纸面计算先算参数量这块算清楚了心里才有底。以 d768、12 层、12 头、p2、潜空间 32×32 为例每个 block 里注意力四个投影矩阵 4d² ≈ 236 万MLP 两个矩阵 8d² ≈ 472 万加起来约 708 万。调制分支 d×d d×6d ≈ 413 万不对算细一点Linear(d, 6d)是 768×4608 ≈ 354 万加上前面的 SiLU 之前的投影可以忽略。所以每个 block 约 1060 万12 层就是 1.27 亿。再看首尾patch 嵌入Conv2d(4, 768, 2, 2)约 1.2 万最后一层Linear(768, p*p*out_ch)约 1.2 万时间步 MLP 大概 40 万。这些加起来占比极小。所以这个配置大概在 1.3 亿参数上下。用 AdamW 训练参数 fp32 是 520MB一阶二阶动量各 520MB梯度 520MB光这些就 2GB 左右。混合精度下参数和梯度用 fp16/bf16 存可以砍一半但优化器状态还是 fp32。激活值的账更麻烦。每个 token 是 768 维256 个 token一个样本一层的激活是 256×768≈19.7 万个元素fp16 下约 0.39MB。一层里需要保存的张量包括注意力的 QKV、注意力输出、两次 LayerNorm 的输出、MLP 中间层这个是大头4d 宽1.57MB加起来一层约 4MB。12 层就是 48MB 一个样本batch64 就是 3GB。加上参数和优化器总计 5 到 6GB24G 卡跑 batch64 是比较宽松的跑到 128 问题也不大。我实测下来的经验值这个配置在 24G 卡上 batch96 是安全线再往上开始 OOM。如果显存吃紧开梯度检查点gradient checkpointing激活显存能砍到 1/3 左右代价是训练速度慢 20% 到 30%。5.2 训练循环的关键几行除了上面的单步逻辑还有几个地方必须处理对。第一EMA指数滑动平均几乎是必开的对采样质量提升非常明显衰减系数取 0.9999。第二是学习率 warmup前 5000 步线性从 0 爬到 1e-4之后余弦衰减这个对稳定性影响很大。第三是梯度裁剪max_norm 取 1.0防止个别 batch 炸掉。ema {k: v.clone().detach() for k, v in model.state_dict().items()} torch.no_grad() def update_ema(model, ema, decay0.9999): for k, v in model.state_dict().items(): if v.dtype.is_floating_point: ema[k].mul_(decay).add_(v.detach(), alpha1 - decay) else: ema[k].copy_(v)还有一个细节我踩过alpha_bar表的精度要用 float64 算完再转 float32不然 t 接近 1000 的时候连乘出来的值下溢到 0sqrt(1-abar)变成 1倒是不影响但1/sqrt(abar)会变成 inf。这个 bug 只在特定 t 上出现非常难查。5.3 训练曲线怎么读我一般看三个信号。第一是 loss 的整体趋势正常的曲线应该是前期陡降、中期缓慢下降、后期在一个平台上小幅波动。如果 1000 步内没有任何下降八成是初始化或者数据的问题别急着调学习率。第二是 loss 的方差。健康的训练里相邻 batch 的 loss 波动应该在一定范围内如果突然出现尖刺又立刻恢复多半是某个 batch 的数据有问题或者梯度爆炸被裁剪掉了。如果波动越来越大那就是学习率太高或者模型开始发散。第三是抽样看效果这个比 loss 曲线重要得多。每 5000 步用固定的几个随机种子采样一次用同样的种子看变化你能很直观地看到模型从“纯噪声”到“有色块”到“有结构”的过程。我一般会存成网格图方便对比。这里有个小技巧用固定的种子序列不要每次随机不然你分不清是模型变好了还是种子抽得好。6. 踩坑记录与速查表6.1 五个最常见的崩点第一个是前面说的 adaLN 零初始化症状是 loss 卡在平台、生成图全是灰蒙蒙的噪声。检查方法很粗暴把模型初始化后直接跑一次前向输出应该接近全零因为最后一层通常也是零初始化或者至少方差很小如果输出方差很大那肯定是某个分支初始化错了。第二个是alpha_bar的索引错位。t 是从 0 开始还是从 1 开始alpha_bar[t]还是alpha_bar[t-1]这个差一位会导致训练时用的噪声水平和采样时对不上症状是训练 loss 很好但采样出来全是糊的。我第一次就栽在这查了两天才发现。第三个是 LayerNorm 没关仿射。如果 adaLN 的调制和 LayerNorm 自带的缩放偏移同时存在两者会打架训练能收敛但效果上限明显低。症状不太明显容易被忽略但对比实验能看出来。第四个是时间步嵌入的频率尺度。用 10000 做底数在 T1000 的时候高频部分的周期远小于步长等于那几维基本没信息。有些实现会把频率范围调窄让 1000 个步长能被充分区分。这个不影响收敛但影响模型对精细噪声水平的判断。第五个是混合精度下的 softmax。fp16 做 softmax 时如果输入有较大的正值exp 会溢出成 inf。现在的框架大多会自动处理把 softmax 转到 fp32但如果你自己写了注意力一定要记得这一条。6.2 问题速查表症状最可能原因排查手段loss 完全不降学习率过大 / 初始化错打印初始前向的输出方差loss 降但采样全糊alpha_bar 索引错位核对训练和采样的 t 定义训练后期突然发散无梯度裁剪 / 学习率未衰减加 max_norm1.0开余弦衰减生成图结构对但纹理脏步数太少 / v 预测未启用采样步数从 20 提到 50显存 OOM 但参数不多激活未做 checkpoint开梯度检查点同样种子结果差异大未开 EMA / 随机性来源多固定种子 开 EMA高噪声步采样崩权重项未去掉确认损失是纯 MSE最后补充一个我自己总结的判断标准如果你不确定某个改动是好是坏别只看 loss一定要用固定种子跑一组采样图对比。扩散模型的 loss 和生成质量之间的相关性没有想象中那么强同一个 loss 水平下采样质量的差异可以很大。我现在的习惯是每改一次结构就存一组 64 张的采样网格用同样的种子这样累积下来就是一个很直观的对比库比任何指标都可靠。
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑