
多模态大模型走到今天架构选型这件事已经不像两年前那样Transformer一把梭了。我最近在复现一个图文生成的小项目时重新把扩散模型这条线捋了一遍发现很多刚入门的朋友卡在一个很具体的问题上做多模态生成Backbone到底该用UNet还是Transformer这两个词在各种论文和开源项目里反复出现但真正动手从零搭一个Stable Diffusion下称SD的时候你会发现选型不是拍脑袋而是由你要解决的任务、算力预算、以及你想不想复用现成权重共同决定的。这篇就围绕这个核心问题展开把SD的骨架拆开给你看顺带把从零搭建一个最小可用SD的完整路径走一遍。不管你是刚接触扩散模型的新手还是已经跑过几个Demo但没深究内部结构的老手应该都能从里面拿到点能直接抄的东西。1. 先把Backbone这个词放回它该在的位置1.1 Backbone在生成模型里到底指什么Backbone这个词在不同语境下含义差别很大。在分类任务里它通常指ResNet、ViT这类特征提取主干但在扩散模型里Backbone指的是承载去噪过程的主网络结构它接收带噪的latent或像素和时间步信息输出对噪声的预测。换句话说它是整个生成流程里真正干活的那部分参数量的大头都在它身上。很多人第一次看SD的架构图会懵因为图里既有VAE又有UNet还有CLIP文本编码器。这三者分工完全不同VAE负责像素空间和latent空间之间的压缩与还原CLIP负责把文本变成条件向量而UNet才是那个在latent空间里一步步去噪的Backbone。所以当有人问SD的Backbone是什么准确答案是UNet而不是Transformer——尽管UNet内部大量使用了Transformer的注意力机制。这个区分很关键后面讲选型的时候会反复用到。我见过不少初学者把SD里的Cross-Attention模块误认为是用了Transformer做Backbone其实那只是UNet中间层插入的注意力块主干依然是卷积下采样加跳连的结构。搞清楚这一点你才不会在选型时被名词带偏。1.2 为什么这个问题现在值得单独拎出来说放在2022年这个问题几乎不用讨论图像生成领域UNet是绝对主流。但到了多模态大模型时代情况变了。DiTDiffusion Transformer这一类工作证明了纯Transformer结构在扩散模型里同样能打甚至在scaling上表现更好Sora这类视频生成模型也把Transformer架构推到了台前。于是就有了一个很现实的困惑我从零搭SD到底该照着经典UNet抄还是直接上Transformer这个选择的背后其实是三笔账算力账、复用账、可控账。算力账决定你能不能训得动复用账决定你能不能站在别人肩膀上可控账决定你调试时能不能快速定位问题。下面几节我会把这三笔账一笔笔算给你看而不是给你一个看情况的万能答案。2. UNet作为SD Backbone的结构拆解2.1 编码器-解码器加跳连UNet的骨架逻辑UNet之所以叫U-Net是因为它的结构画出来像个U左边一路下采样把特征图尺寸压小、通道数拉大右边一路上采样把尺寸还原中间用跳连skip connection把左右两侧同尺度的特征拼起来。这个设计最初是为医学图像分割做的但放到扩散模型里意外地合适。原因在于去噪这个任务本身需要同时理解全局语义和局部细节。下采样路径负责抓全局——这张图整体是什么构图、什么风格上采样路径负责还原细节——边缘、纹理、小结构。跳连则保证在下采样过程中丢失的高频信息能直接传回解码端不至于全靠网络自己回忆。你可以把它理解成编码器是看远解码器是画细跳连是别把看远时记下的细节弄丢。在SD里这个UNet还额外插了两类条件注入时间步embedding通过加法或FiLM方式注入每个残差块文本条件则通过Cross-Attention注入中间层。所以严格说SD的UNet是一个条件UNet不是原始那版纯图像输入的UNet。从零搭建时这两条条件通路是必须自己接的也是最容易出错的地方。2.2 时间步与文本条件是怎么塞进去的时间步的处理相对直接把标量t做正弦位置编码过两层MLP得到embedding然后在每个ResBlock里通过scale-shift的方式调制特征。这个做法借鉴了FiLM的思想好处是几乎不增加多少参数就能让网络感知现在是第几步去噪。文本条件复杂一些。SD用的是Cross-Attention把UNet中间层特征作为Query把CLIP输出的文本token序列作为Key和Value算注意力。这样每个空间位置都能查询文本里哪些词和自己相关。这里有个实操细节——Cross-Attention一般只插在UNet的中低分辨率层高分辨率层不插原因是高分辨率特征图空间位置太多做注意力的计算量会爆炸而且高分辨率层更需要的是局部细节而非语义对齐。提示自己搭的时候Cross-Attention的插入位置直接决定显存占用。我最初图省事在所有层都插了结果batch size只能开到1后来按SD官方配置只在中低分辨率层插显存立刻松了一大截。2.3 从零写一个最小UNet的实操路径真要动手我建议先别追求和SD完全一致先写一个能跑通的最小版本。核心模块就三个ResBlock含时间步调制、Downsample/Upsample、Cross-Attention。下面是一个极简ResBlock的骨架用PyTorch写import torch import torch.nn as nn class ResBlock(nn.Module): def __init__(self, in_ch, out_ch, temb_dim): super().__init__() self.norm1 nn.GroupNorm(8, in_ch) self.conv1 nn.Conv2d(in_ch, out_ch, 3, padding1) self.temb_proj nn.Linear(temb_dim, out_ch) self.norm2 nn.GroupNorm(8, out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, padding1) self.skip nn.Conv2d(in_ch, out_ch, 1) if in_ch ! out_ch else nn.Identity() def forward(self, x, temb): h self.conv1(torch.nn.functional.silu(self.norm1(x))) h h self.temb_proj(torch.nn.functional.silu(temb))[:, :, None, None] h self.conv2(torch.nn.functional.silu(self.norm2(h))) return h self.skip(x)这个块里时间步是通过广播加法注入的比FiLM简化了一些但足够跑通。搭完整UNet时按通道数[320, 640, 1280, 1280]分四级每级堆2到3个ResBlock级间下采样解码端对称上采样并拼接跳连特征。这套配置基本对齐SD 1.x的量级参数量在8亿左右单卡24G显存能训小分辨率。我踩过的一个坑是GroupNorm的组数。一开始随手设了32组结果通道数320时每组才10个通道归一化效果很差训练loss抖得厉害。后来改成8组稳定多了。这个细节文档里很少提但实际影响不小。3. Transformer当BackboneDiT路线的取舍3.1 DiT把UNet的归纳偏置换成了什么DiTDiffusion Transformer的核心思路是既然Transformer在别的领域都能scale扩散模型为什么不行它把UNet的卷积下采样换成了Patchify——把latent切成一个个patch展平成token序列然后直接喂给标准Transformer块。跳连没了U形结构没了取而代之的是纯注意力的全局建模。这个改动带来的最大变化是归纳偏置的转移。卷积天然带局部性和平移等变性这是UNet的先天优势Transformer没有这些假设一切靠数据学。好处是模型容量上限更高scaling曲线更漂亮坏处是同等参数量下Transformer需要更多数据才能学到卷积免费给你的那些先验。我在一个小数据集上对比过两者UNet收敛明显更快前几千步就能出像样的轮廓而DiT要到更后面才开始开窍。这不是说DiT差而是说它的优势要在数据和算力都充足时才显现。你要是只有几张卡、几万张图硬上DiT大概率是给自己找罪受。3.2 什么情况下该认真考虑Transformer有三种情况我会建议你认真评估Transformer路线。第一种是你要做的是视频或多帧生成时序建模用注意力比用卷积自然得多这也是很多视频生成模型选Transformer的现实原因。第二种是你打算做超大规模训练参数量往几十亿走的时候Transformer的scaling行为更可预测工程上也更容易做并行切分。第三种是你想复用现成的大语言模型权重把文本侧和生成侧统一到一个Transformer范式里架构上更干净。反过来如果你的目标是快速出一个能用的图像生成Demo、算力有限、还想直接加载社区现成的SD权重那UNet几乎是唯一理性选择。这不是保守是算账算出来的。3.3 一个常被忽略的中间路线其实还有第三条路UNet为主干但在中间层用Transformer块替换部分卷积块。SD本身就已经这么干了——它的UNet中间层插了Transformer块做空间自注意力和文本交叉注意力。所以UNet还是Transformer这个二选一本身有点伪命题真实工程里更多是混合。我个人的经验是从零搭的时候先按纯UNet跑通等baseline稳了再逐步把中间层换成Transformer块做消融。这样每一步改动的影响都可控出问题也知道是哪一步引入的。一上来就搞混合架构调试成本会高很多。4. 从零搭一个SD完整链路与关键决策点4.1 三个组件各司其职别混为一谈搭SD不是只搭Backbone完整链路包含三块VAE、文本编码器、去噪Backbone。VAE把512×512×3的图像压成64×64×4的latent压缩比48倍这一步决定了后续所有计算都在低维空间进行是SD能在家用卡上跑的根本原因。文本编码器把prompt变成77×768的条件序列。Backbone在latent空间里做去噪。从零搭的时候我的建议是VAE和文本编码器直接加载预训练权重只自己搭Backbone。原因很实在VAE的重建质量直接决定生成图的清晰度上限自己训一个够用的VAE成本极高文本编码器更是如此CLIP那种图文对齐能力不是小项目能训出来的。把精力集中在Backbone上才是性价比最高的做法。4.2 训练目标与噪声调度容易被跳过的地基扩散模型的训练目标说穿了很简单给图像加噪让网络预测加进去的噪声算MSE。但噪声调度noise schedule的选择很讲究。SD用的是scaled linear beta schedulebeta从0.00085线性增到0.012而且是对beta的平方根做缩放。这个细节直接影响模型在不同噪声水平上的表现。我最初照搬了DDPM原始的linear schedule结果生成图在高噪声阶段糊成一团。后来换成SD的scaled linear同样的网络结构效果肉眼可见地好了。这说明调度策略不是可有可无的超参而是和Backbone同等重要的设计决策。从零搭的时候这一块千万别图省事随便填。训练循环的核心逻辑大概是这样# 伪代码突出关键步骤 x0 vae.encode(images) # 图像 - latent t torch.randint(0, T, (bs,)) # 随机时间步 noise torch.randn_like(x0) xt sqrt_alphas_cumprod[t] * x0 sqrt_one_minus[t] * noise pred unet(xt, t, text_emb) # Backbone预测噪声 loss mse_loss(pred, noise)看着简单但每一步都有坑。比如latent要乘一个scaling factorSD里是0.18215这个系数是VAE latent的统计标准差倒数不乘的话训练会不稳定。这种细节论文里一句话带过实操时不知道就会卡很久。4.3 采样阶段Backbone之外的另一半工程训练完只是第一步采样器sampler的选择同样影响最终效果。DDPM采样要上千步慢得没法用DDIM能压到几十步到了DPM-Solver这类高阶求解器20步左右就能出不错的结果。这些采样器本质上都是在用Backbone的噪声预测做数值积分Backbone不变换采样器就能显著改变速度和质量。我的建议是训练阶段用最简单的DDPM目标采样阶段直接上DPM-Solver或多步DDIM。这样训练逻辑干净采样效率也够。自己实现采样器的时候注意classifier-free guidance的实现——它需要对条件和无条件两次前向做插值guidance scale设太高会过饱和设太低条件不起作用7.5左右是个常见起点。5. 实操中真正会卡住你的几个地方5.1 显存与精度的现实约束从零搭SD第一个撞上的墙几乎都是显存。8亿参数的UNetfp32下光权重就3G多加上梯度和优化器状态24G卡训512分辨率很勉强。几个实用的缓解手段用混合精度fp16或bf16训练权重和激活都能省一半用gradient checkpointing用计算换显存batch size压到1到2靠梯度累积凑等效batch。我实测下来bf16比fp16更稳不容易出NaN前提是你的卡支持。另外VAE的latent在训练时最好detach别让梯度回传到VAE否则显存和计算都会翻倍而且没必要——VAE是冻结的。5.2 条件注入的维度对齐问题文本条件注入是最容易出维度错误的地方。CLIP输出的序列长度是77维度768SD 1.x而UNet中间层特征通道可能是320或640。Cross-Attention里Query的维度要和Key/Value的维度通过投影矩阵对齐这个投影层必须自己加不能想当然直接算。我第一次写的时候忘了加投影报了个维度不匹配的错查了半天才发现是这里。还有一个隐蔽的坑文本编码器输出的padding token。CLIP的77长度里实际prompt可能只占十几个token剩下都是padding。Cross-Attention如果不做maskpadding token也会参与注意力计算虽然影响不大但会让结果略微偏移。SD官方实现里是处理了这个的自己写的时候留意一下。5.3 训练不收敛时的排查顺序训练loss不降或者生成图是噪声排查要有顺序别乱试。我的习惯是先确认数据管线——latent是不是正确编码的、scaling factor乘了没再确认噪声调度——加噪公式的系数有没有算错然后看条件注入——把文本条件置空看模型能不能至少学会无条件生成最后才怀疑网络结构。这个顺序的逻辑是从简单到复杂、从数据到模型。数据问题最常见也最好查网络结构问题最少见但最难改。我见过有人一上来就怀疑UNet设计有问题结果折腾一周发现是latent没乘scaling factor。这种坑按顺序排查十分钟就能定位。6. 选型结论没有最优架构只有最合适的账本回到最初那个问题——Backbone用UNet还是Transformer。我的实际体会是这个选择不该脱离你的具体约束来谈。算力紧、要快速出结果、想复用社区权重UNet是理性答案做视频、做超大规模、要统一多模态范式Transformer值得投入。而真实工程里混合架构往往是更务实的落点SD本身就是最好的例子。从零搭一个SD的过程本质上是在理解每个设计决策背后的权衡为什么VAE要冻结、为什么Cross-Attention只插中低层、为什么噪声调度要用scaled linear。这些为什么搞清楚了你换任何Backbone都能快速上手因为变的只是承载去噪的那个壳不变的是扩散模型这套加噪去噪的底层逻辑。我自己搭完一遍之后最大的收获不是得到了一个能出图的模型而是再看到任何新的生成架构都能一眼看出它在哪一层做了什么样的取舍。