多模态大模型这两年火得一塌糊涂但真到动手搭的时候第一个卡住大多数人的问题往往不是数据、不是训练策略而是Backbone到底选什么。我前后折腾过好几个图像生成和图文理解的项目从早期拿Unet硬怼到后来全面转向Transformer架构中间踩的坑能写满一个笔记本。这篇就围绕“从0搭建一个SD”这件事把Backbone选型的来龙去脉、Unet和Transformer各自的脾气、以及一个能跑起来的最小实现路径掰开揉碎讲清楚。不管你是刚接触扩散模型的新手还是已经调过几个模型但没深究过结构的老手应该都能从里面抠出点有用的东西。1. 先搞清楚Backbone在SD里到底扮演什么角色很多人一上来就纠结“Unet还是Transformer”但连Backbone在整条链路里干什么都没想明白选型就成了拍脑袋。我先把这层窗户纸捅破。1.1 扩散模型的三段式结构里Backbone是那个“干重活”的一个标准的Stable Diffusion类模型拆开来看其实是三块文本编码器负责把提示词变成语义向量**Backbone去噪网络**负责在潜空间里一步步把噪声还原成图像特征**解码器VAE Decoder**负责把潜空间特征还原成像素图。文本编码器通常用CLIP或者T5这类现成的解码器也是固定的VAE真正决定模型能力上限、参数量大头、训练成本大头的就是中间那个Backbone。打个比方文本编码器像是“翻译官”把你的话翻译成模型能懂的语言VAE解码器像是“画师的手”负责最后落笔而Backbone是“画师的脑子”它要在几十步里去判断“当前这团噪声里哪些该去掉、哪些该保留、往哪个方向走”。这个判断的质量直接决定了出图的构图、细节、语义一致性。所以Backbone选型不是审美问题是能力问题。1.2 为什么会有Unet和Transformer两条路线之争早期扩散模型DDPM那一批几乎清一色用Unet因为Unet天生适合做“图像到图像”的稠密预测跳跃连接能把浅层细节直接送到深层去噪这种任务简直是量身定做。SD 1.5、SD 2.1都是Unet backbone稳定、成熟、生态好。但后来DiTDiffusion Transformer出来了Sora那一波又把Transformer路线推到了台前。Transformer的优势在于可扩展性——参数量往上堆的时候性能提升曲线比Unet更平滑而且天然适合处理多模态输入文本token和图像token可以拼在一起做注意力。这就是为什么新一代模型越来越多往Transformer靠。提示选型不是“新的就一定好”。Unet在中小规模、单模态图像生成上依然能打Transformer在超大规模、多模态联合建模上更有优势。你的项目规模决定了你该站哪边。1.3 从0搭SD先明确你要的是“能跑”还是“能扩”我见过太多人一上来就想复现一个Sora级别的架构结果卡在显存和训练成本上动弹不得。从0搭建我的建议是先分清目标如果你是想理解原理、跑通一个能出图的最小闭环Unet路线成熟度高、调试成本低两三天能出第一张图如果你是想做多模态联合训练、为后续扩展留空间那从一开始就上Transformer架构哪怕先做小规模也比后期重构省事。下面这张表是我自己总结的两条路线对比选型时可以直接对照维度Unet BackboneTransformer BackboneDiT类参数量规模通常0.8B~3B可轻松扩展到数十B训练稳定性高收敛行为可预测需要更精细的学习率调度显存占用相对友好序列长度一长就吃显存多模态扩展需要额外设计融合模块天然支持token拼接生态成熟度极高参考实现多中等仍在快速演进适合场景单模态图像生成、快速验证多模态联合、大规模扩展2. Unet路线成熟稳重但别指望它无限扩展既然要从0搭我先把Unet这条路走一遍因为它是理解扩散模型去噪机制最好的入口。你把Unet啃透了再去看DiT会发现很多设计思想是一脉相承的。2.1 Unet的核心结构拆解编码、瓶颈、解码三段Unet之所以叫U是因为它的结构像个字母U左边一路下采样编码器中间一个瓶颈层右边一路上采样解码器左右之间用跳跃连接skip connection连起来。在SD里这个Unet还额外接入了时间步嵌入timestep embedding和文本条件cross-attention所以它其实是个“条件Unet”。编码器负责逐层提取特征、降低分辨率把空间信息压缩成语义信息瓶颈层在最低分辨率上做全局推理解码器再逐层上采样恢复分辨率。跳跃连接的作用是把编码器浅层的高频细节直接送到解码器对应层避免上采样过程中细节丢失。这个设计对去噪任务特别关键——因为去噪既要理解全局语义这团噪声整体该往哪个方向走又要保留局部细节边缘、纹理不能糊。我实际搭的时候编码器和解码器各用4个stage每个stage包含若干个ResBlock加注意力层通道数从128起步逐级翻倍到512瓶颈层再翻一次。这个配置在单卡24G显存上能跑batch size 4左右属于能验证效果的入门配置。2.2 时间步和文本条件是怎么注入的这是Unet做条件生成的关键也是新手最容易写错的地方。时间步注入用的是正弦位置编码加MLP把标量时间步t映射成一个高维向量然后通过全连接层投影到每个ResBlock的通道维度以加法方式注入。文本条件注入用的是cross-attention在Unet的中间层和上采样层里插入cross-attention模块query来自图像特征key和value来自文本编码器的输出。这里有个实操细节cross-attention不要加得太早。我一开始在每个stage都塞了cross-attention结果训练时文本条件“压过”了图像自身的去噪信号出图语义对了但构图崩了。后来改成只在瓶颈层和上采样阶段加效果立刻正常。这个经验在官方文档里不会写但实测很关键。# 时间步嵌入的简化实现 import torch import torch.nn as nn import math class TimestepEmbedding(nn.Module): def __init__(self, dim, freq_dim256): super().__init__() self.freq_dim freq_dim self.mlp nn.Sequential( nn.Linear(freq_dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim) ) def forward(self, t): # t: [batch] half self.freq_dim // 2 freqs torch.exp( -math.log(10000) * 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)2.3 从0搭Unet的实操步骤与参数选择我把搭建流程拆成五步按顺序走基本不会乱确定潜空间分辨率。SD用的是VAE把512x512图像压到64x64x4的潜空间所以Unet的输入是64x64、4通道。这个分辨率决定了你的下采样次数——64到32到16到84个stage刚好再多就压没了。定通道数基数。入门用128显存紧张用64追求效果用192。通道数基数乘以2的stage次方就是各层通道。堆ResBlock和注意力。每个stage至少2个ResBlock注意力层在低分辨率stage加即可高分辨率stage加注意力显存吃不消。接条件注入。时间步嵌入全局注入文本条件用cross-attention选择性注入。写输出层。最后一层用卷积把通道数映射回4潜空间通道数输出预测的噪声。参数选择上有个经验公式总参数量控制在0.8B到1.2B之间这个区间在单卡或小集群上训练成本可控效果也够看。超过这个量级Unet的扩展收益开始递减不如换Transformer。注意Unet的下采样用卷积stride2还是用池化我建议用卷积因为可学习。池化会丢信息且不可逆在去噪任务里是减分项。2.4 Unet路线的坑我踩过的三个典型问题第一个坑是跳跃连接的特征对齐。编码器和解码器对应层的通道数、分辨率必须严格一致否则拼接会报错。我一开始通道数设计成不对称调试了半天才发现是拼接维度对不上。第二个坑是注意力层的显存爆炸。在64x64分辨率上做全局注意力序列长度4096注意力矩阵是4096x4096单这一层就能吃掉好几个G。解决办法是在高分辨率用窗口注意力或者干脆不加只在16x16、8x8这些低分辨率层加全局注意力。第三个坑是训练初期loss震荡。Unet对学习率比较敏感我实测用1e-4起步、配合cosine退火比较稳用1e-3直接发散。另外EMA指数移动平均对Unet特别重要不加EMA出图质量明显差一截。3. Transformer路线为扩展而生但要伺候好它Unet跑通之后我把同样的任务用DiT架构重写了一遍感受完全是另一个世界——灵活是真灵活难调也是真难调。3.1 DiT的核心思想把图像切成patch当token处理DiTDiffusion Transformer的思路很直接既然Transformer擅长处理序列那就把潜空间特征图切成一个个patch每个patch展平成一个token加上位置编码然后丢进标准的Transformer block里。文本条件通过cross-attention或者adaLN自适应层归一化注入。和Unet最大的区别在于DiT没有下采样和上采样全程保持同一分辨率。这意味着它的计算量随分辨率是平方级增长的所以patch size的选择很关键。SD潜空间是64x64patch size取2的话就是32x321024个token这个长度在可接受范围取1就是4096个token显存直接翻几倍。我实测下来patch size2是性价比最高的选择1024个token的序列长度配合12到24层Transformer block参数量能控制在0.6B到1.5B之间单卡能训。3.2 adaLN条件注入比cross-attention更优雅的方案DiT论文里提了一个很聪明的设计用adaLN-Zero注入条件。具体做法是把时间步嵌入和类别/文本嵌入相加通过MLP预测出每个Transformer block的scale和shift参数对LayerNorm的输出做仿射变换。最后一层还额外预测一个gate参数初始化为0这样训练初期整个block相当于恒等映射训练非常稳定。这个设计比cross-attention优雅在哪cross-attention要额外引入一组注意力参数计算量不小adaLN只是对归一化层做调制参数少、计算省而且初始化成0保证了训练稳定性。我在项目里两种都试过adaLN收敛更快、显存更省唯一缺点是文本条件的细粒度控制稍弱——如果你的任务需要精确的局部文本对齐cross-attention还是更合适。# adaLN-Zero 的简化实现 class DiTBlock(nn.Module): def __init__(self, dim, num_heads): super().__init__() self.norm1 nn.LayerNorm(dim, elementwise_affineFalse) self.attn nn.MultiheadAttention(dim, num_heads, batch_firstTrue) self.norm2 nn.LayerNorm(dim, elementwise_affineFalse) self.mlp nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim) ) # 6个调制参数attn的shift/scale/gate mlp的shift/scale/gate self.adaLN_modulation nn.Sequential( nn.SiLU(), nn.Linear(dim, 6 * dim) ) # 零初始化gate保证训练初期恒等 nn.init.zeros_(self.adaLN_modulation[-1].weight) nn.init.zeros_(self.adaLN_modulation[-1].bias) def forward(self, x, c): shift1, scale1, gate1, shift2, scale2, gate2 \ self.adaLN_modulation(c).chunk(6, dim-1) h self.norm1(x) * (1 scale1.unsqueeze(1)) shift1.unsqueeze(1) h, _ self.attn(h, h, h) x x gate1.unsqueeze(1) * h h self.norm2(x) * (1 scale2.unsqueeze(1)) shift2.unsqueeze(1) h self.mlp(h) x x gate2.unsqueeze(1) * h return x3.3 从0搭DiT的关键参数与训练配置DiT的搭建比Unet简单——没有复杂的上下采样结构就是patchify、堆block、unpatchify三步。但训练配置比Unet讲究得多学习率DiT对学习率更敏感我用1e-4配合warmup 5000步比Unet的warmup要长。权重衰减0.01到0.05之间太高会欠拟合。梯度裁剪必须开阈值1.0Transformer的梯度比Unet更容易爆。混合精度bf16比fp16稳fp16在DiT上容易出现梯度下溢。patchify的实现要注意用卷积实现比用reshapelinear更高效kernel size和stride都等于patch size输出通道就是hidden dim。unpatchify用转置卷积或者pixel shuffle都行我用的转置卷积实测更稳。3.4 Transformer路线的坑显存、收敛、位置编码第一个坑还是显存。1024个token、24层、hidden dim 1152的DiT单卡24G跑batch size 2都紧张。解决办法是梯度检查点gradient checkpointing用时间换显存能省40%左右代价是训练慢20%。第二个坑是收敛慢。DiT前期loss下降比Unet慢我一度以为是自己写错了后来对比论文才发现这是正常现象——Transformer需要更多步数才能“找到感觉”。耐心训到2万步左右loss曲线才开始和Unet持平。第三个坑是位置编码。DiT用的是固定的2D正弦位置编码我试过换成可学习的效果反而差一点因为可学习位置编码在分辨率变化时泛化性差。如果你的项目需要多分辨率训练固定编码更稳妥。4. 两条路线的融合与选型决策别非黑即白搭完两套之后我最大的感受是Unet和Transformer不是对立关系实际项目里经常是混着用的。4.1 混合架构Unet的骨架加Transformer的注意力现在不少工作走的是混合路线——保留Unet的上下采样结构但把里面的卷积注意力换成Transformer block。这样既有Unet的多尺度归纳偏置又有Transformer的长程建模能力。我试过在Unet的瓶颈层和最低分辨率stage用Transformer block替换出图的全局一致性有提升显存增加在可接受范围。这种混合架构的实操要点是只在低分辨率层换Transformer高分辨率层保持卷积否则显存扛不住。另外Transformer block的初始化要小心建议用零初始化输出投影保证训练初期不破坏原有Unet的行为。4.2 选型决策表按你的实际约束来选我把选型时该问自己的问题整理成一张表照着填基本不会选错你的约束推荐路线理由单卡24G以内快速出效果Unet成熟、省显存、调试快需要多模态联合训练Transformertoken拼接天然支持参数量要上到3B以上Transformer扩展曲线更优团队没有Transformer调参经验Unet收敛行为更可预测要做视频/长序列生成Transformer序列建模能力强推理延迟要求极低Unet结构轻、推理快4.3 从0搭SD的完整最小闭环不管选哪条路线一个能跑的最小闭环都包含这几块VAE可以直接用现成的、文本编码器CLIP现成的、Backbone自己搭、训练循环噪声采样、前向、loss、反向、采样器DDIM或DPM-Solver。我的建议是先把VAE和文本编码器用现成的冻结住只训Backbone这样能把问题聚焦。训练循环里有个容易忽略的点噪声调度。linear schedule和cosine schedule出图风格不一样cosine在低噪声段更细腻linear在高噪声段更激进。我默认用cosine配合v-prediction实测比epsilon-prediction在细节上更好。提示第一次跑通不要追求出图质量先确认loss能稳定下降、采样能出有结构的图再逐步调参。我第一版跑出来的图全是噪点排查后发现是噪声采样和设备没对齐这种低级错误新手很容易犯。5. 实操排查那些文档里不会写的坑搭模型的过程里报错和异常是常态。我把高频问题和排查思路整理出来省得你重复踩。5.1 训练不收敛的排查顺序遇到loss不降按这个顺序查先看数据管道有没有问题噪声和图像是否对齐、归一化范围对不对再看学习率和warmup然后看条件注入有没有写错时间步嵌入维度、文本mask最后看初始化。我遇到过一次loss完全不降查了两天才发现是时间步嵌入的freq_dim和MLP输入维度对不上静默广播成了错误形状这种bug最坑。5.2 出图质量差的典型原因对照表现象可能原因排查方向图全糊训练步数不够或学习率太低加步数、调学习率语义对但构图乱文本条件注入过强减少cross-attention层数颜色异常VAE解码问题或归一化错误检查VAE输入范围局部重复纹理注意力坍缩加dropout或调温度高噪声步出图崩噪声调度不匹配换cosine schedule5.3 显存优化的几个实用手段显存不够是常态我常用的手段按性价比排序梯度检查点省40%慢20%、bf16混合精度省30%左右、把优化器状态换成分片存储、减小batch size配合梯度累积。最后这个最容易被忽略——batch size 1配合梯度累积8步效果和batch size 8接近显存却省了8倍。6. 我个人在Backbone选型上的几点体会折腾完这几轮我最大的体会是Backbone选型本质是在“归纳偏置”和“可扩展性”之间做权衡。Unet的卷积和跳跃连接是强归纳偏置它“知道”图像是局部相关的、多尺度的所以小数据小算力下表现好Transformer的归纳偏置弱它靠数据和规模去学这些规律所以规模上去了才显出优势。如果你的项目是验证想法、快速迭代别犹豫Unet起步两三天能出结果。如果你明确要做多模态、要往大规模走那从一开始就上Transformer哪怕先做小规模架构上的债早还比晚还省事。混合架构是个不错的折中但别为了混合而混合先想清楚你要解决的是全局一致性问题还是局部细节问题。最后分享一个我调试时的小习惯每次改Backbone结构先用一个极小的数据集比如几十张图过拟合一遍。如果连几十张图都过拟合不了说明结构或训练配置有硬伤别急着上大数据集浪费时间。这个“过拟合小数据集”的土办法帮我省下的调试时间比任何工具都多。