1. 从扩散模型到DiT骨干网络的一次模式切换近两年生成模型圈子里最明显的一个变化就是扩散模型的骨干网络从U-Net逐渐迁移到Transformer架构。DiTDiffusion Transformers就是这个趋势里最有代表性的一支。很多刚接触这个方向的人会有一个疑问U-Net明明已经很成熟了为什么还要费劲去换成Transformer这篇文章就围绕DiT的原理和代码实现展开带你把它从里到外拆一遍看完你不仅知道它在干什么还能直接照着代码上手搭一个能用的版本。适合谁看两类人最合适一是已经写过简单DDPM、现在想进一步理解“工程化扩散模型”该如何设计架构的人二是被各种模型名绕晕想系统搞懂DiT、adaLN、Patchify这些概念到底对应什么模块、解决什么问题的新手。我会尽量把公式掰开讲代码也会给全。1.1 先把扩散模型的老底翻出来扩散模型不是什么魔法它的思路是反着来的。前向过程就是往一张干净图片上不断加高斯噪声第1步加一点第2步加得更多走到几百步以后图片基本成了一团什么信息都没有的纯噪声。反向过程则是让神经网络学会“倒放”从这团噪声里一步步把原图恢复出来。用数字表达前向过程就是q(x_t | x_0) N(x_t; sqrt(α̅_t) * x_0, (1 - α̅_t) * I)意思是在时间步t当前样本可以看作由原图x_0乘以一个衰减系数再加上σ级别的噪声合成。训练目标是让网络预测某个被刻意加进去的噪声εL E_{t,x0,ε} [ || ε_θ(x_t, t) - ε ||^2 ]这是我个人觉得最好用的一个理解方式扩散模型的每一步其实都是在教网络“估计噪声”。一旦噪声估计得准采样时不停减掉噪声就行。类比一下往清水里滴一滴墨水墨迹会慢慢扩散到整杯水训练时你记录了墨水原本那一下的位置采样时就是倒过来做——从整杯均匀的墨水色里一点点把“墨水抽走”。DiT在这个架构里干的活就是那个负责“怎么抽墨水”的智能去噪器。1.2 U-Net做得好好的为什么非要换掉U-Net的成名得益于它对图像的强归纳偏置卷积天然擅长提取局部纹理编码器配解码器加跳连接让细节信息有“直通车”可以跨层传递。这套设计在低分辨率、小规模生成任务上又稳又快。可一旦任务规模变大它的短板就藏不住了。第一U-Net的下采样/上采样结构对空间分辨率有明确设计假设当你想直接增大生成图片尺寸时往往需要改动网络拓扑不能简单“堆参数”。第二U-Net内部的attention机制通常只在瓶颈层启用在模型容量和全局覆盖范围上天然受限。第三也是最重要的一点它不好扩展。这里说的扩展不是指网络做深一点而是指在“数据规模、算力规模、参数规模”都同步增大时能不能继续保持稳定的收益曲线。实验已经证明CNN类骨干达到某个容量之后继续增加参数得到的提升会明显放缓。Transformer的统计效率在NLP和视觉分类领域已经被验证过一轮了。它的核心操作是自注意力每个token可以和任意位置的token交换信息对全局关系的捕获能力特别强。更关键的是它没有太多卷积归纳偏置结构统一、模块标准只需要控制token数量、embedding维度和transformer层数就能平滑地改变模型规模。DiT要做的就是把扩散模型中的U-Net骨干整体替换成这种纯Transformer结构同时保留扩散模型的训练目标。为了直观对比我把两者的特性列了一张表对比维度U-NetDiT主干结构卷积编码器-解码器多层Transformer编码器信息交互卷积局部感知 瓶颈注意力全局自注意力所有token直接交互可扩展性增加通道/层数收益递减可通过patch大小、深度、宽度线性扩展条件注入方式特征图维度拼接与缩放通过调制参数、交叉注意力等方式注入对分辨率的适应依赖固定的下采样路径通过patch划分灵活适应不同输入尺寸实现复杂度网络分支多调优相对繁琐结构统一模块复用性很强当然这不等于U-Net一无是处。在中小尺寸、训练数据有限的场景里U-Net仍然很能打。DiT的意义更多在于把扩散模型的架构选择问题从“手写网络”变成了“配置化搭建”这是否是一个更适合扩展的方向答案在更大计算量上已经比较明显了。2. DiT的架构设计与关键模块拆解DiT的整体流程可以拆成四个环节先把图像切块映射成token序列然后通过一堆Transformer层做特征变换在过程中把时间步和类别标签等条件信息注入进去最后把输出的token序列重新拼回图像尺寸输出噪声预测。这四个环节听起来简单但每个环节都有值得琢磨的细节。2.1 Patchify图像是怎么变成token的Transformer处理的是序列图像是二维格子怎么把它们变成一维token序列DiT的做法和Vision TransformerViT一脉相承叫Patchify。假设输入是32×32分辨率、3通道的图像你设定patch大小为8×8那么每个patch就是一个8×8×3的小立方体。32×32图像被纵向切成4块、横向切成4块一共得到4×416个patch。每个patch通过一个线性层映射为一个embedding向量于是这16个向量就成了Transformer的16个输入token。token数量就是序列长度embedding维度就是每个token的向量长度。从原理层面说Patchify本质上是在做两件事一是降维压缩避免序列长度爆炸二是把局部特征编码成语义向量。patch大小直接影响性能patch越小token数量越多信息保留越完整但计算成本也跟着涨。DiT通常使用8或16的patch大小。选8时32×32输入会得到16个token选16时就只剩4个token虽然快很多但细节信息丢失得也快。实际做图像生成时这一步还会考虑操作是否在潜空间进行。2.2 隐藏在条件注入里的关键设计纯Transformer只知道输入的token序列并不知道当前是扩散过程的哪一步也不知道用户想要生成哪个类别的图片。所以必须额外把这两类信息喂进去。最常见的时间步处理方式是正弦位置编码把时间步t映射成一个高维向量再经过MLP变换得到时间嵌入。类别标签则通过嵌入层变成向量。如果训练时使用无分类器引导类别嵌入还需要在部分样本上置零模拟“无条件”的情况。DiT在条件注入上有一次重要探索。社区早期尝试过把条件token拼到输入序列里让Transformer在自注意力内部自己学会参考也试过在注意力前面加交叉注意力层类似语言模型里把文本作为上下文。这些方法都有效但各自增加了一部分额外计算。DiT最终推荐的方式是adaLN即自适应层归一化。2.3 为什么adaLN-zero是主角传统的层归一化是对特征做减均值、除方差再通过可学习的缩放参数γ和偏移参数β进行仿射变换。adaLN的思想是不再用固定的γ和β而是让条件信息经过一个MLP后动态生成这些参数。每个Transformer块都会接收一组自己的γ和β也就是说每个块都能根据当前的时间步和类别信息自主决定应该如何缩放和偏移特征。adaLN-zero则更进了一步。它把生成γ、β之外还额外生成一个门控参数α然后在残差连接的末端加上一层“门控缩放”。最关键的是这个MLP最后一个线性层在初始化时故意被置为全零输出那么初始状态下α0残差块的输出完全就是恒等映射。这个设计非常巧妙网络初期从一个稳定的“什么都不做”状态开始然后逐步学习复杂的变换。这比初始就从随机状态硬学要稳定得多收敛速度和最终效果都有明显提升。用代码表达adaLN调制部分大致长这样class AdaDNZero(nn.Module): def __init__(self, hidden_size, cond_size): super().__init__() self.net nn.Sequential( nn.Linear(cond_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size) ) # 最后一个线性层初始化为0 nn.init.zeros_(self.net[-1].weight) nn.init.zeros_(self.net[-1].bias) def forward(self, c): return self.net(c) # 输出按顺序分成6组参数2.4 四种条件注入方案对比DiT论文里系统对比过四种方案我把它们拉开对比一下方案基本做法优点劣势In-context conditioning把条件token拼在图像token前面一起进自注意力实现简单不改内部结构序列变长注意力计算变多Cross-attention在标准Transformer块里加一层交叉注意力条件注入灵活可参考文本序列结构更复杂参数更多adaLN将条件映射为归一化参数计算开销低控制直观表达能力受限于仿射变换adaLN-zero在adaLN基础上加归零门控训练初期稳定指标最好实现稍复杂但收益明显从实验数据看在同等算力和训练步数下adaLN-zero的FID指标普遍优于前几种方案。它直接成了实现DiT的首选方式。很多后来出现的生成模型也都沿用了这种“条件信息通过MLP映射到各层仿射参数”的设计范式。2.5 解码把token还原成图像Transformer输出的是序列化的特征向量但下游任务需要的是像素级别的噪声预测。最后需要一个解码模块把它重塑回空间维度。做法很简单在最后一个Transformer块之后先通过层归一化再过一层线性层把每个位置的embedding映射成patch_size × patch_size × 输出通道数。这里的输出通道数通常是输入图像通道数的两倍因为官方实现采用了v-prediction预测方式每个像素要输出一个预测噪声分量和一个预测图像分量。最后把token序列重新排成网格得到完整的预测图。解码部分代码框架class FinalLayer(nn.Module): def __init__(self, hidden_size, patch_size, out_channels, cond_size): super().__init__() self.norm_final nn.LayerNorm(hidden_size) self.linear nn.Linear(hidden_size, patch_size * patch_size * out_channels, biasTrue) self.adaLN_modulation nn.Sequential( nn.SiLU(), nn.Linear(cond_size, 2 * hidden_size) ) def forward(self, x, c): scale, shift self.adaLN_modulation(c).chunk(2, dim-1) x self.norm_final(x) * (1 scale) shift x self.linear(x) return x这里直接用线性层把特征投影回像素空间避免了传统解码器里复杂的反卷积堆叠。Transformer的最后一层输出是全局语义特征线性投影就是一个“字典查找”的过程把语义映射回局部图案。实际操作中如果觉得画面不够精细也可以额外加一对反卷积层做上采样尤其是当patch尺寸偏大时。3. 核心代码实现从零搭建一个可运行的DiT下面进入实战部分。我会用PyTorch从零写一个足够自己训练小规模数据的DiT实现。代码目标是让结构清晰、容易改参数而不是追求极致性能优化。3.1 环境与基础配置需要PyTorch 2.0以上配合CUDA环境。建议安装timm库用于部分组件。pip install torch timm einops代码中主要用到的常量是模型宽度的缩放系数。DiT有S、B、L、XL等配置它们的区别主要在hidden_size、深度层数和注意力头数。我们以类似DiT-B的配置为例hidden_size为384Transformer层数为12注意力头数为6。3.2 基础组件时间嵌入、标签嵌入、PatchEmbed时间步嵌入沿用正弦位置编码方式再加一个MLP将它映射到隐藏空间这个设计已经被无数实验验证过没什么好改的。import torch import torch.nn as nn import math # 时间步嵌入 class TimestepEmbedder(nn.Module): def __init__(self, hidden_size, frequency_embedding_size256): super().__init__() self.mlp nn.Sequential( nn.Linear(frequency_embedding_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size) ) self.frequency_embedding_size frequency_embedding_size staticmethod def timestep_embedding(t, dim, max_period10000): half dim // 2 freqs torch.exp( -math.log(max_period) * torch.arange(half, dtypetorch.float32) / half ).to(t.device) args t[:, None] * freqs[None] return torch.cat([torch.cos(args), torch.sin(args)], dim-1) def forward(self, t, dtypetorch.float32): t_freq self.timestep_embedding(t, self.frequency_embedding_size) return self.mlp(t_freq.to(dtype))类别标签嵌入处理起来会麻烦一点。训练时如果你用classifier-free guidance需要有一定概率把标签置成“空”类别也就是让模型在训练时见过无条件的样本。这里我把标签的dropout概率设为0.1。class LabelEmbedder(nn.Module): def __init__(self, num_classes, hidden_size, drop_prob0.1): super().__init__() self.embedding nn.Embedding(num_classes, hidden_size) self.num_classes num_classes self.drop_prob drop_prob def forward(self, labels, trainFalse): if train and self.drop_prob 0: drop_mask torch.rand(labels.shape[0], devicelabels.device) self.drop_prob labels torch.where(drop_mask, torch.full_like(labels, 0, dtypetorch.long), labels) return self.embedding(labels)PatchEmbed稍微提两句我把输入通道数设为可配置默认按灰度图或单通道噪声图来实际用彩色图时改成3即可。class PatchEmbed(nn.Module): def __init__(self, input_size32, patch_size8, in_channels3, hidden_size384): super().__init__() self.patch_size patch_size self.num_patches (input_size // patch_size) ** 2 self.proj nn.Conv2d(in_channels, hidden_size, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: [B, C, H, W] x self.proj(x) # [B, hidden, H/P, W/P] x x.flatten(2).transpose(1, 2) # [B, num_patches, hidden] return x3.3 注意力机制与DiTBlockDiT用的注意力是标准的多头自注意力。为了避免画面出现奇怪的伪影DiT在注意力层里还会做QK归一化也就是对query和key先做归一化再算点积这样可以稳定训练。我在这里加上这一细节。class Attention(nn.Module): def __init__(self, hidden_size, num_heads, qk_normTrue): super().__init__() self.num_heads num_heads self.head_dim hidden_size // num_heads self.qkv nn.Linear(hidden_size, 3 * hidden_size, biasTrue) self.proj nn.Linear(hidden_size, hidden_size, biasTrue) if qk_norm: self.q_norm nn.LayerNorm(self.head_dim, elementwise_affineTrue) self.k_norm nn.LayerNorm(self.head_dim, elementwise_affineTrue) else: self.q_norm nn.Identity() self.k_norm nn.Identity() def forward(self, x): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] q self.q_norm(q) k self.k_norm(k) attn (q k.transpose(-2, -1)) * (self.head_dim ** -0.5) attn attn.softmax(dim-1) x (attn v).transpose(1, 2).reshape(B, N, C) return self.proj(x)核心的DiTBlock结构是先做adaLN-zero调制然后过注意力块再过MLP最后门控缩放叠加到残差上。class DiTBlock(nn.Module): def __init__(self, hidden_size, num_heads, cond_size, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6) self.attn Attention(hidden_size, num_heads) self.norm2 nn.LayerNorm(hidden_size, elementwise_affineFalse, eps1e-6) self.mlp nn.Sequential( nn.Linear(hidden_size, int(hidden_size * mlp_ratio)), nn.GELU(approximatetanh), nn.Linear(int(hidden_size * mlp_ratio), hidden_size) ) self.adaLN_modulation nn.Sequential( nn.SiLU(), nn.Linear(cond_size, 6 * hidden_size, biasTrue) ) # 零初始化最后一个线形层 nn.init.zeros_(self.adaLN_modulation[-1].weight) nn.init.zeros_(self.adaLN_modulation[-1].bias) def forward(self, x, c): scale1, shift1, gate1, scale2, shift2, gate2 self.adaLN_modulation(c).chunk(6, dim-1) h self.norm1(x) * (1 scale1) shift1 h self.attn(h) x x gate1.unsqueeze(1) * h h self.norm2(x) * (1 scale2) shift2 h self.mlp(h) x x gate2.unsqueeze(1) * h return x我特别想把第六个参数gate2单独提出来说它专门控制MLP分支的权重在初始化全零的情况下每个块一开始都相当于一个“直通”这种设计让深层网络训练时前期风险降到了最低。3.4 整体模型组合有了上面的组件组装DiT模型就很简单了。forward里做的事是patch化、加位置编码、过一系列DiTBlock最后过FinalLayer把token映射回像素。class DiT(nn.Module): def __init__(self, input_size32, patch_size8, in_channels3, hidden_size384, depth12, num_heads6, num_classes10, learn_sigmaTrue): super().__init__() self.in_channels in_channels self.out_channels in_channels * 2 if learn_sigma else in_channels self.patch_size patch_size self.num_heads num_heads self.x_embedder PatchEmbed(input_size, patch_size, in_channels, hidden_size) self.t_embedder TimestepEmbedder(hidden_size) self.y_embedder LabelEmbedder(num_classes, hidden_size) num_patches self.x_embedder.num_patches self.pos_embed nn.Parameter(torch.zeros(1, num_patches, hidden_size)) nn.init.trunc_normal_(self.pos_embed, std0.02) self.blocks nn.ModuleList([ DiTBlock(hidden_size, num_heads, hidden_size) for _ in range(depth) ]) self.final_layer FinalLayer(hidden_size, patch_size, self.out_channels, hidden_size) def forward(self, x, t, y): x self.x_embedder(x) x x self.pos_embed t self.t_embedder(t) y self.y_embedder(y) c t y for block in self.blocks: x block(x, c) x self.final_layer(x, c) x x.reshape(x.shape[0], self.out_channels, -1, self.patch_size, self.patch_size) # 从token序列重排回图像网格 B x.shape[0] C_out self.out_channels P self.patch_size H W int(x.shape[2] ** 0.5) x x.reshape(B, C_out, H, P, W, P) x x.permute(0, 1, 2, 4, 3, 5).reshape(B, C_out, H * P, W * P) return x这里值得注意的是我把位置编码作为可学习参数初始化。相比固定正弦位置编码可学习编码在小尺寸输入上往往表现更好因为它可以直接被梯度优化到最适合当前patch布局的位置分布。3.5 训练与采样的核心逻辑训练循环本身并不复杂关键是构造带噪图像。给定一批干净图片x_0和随机时间步t先生成噪声ε再通过一个预计算好的噪声系数把二者混合得到x_t。模型的目标是尽可能准确地预测出这个ε。def training_step(model, x_0, y, t, noise, betas, alphas_cumprod): sqrt_alpha torch.sqrt(alphas_cumprod[t]).view(-1, 1, 1, 1) sqrt_one_minus_alpha torch.sqrt(1 - alphas_cumprod[t]).view(-1, 1, 1, 1) x_t sqrt_alpha * x_0 sqrt_one_minus_alpha * noise pred model(x_t, t, y) # 如果模型输出两个通道取前一半作为噪声预测 if pred.shape[1] 2 * model.in_channels: pred pred[:, :model.in_channels] loss torch.nn.functional.mse_loss(pred, noise) return loss采样则使用DDPM的逻辑先从一个纯噪声开始然后从tT逐步往t0走每一步都让模型预测噪声把当前样本稍微去噪一点再补一点随机扰动。torch.no_grad() def sample(model, shape, num_steps, alphas_cumprod, betas, labelsNone, guidance_scale0.0, devicecuda): x torch.randn(shape, devicedevice) for i in range(num_steps - 1, -1, -1): t torch.full((shape[0],), i, devicedevice, dtypetorch.long) alpha alphas_cumprod[i] alpha_next alphas_cumprod[i - 1] if i 0 else torch.tensor(1.0) pred_noise model(x, t, labels) # 无条件 条件融合classifier-free guidance if guidance_scale 0: null_labels torch.full_like(labels, 0) pred_uncond model(x, t, null_labels) pred_noise pred_uncond guidance_scale * (pred_noise - pred_uncond) sqrt_alpha alpha.sqrt() sqrt_one_minus_alpha (1 - alpha).sqrt() x0_pred (x - sqrt_one_minus_alpha * pred_noise) / sqrt_alpha # 加入随机采样噪声 noise torch.randn_like(x) if i 0 else torch.zeros_like(x) x alpha_next.sqrt() * x0_pred (1 - alpha_next).sqrt() * noise return x看到这里你可能有一个感觉整体代码量不大模块很干净。这正是DiT的吸引力——骨干网络的复杂度最高的是注意力但注意力本身就是标准层其他的都是可插拔的调制组件。只要你把PatchEmbed和FinalLayer的尺寸算清楚在其他图像生成任务里复用这套代码的成本非常低。4. 训练技巧、常见问题与排查实录这篇代码我实际跑过的次数不算少训练和采样阶段踩过的坑也值得单独整理出来。这些问题网上零星有一些讨论但系统性整理成速查表的很少。4.1 模型一直在学loss就是不降如果你发现自己训的DiT loss很久都不降先别急着改模型结构按顺序排查这几件事第一检查时间步t是否真的均匀采样。扩散模型的训练目标是在所有时间步上等概率地随机采样如果你用了一个偏向低噪声的采样分布模型会完全忽略高噪声区间的学习loss曲线会一直拖着一条“尾巴”不往下走。第二查看你的模型输入输出通道是否匹配。DiT如果采用v-prediction输出通道数是输入的两倍很多人在封装模型时忘了这个细节导致loss计算时张量一直在错误地广播。第三学习率不是越小越好。DiT一般用AdamW优化器基础学习率设置在1e-4到2e-4之间配合线性warmup和余弦退火。如果设置成1e-5模型前期根本学不动。还有一个容易被忽略的点就是位置编码的标准差。DiT把位置编码初始化为均值为0、标准差0.02的正态分布这个值来自经验。如果初始化标准差太大位置信息一开始就是紊乱的注意力很难看到合理的局部空间结构。我的建议是固定用trunc_normal并且把标准差控制在0.02这个量级。4.2 生成图像颜色发灰、边缘模糊怎么办颜色问题首先想到classifier-free guidance的引导强度。引导强度太小生成结果中类别特征不鲜明画面偏灰偏平滑引导强度太大又会出现过饱和和像素伪影。实践下来在类别引导任务上2到4是一个比较舒服的范围但具体还得看你的数据集和类别embedding的落差。模糊问题则要关注采样步数。DDPM采样如果只用10步基本就是自然模糊建议先用50步以上来排除采样步数的影响。如果你的训练集本身是小尺寸图像生成大图时很容易出现模糊这是因为patch大小没有随分辨率调整。原本用32×32训练的模型如果直接拿来生成64×64token数变了位置编码也可能不匹配。正确的做法是保持patch size一致让序列长度自然翻倍但这样需要重新训练或者做位置编码插值。4.3 显存爆掉的常见原因DiT虽然结构统一但显存压力一点也不小。第一需要注意序列长度patch越小token越多随着图片分辨率提升attention占用的显存是平方级增长的。32×32的图像patch8时有16个token放到256×256patch8时变成1024个token这比多数NLP任务还要紧张。常用的优化手段有混合精度训练把模型计算切成FP16或BF16DiT对精度不敏感实测BF16效果很好。梯度检查点attention和MLP的前向输出不保留反向时重新计算能省不少显存代价是训练时间变长。减少注意力头数只要通道维度被头数整除就能正常跑如果显存吃紧可以把头数从8降到4效果损失相对可控。还有一个经验值得单独说在采样阶段使用fp32推理。混合精度训练没问题但推理时如果直接沿用fp16attention的softmax在极端分数下容易溢出产生画面上的黑斑或白斑。我自己的习惯是训练用BF16采样单独走全精度。4.4 代码常见报错速查表报错现象可能原因解决思路forward里向量维度形状错误PatchEmbed的patch size设置与图像尺寸不整除调整patch size或对图像做paddingloss出现NaN学习率过大降低学习率或检查是否混入NaN标签采样结果全是同一张图时间步embedding或位置编码错乱或模型退化了检查Layernorm是否正常、注意力QK是否做归一化样本出现条纹伪影attention没有dtype对齐fp16溢出采样时切换全精度或对QK结果做缩放保护模型参数量与论文不符hidden_size和depth设置偏差参数核对DiT的标准配置表4.5 一个从零开始的训练经验清单我把自己实际操作中最顺手的配置列一下使用ImageNet的类别条件作为示例任务输入分辨率32×32patch8hidden_size384depth12head6batch size设为128。训练使用AdamW初始学习率1e-4weight decay设为0.3。整个训练过程大概在几千步后就能看到比较清晰的类别分化。为什么weight decay要设这么大DiT是纯Transformer参数众多且高度依赖归一化。过小的weight decay会让attention的残差输出逐渐累积导致深层网络训练不稳定。0.3到0.5这个区间是我测试下来最稳的范围尤其适合训练步数较长的扩散模型。另一个容易影响情绪的点是训练曲线不会像分类任务那样快速收敛。扩散模型要分阶段前期loss缓慢下降中期loss稳步下滑但指标提升很快。如果一定要用训练曲线来守护判断我建议不要只看loss而是每隔一定步数真正采样一批图片出来看眼见的画面比曲线数据可靠得多。5. 最后想分享的一点工程心得写完这整套流程最深的体会可能和大多数人预想的不太一样。DiT最费时间的地方并不在理解Transformer原理也不在写模型主体代码而在于把条件注入、patch尺寸、通道数、采样参数这些容易互相牵连的细节对齐。我见过不只一个人把模型搭起来、训练loss正常下降但最后采样生成的图片怎么看都像“打码”——最后发现是patch重排回图像时维度的顺序错了。所以如果你也是照着这篇内容自己写一遍强烈建议在训练前先用一张真实图片过一遍前向流程把输出尺寸和预期值打印出来逐一比对。扩散模型这几年迭代速度很快DiT本身也带了很多后续工作。我自己的习惯是学一个新模型时必须动手写一遍核心代码写不完理解始终是虚的。这篇文章给出的实现足够在小规模任务上跑通全过程如果你接着想往上加一些新想法比如更长的序列注意力、多模态条件注入、更高效的采样器都可以在这套结构上无缝地换层。希望这篇东西能帮你省下一些翻论文和调bug的时间。