文章目录一、为什么图像生成不再是“写作文”而是“洗照片”1. 核心比喻DiT 像一群修复师在擦一幅被雪花盖住的海报二、先建立全局观DiT 到底由哪几部分组成1. VAE先把高清大图压成“浓缩草稿”2. Forward Diffusion故意往草稿上泼噪声3. Patchify把 latent 切成一个个 patch token4. Position Embedding告诉模型“每块在画面的哪里”5. Timestep / Condition Embedding告诉模型“现在是第几轮去噪”6. DiT Blocks让所有 patch 进行全局群聊三、DiT 最核心的比喻不是“逐字续写”而是“整图会诊”1. 图像里的 Q、K、V 到底在表达什么四、数学上它到底在预测什么不是类别而是“噪声”五、为什么 DiT 不用 causal mask这点和 GPT 完全不同1. GPT 为什么必须 mask2. DiT 为什么不需要六、DiT 最关键的条件注入adaLN 和 adaLN-Zero 到底是什么1. 先回忆普通 LayerNorm 是什么2. adaLN把“老师指令”变成每层的调音旋钮3. adaLN-Zero为什么最后还要加个 Zero4. 直观比喻新员工先旁听再逐步接手工作七、DiT Block 长什么样和标准 Transformer Block 有什么不同八、DiT 和 U-Net 到底有什么结构差异1. U-Net像金字塔式的卷积施工队2. DiT像一群平级专家在全图会议室里反复会诊3. 粗暴对比九、训练和采样流程一口气串起来看1. 训练时故意弄脏再逼模型学会清洗2. 采样时从纯噪声开始反复显影十、极简 PyTorch手搓一个 Mini DiT 骨架观察这段代码时你重点盯 4 件事十三、下一步该学什么本文通过直观的比喻和极简的 PyTorch 代码彻底拆解 DiTDiffusion Transformer的核心——它到底是如何把“满屏雪花噪声”一步步变成清晰图像的。一、为什么图像生成不再是“写作文”而是“洗照片”如果你已经理解了 Transformer 做文本生成那么你脑子里很容易先冒出一个想法既然文本可以一个 token 一个 token 地往后写那图像是不是也能一个像素一个像素地往后写理论上可以但图像和文本有一个巨大的不同文本天然是一条序列词有明显的前后顺序。图像天然是一个二维平面左上角、右下角、远处背景和近处主体是同时存在、彼此耦合的。如果你硬要像写作文一样一个像素一个像素地生成图片就会非常别扭生成左眼时还没看到右眼脸容易歪。生成前景人物时还没全局理解背景光影容易乱。图像的“整体一致性”很难保证。所以扩散模型Diffusion Model走了另一条完全不同的路线它不从左到右“写”图而是从一张全是噪声的脏图开始一轮一轮把噪声擦掉让图像逐渐“显影”出来。1. 核心比喻DiT 像一群修复师在擦一幅被雪花盖住的海报想象你面前有一张海报但海报表面被泼满了白色雪花噪点几乎看不清内容。这时来了一群修复师他们不是每次只修一个像素而是先把海报切成很多小方块每个修复师负责看一块所有修复师先开会互相交流全局情况然后每人决定“我这一块这一轮该擦掉多少噪声”擦掉一点后再开下一轮会一轮轮下来原本满是雪花的海报就逐渐显出轮廓、颜色、细节。这就是 DiT 的本质DiT 不是直接“画图”而是在“反复开会协同去噪”。二、先建立全局观DiT 到底由哪几部分组成如果粗暴地用一句话概括 DiTDiT 在 latent 空间里把图像切成 patch token再用 Transformer 反复预测噪声。它大致可以拆成 6 个零件1. VAE先把高清大图压成“浓缩草稿”真实图片太大了。比如一张256 x 256 x 3的 RGB 图直接让 Transformer 对它做全局注意力成本很高。所以 DiT 通常不会直接处理原始像素而是先用一个VAE Encoder把图片压缩到 latent 空间原图256 x 256 x 3latent32 x 32 x 4这就像先把高清海报压成一张“浓缩草稿纸”分辨率小了很多但主要信息还保留着后面只需要在这张草稿上做去噪2. Forward Diffusion故意往草稿上泼噪声训练时模型不是直接看干净图而是故意把 latent 草稿弄脏。设干净 latent 是z zz随机采样一个时间步t tt往里面加上高斯噪声得到z t z_tzt​。你可以把这一步理解成老师先把一张草稿图泼脏再拿给学生看“你来猜猜我刚才泼了多少脏水。”3. Patchify把 latent 切成一个个 patch token到了 Transformer 这里不能直接吃二维特征图它更擅长处理 token 序列。所以要把 latent 切成小块 patch。例如latent 大小32 x 32 x 4patch size2 x 2那就会得到patch 数量(32/2) * (32/2) 256每个 patch 的原始维度2 * 2 * 4 16然后再通过一个线性层把每个 patch 投影成统一维度的 token embedding。于是一张图就变成了一串图像 token。这一步非常像 NLP一句话 - 一串词 token一张 latent 图 - 一串 patch token4. Position Embedding告诉模型“每块在画面的哪里”如果没有位置编码Transformer 只知道自己拿到了 256 个 patch但不知道谁在左上角谁在中心谁在右下角这就像你把拼图块全倒在桌子上但没告诉模型哪块原来属于天空、哪块属于脸。所以 DiT 也必须给 patch token 加上位置编码。区别只是文本位置编码强调一维顺序图像位置编码强调二维空间位置5. Timestep / Condition Embedding告诉模型“现在是第几轮去噪”这一点是 DiT 和普通 ViT 最大的不同之一。DiT 每次输入的不是固定图像而是某个噪声阶段的图z t z_tzt​。所以模型必须知道现在是早期噪声很重还是后期图已经很清晰这张图要求生成“猫”还是“宇航员”因此 DiT 会把这些条件编码成向量时间步 embedding告诉模型现在在第几轮类别/文本 embedding告诉模型想生成什么你可以把它理解成每轮开会前老师先发指令“现在噪声很大先抓轮廓。”“现在接近收尾重点修眼睛、纹理和光影。”“这张图的目标是金毛犬不是猫。”6. DiT Blocks让所有 patch 进行全局群聊进入主体后就是一层层 Transformer Block。每个 patch 都会生成自己的Q我需要什么信息K我能提供什么线索V我真正携带的内容然后所有 patch 一起做 Self-Attention全局互相交流。这意味着左上角的一块天空可以直接参考右边的太阳脸部 patch 可以直接和肩膀、头发、背景光线交流前景和背景可以在同一层里完成全局协调这就是 Transformer 放到图像生成里的最大魅力全图所有区域都能直接“开群聊”不是只能看局部邻居。三、DiT 最核心的比喻不是“逐字续写”而是“整图会诊”文本 Transformer 更像一群词在讨论下一句该写哪个词。DiT 更像一群 patch 在讨论这一轮每块应该擦掉多少噪声。1. 图像里的 Q、K、V 到底在表达什么假设画面里有一只猫背景是沙发。某个 patch 位于猫眼附近它的 Q 可能在“想”“我现在需要确认自己到底属于眼睛、脸毛还是背景阴影”沙发区域的 patch 的 K 可能在“展示”“我是大块的平坦纹理像背景不像五官。”猫耳朵附近的 patch 的 K 可能在“展示”“我属于猫头的边缘结构和眼睛区域关系密切。”于是猫眼 patch 在做注意力时会重点关注邻近的脸部 patch耳朵 patch光照相关的背景 patch而不会把大量注意力浪费在无关区域。最终这个 patch 就能更清楚地判断“我这里应该保留猫眼结构不该继续像噪声一样模糊。”所以 Self-Attention 在 DiT 里干的事情不是语言里的“上下文消歧”而是在全局画面中确定每个局部到底应该长成什么样。四、数学上它到底在预测什么不是类别而是“噪声”这是初学 DiT 最容易混淆的地方。分类网络最后预测的是这是一只猫这是一辆车GPT 最后预测的是下一个 token 是谁但 DiT 最后预测的通常是当前图里加进去的噪声ϵ \epsilonϵ或者等价形式v vv它不是直接输出“最终图像”而是输出“这张带噪图里哪些部分像噪声应该减掉多少。”于是训练目标通常是一个简单的均方误差L ∥ ϵ ^ − ϵ ∥ 2 \mathcal{L} \| \hat{\epsilon} - \epsilon \|^2L∥ϵ^−ϵ∥2其中ϵ \epsilonϵ真实加进去的噪声ϵ ^ \hat{\epsilon}ϵ^DiT 预测出来的噪声这就像老师真的往图上泼了一桶脏水而模型要尽量精确地回答“你刚才泼的是这一桶脏水对吧”如果它猜得越来越准反过来就越来越会“洗照片”。五、为什么 DiT 不用 causal mask这点和 GPT 完全不同你学过文本 Transformer 后脑子里很容易有个惯性Transformer 不是都要 mask 吗答案是DiT 不需要 GPT 那种 causal mask。1. GPT 为什么必须 mask因为 GPT 是做自回归生成预测第 5 个词时不能偷看第 6 个词否则训练时等于作弊所以它必须使用下三角 mask把未来位置遮住。2. DiT 为什么不需要因为 DiT 的任务不是“预测未来 token”而是给定一整张带噪图同时判断所有位置的噪声。也就是说在 DiT 的每一步里所有 patch 都已经同时存在大家可以彼此看见没有“未来 token 不能看”的约束这就像全体修复师围着同一张海报开会左上角的修复师当然可以看右下角眼睛区域当然可以参考嘴巴和耳朵背景当然可以参考主体的轮廓所以 DiT 使用的是普通双向自注意力不是自回归的单向 mask 注意力。一句话记忆GPT 的时间轴在“词序”上所以怕偷看未来。DiT 的时间轴在“去噪轮数”上所以同一轮里所有 patch 都能全局互看。六、DiT 最关键的条件注入adaLN 和 adaLN-Zero 到底是什么这部分是 DiT 真正有“灵魂细节”的地方。因为仅仅把时间步 embedding 加到输入里还不够强。模型更希望在每一层都知道当前噪声阶段是多少当前条件是什么于是 DiT 引入了Adaptive LayerNormadaLN以及更常见的adaLN-Zero。1. 先回忆普通 LayerNorm 是什么普通 LayerNorm 会对一个 token 的特征做标准化LN ( x ) γ ⋅ x − μ σ β \text{LN}(x) \gamma \cdot \frac{x - \mu}{\sigma} \betaLN(x)γ⋅σx−μ​β其中μ , σ \mu, \sigmaμ,σ由当前 token 自己算出来γ , β \gamma, \betaγ,β是固定可学习参数作用就是把数值拉稳。2. adaLN把“老师指令”变成每层的调音旋钮在 adaLN 中γ \gammaγ和β \betaβ不再是固定参数而是由条件向量动态生成的。条件向量可以来自timestep embeddingclass embeddingtext embedding也就是说不同时间步、不同生成目标会让 LayerNorm 的缩放和平移方式都不同。这就像每层前面多了一个“总指挥调音台”现在是早期噪声阶段 - 把模型调到更关注大轮廓现在是后期精修阶段 - 把模型调到更关注局部细节现在目标是狗 - 强化和狗相关的结构模式3. adaLN-Zero为什么最后还要加个 ZeroDiT 论文里常用的是adaLN-Zero。它的关键思想是让每个块一开始先“几乎什么都不做”训练时再慢慢学会如何介入。具体做法可以粗略理解成条件向量不只生成 shift / scale还生成一个 gate门控系数这个 gate 初始接近 0于是某个 block 的残差分支一开始近似是关闭的x ← x gate ⋅ Sublayer ( adaLN ( x , c ) ) x \leftarrow x \text{gate} \cdot \text{Sublayer}(\text{adaLN}(x, c))x←xgate⋅Sublayer(adaLN(x,c))当gate ≈ 0时这层一开始对主干影响很小训练更稳定。4. 直观比喻新员工先旁听再逐步接手工作普通残差块像一个新员工第一天上班就直接大幅改流程容易把系统搞乱。adaLN-Zero 像是新员工第一天先只旁听不随便动系统随着训练进行模型逐渐学会什么时候该介入、介入多少这能显著提升大模型训练的稳定性。七、DiT Block 长什么样和标准 Transformer Block 有什么不同如果用最简化视角看DiT Block 还是那两个老朋友Self-AttentionMLP但每个子层前面会加上条件调制adaLN并且通常带门控。一个非常典型的简化版流程是对输入 token 做 adaLN送入 Self-Attention用 gate 控制这条残差分支的强度再对主干做一次 adaLN送入 MLP再用 gate 控制 MLP 分支的强度你可以把它理解成标准 Transformer Block 是“大家开会 各自思考”。DiT Block 则是“老师先发这轮指令再开会再思考”。八、DiT 和 U-Net 到底有什么结构差异在 DiT 爆火之前扩散模型最经典的 backbone 是 U-Net。两者都能做扩散去噪但世界观不太一样。1. U-Net像金字塔式的卷积施工队U-Net 的典型特点是下采样不断压缩分辨率提取高层语义上采样逐步恢复分辨率跳跃连接把浅层细节直接传给深层它像一支分层施工队小工先处理局部纹理中层处理结构高层处理全局语义最后再逐级还原回来2. DiT像一群平级专家在全图会议室里反复会诊DiT 的核心则是把图切成 patch直接作为 token 序列处理每层都用全局 self-attention它不像 U-Net 那样强依赖多尺度卷积金字塔而是更像所有 patch 从一开始就处在同一个大会场里每一层都能做全局交流。3. 粗暴对比U-Net 强在卷积归纳偏置强天生适合图像局部结构DiT 强在 Transformer 扩展性好模型做大后全局建模能力强U-Net 像“图像工程老将”DiT 像“Transformer 化的新一代主力”所以很多人会把它们类比成U-Net扩散时代的卷积王者DiT扩散时代的 Transformer 主力九、训练和采样流程一口气串起来看1. 训练时故意弄脏再逼模型学会清洗训练流程可以概括为输入真实图片x xx用 VAE 编码成 latentz zz随机采样时间步t tt加噪得到z t z_tzt​patchify 位置编码加入时间步和类别/文本条件送入 DiT预测噪声ϵ ^ \hat{\epsilon}ϵ^和真实噪声ϵ \epsilonϵ做 MSE loss2. 采样时从纯噪声开始反复显影真正生成图像时没有真实图片只有随机噪声从纯高斯噪声 latent 开始把当前z t z_tzt​送入 DiTDiT 预测这一轮的噪声调度器scheduler据此把图变干净一点重复很多轮最后得到干净 latent用 VAE Decoder 解码成最终图片一句话总结训练是“老师泼脏水学生学会认脏水”。采样是“学生独立洗照片”。十、极简 PyTorch手搓一个 Mini DiT 骨架下面这段代码不是工业级实现但足够帮助你把 DiT 的核心骨架真正串起来。importtorchimporttorch.nnasnnclassPatchEmbed(nn.Module):def__init__(self,in_channels4,patch_size2,hidden_size256):super().__init__()self.patch_sizepatch_size# 用卷积做 patchify等价于切块后线性投影self.projnn.Conv2d(in_channels,hidden_size,kernel_sizepatch_size,stridepatch_size)defforward(self,x):# x: (B, C, H, W)xself.proj(x)# (B, D, H/P, W/P)xx.flatten(2).transpose(1,2)# (B, N, D)returnxclassTimestepEmbedder(nn.Module):def__init__(self,hidden_size):super().__init__()self.mlpnn.Sequential(nn.Linear(hidden_size,hidden_size),nn.SiLU(),nn.Linear(hidden_size,hidden_size),)defforward(self,t_embed):# 这里假设外部已经把 t 变成 hidden_size 维向量returnself.mlp(t_embed)classAdaLNModulation(nn.Module):def__init__(self,hidden_size):super().__init__()# 这里一次性生成 6 组参数# attn 的 shift/scale/gate mlp 的 shift/scale/gateself.netnn.Sequential(nn.SiLU(),nn.Linear(hidden_size,6*hidden_size))defforward(self,cond):returnself.net(cond).chunk(6,dim-1)classDiTBlock(nn.Module):def__init__(self,hidden_size256,num_heads8,mlp_ratio4.0):super().__init__()self.norm1nn.LayerNorm(hidden_size,elementwise_affineFalse)self.norm2nn.LayerNorm(hidden_size,elementwise_affineFalse)self.attnnn.MultiheadAttention(hidden_size,num_heads,batch_firstTrue)mlp_hiddenint(hidden_size*mlp_ratio)self.mlpnn.Sequential(nn.Linear(hidden_size,mlp_hidden),nn.GELU(),nn.Linear(mlp_hidden,hidden_size))self.adaLNAdaLNModulation(hidden_size)defmodulate(self,x,shift,scale):returnx*(1scale.unsqueeze(1))shift.unsqueeze(1)defforward(self,x,cond):shift_msa,scale_msa,gate_msa,shift_mlp,scale_mlp,gate_mlpself.adaLN(cond)# Attention 子层x_normself.modulate(self.norm1(x),shift_msa,scale_msa)attn_out,_self.attn(x_norm,x_norm,x_norm)xxgate_msa.unsqueeze(1)*attn_out# MLP 子层x_normself.modulate(self.norm2(x),shift_mlp,scale_mlp)mlp_outself.mlp(x_norm)xxgate_mlp.unsqueeze(1)*mlp_outreturnxclassFinalLayer(nn.Module):def__init__(self,hidden_size256,patch_size2,out_channels4):super().__init__()self.normnn.LayerNorm(hidden_size,elementwise_affineFalse)self.linearnn.Linear(hidden_size,patch_size*patch_size*out_channels)self.adaLNnn.Sequential(nn.SiLU(),nn.Linear(hidden_size,2*hidden_size))defforward(self,x,cond):shift,scaleself.adaLN(cond).chunk(2,dim-1)xself.norm(x)xx*(1scale.unsqueeze(1))shift.unsqueeze(1)xself.linear(x)returnxclassMiniDiT(nn.Module):def__init__(self,input_size32,in_channels4,patch_size2,hidden_size256,depth6,num_heads8):super().__init__()self.patch_sizepatch_size self.in_channelsin_channels self.x_embedderPatchEmbed(in_channels,patch_size,hidden_size)num_patches(input_size//patch_size)**2self.pos_embednn.Parameter(torch.zeros(1,num_patches,hidden_size))self.t_embedderTimestepEmbedder(hidden_size)self.blocksnn.ModuleList([DiTBlock(hidden_size,num_heads)for_inrange(depth)])self.final_layerFinalLayer(hidden_size,patch_size,in_channels)defunpatchify(self,x):# x: (B, N, P*P*C)B,N,_x.shape Pself.patch_size Cself.in_channels HWint(N**0.5)xx.view(B,H,W,P,P,C)xx.permute(0,5,1,3,2,4).contiguous()xx.view(B,C,H*P,W*P)returnxdefforward(self,z_t,t_embed):# 1. 图像转 patch tokenxself.x_embedder(z_t)self.pos_embed# 2. 时间步条件向量condself.t_embedder(t_embed)# 3. 一层层 DiT Blockforblockinself.blocks:xblock(x,cond)# 4. 输出每个 patch 的噪声预测再拼回二维 latentxself.final_layer(x,cond)xself.unpatchify(x)returnx观察这段代码时你重点盯 4 件事PatchEmbed它把二维 latent 变成 token 序列。pos_embed它告诉模型每个 patch 在图上的空间位置。t_embedder adaLN它把“第几轮去噪”这个条件注入到每一个 block。unpatchify它把 token 再还原回二维 latent 噪声预测图。十三、下一步该学什么如果你已经看懂了这篇那么下一步最值得继续深挖的是 4 个点扩散公式本身彻底搞懂q ( x t ∣ x 0 ) q(x_t|x_0)q(xt​∣x0​)、反向采样、scheduler、DDPM/DDIM。Classifier-Free Guidance为什么一句提示词能把图“拉向”你想要的方向。Latent Diffusion为什么先进 VAE latent 再扩散会极大降低计算量。Stable Diffusion / Flux / PixArt 这类模型看它们如何把文本编码器、VAE、DiT/U-Net 拼成完整系统。如果你愿意我下一条可以继续按同样风格给你写《为什么扩散模型能从噪声里还原图像》或者《Classifier-Free Guidance 到底在“拉”什么》或者《DiT vs U-Net一篇图像生成骨干网络进化史》