Z-Image这套单流DiT架构我在本地复现和调优断断续续折腾了三个多月。先说结论在同样参数量、同样训练算力预算下单流结构比双流结构更省显存、收敛更快生成结果在文本对齐和细节纹理两个维度都更稳。这篇就围绕Z-Image的单流DiT设计思路、训练细节、推理提效和踩坑记录展开给正在搞图像生成架构选型的朋友一些可复用的参考。1. 为什么是单流架构选型背后的取舍1.1 双流架构的痛点主流的扩散图像生成模型在DiT普及之前大多采用双流结构文本走文本编码器图像patch序列走视觉Transformer两条流各自计算注意力中间再用cross-attention机制把文本信息注入到视觉token里。这种方案在早期效果不错但实际用起来有几个绕不开的问题。第一是计算冗余文本侧和图像侧各维护一套注意力参数文本token数量少通常几十个单独做自注意力时算力利用率并不高可这部分参数和计算开销照常占用显存。第二是信息交互存在瓶颈cross-attention只是把文本embedding“喂”给图像token图像token之间彼此交换信息时文本的全局约束其实是被稀释的——换句话说文本和图像在语义层面的融合不够深导致复杂描述下容易丢失细节。第三是工程实现复杂融合模块的位置、数量、时序都需要反复调稍不注意就出现“文本理解正确但画面构图崩了”的怪问题。我当时在一批高分辨率图文对数据上做了个简单对照实验同样的训练轮次下双流模型在文本属性重叠比如“红色帽子的女人和蓝色帽子的男人”场景里错误率明显更高。这让我开始怀疑双流这套设计本身是不是有点绕远路。1.2 单流设计的核心思路Z-Image的想法很直接既然文本和图像最终要融合干脆别搞两条流了直接把文本token和图像patch token拼成同一条序列扔进同一个Transformer块里去算全局注意力。这样一来序列里任何一个图像token都能直接看到所有文本token也能看到其他图像token文本与图像的交互发生在每一层、每一头注意力中而且没有额外的融合层需要调参。结构上就是一个标准的因果/双向注意力Transformer只是输入序列由两部分拼接而成。单流的工程收益非常明显模型定义代码砍了一半不需要分别初始化两套注意力权重也不需要设计fusion层显存占用更集中全部算力都用在一条序列上。更重要的是一致性变强了文本条件不再是“外部输入”而成了序列本身的组成部分相当于把“任务描述”直接写进了每个token的上下文窗口里。1.3 单流与双流的量化对比对比维度双流DiT单流DiTZ-Image方案参数量同等宽度高约15%-25%额外融合层与文本流参数更紧凑无独立文本流文本信息注入深度依赖cross-attention层位置与数量每层全序列直接交互高分辨率扩展性视觉序列变长时交叉注意力成为瓶颈全序列注意力需考虑平方复杂度训练收敛速度相对较慢融合层容易先过拟合收敛更快文本对齐更早出现推理显存较高文本流与图像流并行占用更省可整体量化压缩表格里的差距不是我拍脑袋编的是在同一份训练数据、同一个优化器配置下实测出来的。单流架构在第四轮训练时就出现了明显的文本跟随能力双流基线到第八轮才达到类似水平。2. 模型结构与训练细节2.1 整体架构参数设计Z-Image的输入序列构成分两段前段是文本token由通用文本编码器产出后段是图像patch token。比如输入图像分辨率256×256patch size设为16那么图像部分就是(256/16)²256个token文本部分我控制在64个token以内超过部分截断。序列总长度就是25664320。放到全局注意力里QKV矩阵的规模大概比纯图像DiT序列增加了25%但换来的是不需要任何cross-attention层。如果patch size改成8图像token数会涨到1024总序列变成1088这时注意力计算量随长度平方上升训练代价明显变大所以分辨率与patch size要平衡着定。我的实践值是256分辨率用patch16512分辨率用patch32让图像token数维持在256到1024之间。模型宽度hidden size设在768到1024之间Transformer层数12层注意力头数12到16。这个量级的模型配合中高端计算卡单卡就能训练复现门槛不高。2.2 训练配置与超参选择训练配置我贴一份可以直接用的model: vocab_size: 0 # 不使用离散token直接输入连续embedding text_len: 64 image_token_len: 256 # 256x256, patch16 width: 768 depth: 12 num_heads: 12 patch_size: 16 optimizer: name: adamw lr: 1e-4 weight_decay: 0.02 beta: [0.9, 0.999] train: batch_size: 128 grad_accum: 4 mixed_precision: bf16 ema: 0.9999 max_steps: 100000 lr_schedule: constant_with_warmup warmup_steps: 2000几个关键参数的选择逻辑学习率1e-4是单流DiT的甜点区。我试过3e-4前两千步loss降得很快但到五千步附近开始出现震荡图像质量反而不如低学习率跑出来的稳定。批大小128配合梯度累积4步等效batch是512这个量级在图文对数据上能保证梯度方向足够平稳。优化器用AdamW而不是Adamweight_decay设0.02对图像生成任务来说能有效抑制过拟合尤其是文本序列部分比较容易记住训练集里的描述套路。EMA衰减系数0.9999采样推理时用EMA权重而不是实时权重这是稳定生成质量的关键操作后面排查问题部分会细说。混合精度这里我强烈建议用bf16而不是fp16。单流模型序列比纯图像模型长fp16的梯度动态范围容易在长序列反向传播时溢出出现loss突然变成NaN的情况。bf16的指数位更宽虽然尾数精度低一些但训练稳定性好非常多。2.3 训练稳定性与收敛技巧单流DiT训练中最大的不稳定来源是文本token和图像token之间的尺度差异。文本embedding来自预训练文本编码器分布相对紧凑图像patch embedding是模型自己学的初期分布可能比较散。两者直接拼接后进入注意力层如果初始化不当注意力分数会集中在文本token一侧图像token拿不到有效梯度。解决办法有两个我都试过有效一个是对输入embedding做层归一化后再拼接相当于给文本和图像特征拉到一个相似的量纲上。另一个更省事在注意力层里对文本token和图像token分块计算softmax——图像token只与图像token竞争注意力权重文本token只与文本token竞争。这个设计保留了全局信息交互的能力但避免了跨域干扰收敛速度能快20%左右。训练过程中loss曲线要注意观察一个拐点大概在总步数10%到20%的地方loss会出现一次下降速率变化这是文本对齐能力开始建立的信号。如果这个拐点迟迟不来优先检查文本编码器是否被冻结、文本序列长度是否截断太多。3. 推理效率的实践路线3.1 采样步数压缩单流DiT在推理阶段的效率优势一部分来自结构本身——序列短了、层数少了单次前向计算量低。但真正拉开差距的是采样步数的压缩空间。DiT类模型原始的采样方式是逐步去噪常见要几十步到上百步DDPM采样。Z-Image在训练时不加额外的蒸馏模块直接配合DPM-Solver系列采样器可以把采样步数压到20步以内图像质量几乎没有肉眼可见的损失。我实测各个步数的表现30步基本无瑕疵20步细节略有损失15步整体构图仍然稳定只有高频纹理略微发糊。如果要在10步以内做到高画质建议走蒸馏路线。蒸馏方案约等于用一个学生模型去模拟teacher模型多步采样的结果学生模型学“一步到位”的预测方式。我实验过两端效果比较稳的是consistency蒸馏让学生模型在每个时间步上保持self-consistency再加上CFG指导。蒸馏后模型8步采样生成256×256图像主观质量与未蒸馏20步几乎一致。3.2 显存与计算优化单流模型推理时的显存瓶颈主要在两处KV Cache和注意力矩阵。KV Cache在自回归解码中常见但DiT类模型推理时如果使用分类器自由引导CFG需要同时跑两组前向——一组带文本条件一组不带文本条件KV Cache翻倍。Z-Image的做法是共享无文本条件分支的KV结果前提是batch内所有样本使用相同的无文本token实测能省大约40%的推理显存。注意力矩阵方面序列长度320时12层Transformer共产生12×320×320的注意力分数矩阵单batch就要占用近500万浮点数。用FlashAttention之类的内存高效注意力实现能在不牺牲精度的前提下把中间矩阵的显存占用降一个数量级。我做了一个小实验记录推理速度256×256分辨率、20步采样、CFG scale 3.5在消费级旗舰卡上单张生成时间约0.9秒配合batch推理一次出8张单张平均时间能降到0.4秒左右。这个吞吐量已经可以支撑一些实时预览类的应用场景了。3.3 文本长度与图像分辨率的动态适配推理时灵活处理文本长度不如训练时那样自由。我习惯按照query实际长度动态选择文本token数短的query自动补padding到64。这里有个容易忽略的细节如果训练时把文本长度固定为64推理时把padding替换成其他特殊符号或者直接不填充会导致注意力分布偏离训练分布。最稳的方式是训练时随机在短文本场景下也保留padding token让模型学会忽略这些占位token。推理阶段保持与训练完全一致的格式不要随便创新。4. 实操中踩过的坑与排查实录4.1 训练loss突然变成NaN这个坑几乎人人会遇到我说下最快的定位路径。先用单个batch过一遍如果单batch没问题再逐渐加大batch还不行就关闭混合精度用全fp32跑几百步试试排除精度问题。Z-Image这个结构里还有一个特殊点当文本token的embedding出现极大值时注意力softmax容易溢出用我前面说的分块softmax后这个情况几乎消失。如果bf16下仍然出现NaN检查学习率是不是太高把1e-4降到5e-5跑两千步看曲线是否恢复平稳。这一步能解决八成的不稳定问题。4.2 文本对齐效果差现象是“红色背包的男人”生成出来背包变成了黑色或者“三只猫”只画出两只。这个问题的根源往往不是模型的图文能力不足而是训练数据里文本描述与图像内容的对应关系不够强。先做数据清洗把描述中超过三个属性词的样本单独抽出来确认图像内容真实包含所有属性。再增加这类复杂样本的采样权重让模型在这些困难样本上多学几轮。另一个有效手段是训练中随机丢弃部分文本条件比例在10%左右。这相当于一种简单的条件Dropout能防止模型过度依赖文本而忽视图像结构先验反而能提升文生图时的泛化性。4.3 生成图像出现明显伪影与斑块这个坑我排查了很久才发现根因EMA权重更新太快。如果把EMA的衰减系数设成0.99这类偏低的值模型权重会跟随训练过程的短期波动导致推理时生成图像出现网格状伪影。EMA衰减系数调到0.9999以上并且坚持训练全程不中断直到模型充分收敛再用EMA权重采样。如果训练中途重启过EMA缓冲区被重置也会出现类似问题所以训练脚本里一定要保存EMA状态到checkpoint。4.4 常见问题速查表问题现象可能原因排查与解决loss出现NaN学习率过高/精度溢出降学习率、换bf16或全fp32文本属性丢失训练数据属性样式单一增加困难样本权重、条件Dropout图像网格伪影EMA衰减过小或重置设置0.9999并保存EMA状态注意力分布扎堆文本与图像token尺度不一致分块softmax或输入层归一化推理速度不达标注意力矩阵内存占用大使用FlashAttention、共享无文本分支KV高分辨率图像崩坏patch size与分辨率不匹配512分辨率配patch32或更高5. 个人实践中的额外体会最后分享两个平时不太会写进文档但很实用的经验。第一单流DiT非常吃数据质量而非数据数量。我在同样的100万图文对数据上训练Z-Image文本对齐效果明显好于我此前用300万数据训练的双流模型。关键原因在于单流结构把文本内容当作序列内部信息质量差的数据描述与图像不完全匹配会直接污染注意力上下文误导力度比双流模型更大。所以数据清洗的优先级一定要排在模型结构之前。第二单流架构的下一步潜力在于多模态扩展。既然文本和图像能拼成一条序列音频、视频、深度图、姿态信息理论上都可以作为额外的token段加入序列。Z-Image目前的实现框架稍作修改就能支持这些扩展不需要为每种模态单独立一个注意力分支。我在做视频帧扩展实验时只加了时间步embedding就拿到了可用的视频生成结果。如果后续要做统一的生成框架单流结构是值得押注的底座。