1. 从标题拆解PixelDiT2到底想解决什么问题1.1 像素空间扩散的老毛病与DiT的新瓶颈PixelDiT2这个名字拆开来看就是三个关键词Pixel、DiT、2。Pixel代表它工作在像素空间DiT代表Diffusion Transformer架构2代表这是第二代方案。把这三个词串起来核心命题就清楚了——在像素空间直接做扩散生成并且用Transformer作为主干网络同时引入表示学习来约束生成过程。先说像素空间扩散这件事。早期扩散模型大多在像素空间直接操作后来大家发现直接在像素上做去噪计算量随分辨率平方增长生成一张1024×1024的图注意力矩阵的规模会大到让人头皮发麻。于是有了Latent Diffusion先把图像压到潜空间在低维空间里做扩散最后再解码回像素。这条路走通了Stable Diffusion也成了过去两年的主流。但潜空间扩散有个绕不开的问题自编码器的压缩是有损的。VAE把图像压到潜空间再解回来细节会丢文字会糊高频纹理容易失真。你在潜空间里辛辛苦苦去噪最后解码出来发现边缘发虚、小字糊成一团这种体验做过生成的人都懂。所以一直有人想回到像素空间把这条损失链路砍掉。DiTDiffusion Transformer的出现给了像素空间扩散新的可能性。Transformer的scaling能力比U-Net强参数量堆上去生成质量能持续提升。但DiT原版是在潜空间做的直接搬到像素空间计算量的问题又回来了。PixelDiT2要做的就是在像素空间把DiT跑起来同时用表示学习来兜底让模型在像素层面去噪时不至于丢掉语义结构。1.2 Representation-Grounded这个定语才是真正的题眼标题里最值得琢磨的是Representation-Grounded这个定语。Grounded的意思是有根基的、有依据的Representation-Grounded就是说生成过程要被表示学习的结果所约束和引导。这背后的逻辑是纯像素空间的扩散模型去噪时只看局部像素的统计规律容易陷入局部合理、全局崩坏的困境。比如生成一张人脸眼睛鼻子嘴巴单独看都挺像但拼在一起比例失调。原因就是模型缺少一个全局的、语义层面的表示来约束生成。表示学习Representation Learning在这里扮演的角色相当于给扩散过程装了一个语义导航。模型在去噪的每一步不仅要看当前像素噪声还要参考一个预训练好的表示编码器提取的语义特征。这个特征告诉模型你现在生成的应该是一张猫的图猫的语义结构是这样的从而把生成过程往正确的语义方向上拉。我个人的理解是PixelDiT2本质上是在做一件事把表示学习从事后评估变成过程约束。以前我们训练完生成模型再用一个分类器或CLIP去评估生成质量表示是外挂的。PixelDiT2把表示内嵌到扩散过程中让表示成为生成的一部分。这个思路如果走通对像素空间扩散的实用性是质的提升。1.3 适合谁来读这篇内容这篇内容适合三类人。第一类是正在做扩散模型训练和调优的算法工程师尤其是被潜空间压缩损失困扰、想尝试像素空间方案的。第二类是对DiT架构感兴趣、想了解Transformer在生成任务上最新进展的研究者。第三类是想把生成模型落地到对细节要求高的场景比如文字生成、精细纹理、医学图像的产品和技术负责人。如果你只是想知道怎么用现成的文生图工具这篇内容可能偏硬核。但如果你想搞清楚像素空间扩散到底能不能打、表示学习怎么和扩散结合那接下来的内容应该对你有用。2. 核心架构设计像素DiT与表示约束怎么捏在一起2.1 为什么非要在像素空间硬刚先回答一个很实际的问题潜空间扩散已经够用了为什么还要折腾像素空间答案藏在应用场景里。潜空间扩散的压缩比通常是8倍甚至更高512×512的图压到64×64的潜表示。这个压缩过程对自然图像的大部分区域是友好的但对高频细节是灾难性的。我做过一个测试用潜空间扩散生成包含小字号文字的图片文字区域几乎必然糊掉因为文字的高频信息在压缩时被丢掉了。同样的问题出现在精细纹理比如织物经纬、毛发细节和医学图像比如X光片里的微小病灶上。像素空间扩散没有这个压缩损失理论上能保留全部细节。但代价是计算量。一张512×512的图像素空间有262144个token如果按patch切分patch size为2的话是65536个token注意力矩阵是token数的平方。这个规模直接做全局注意力显存和算力都扛不住。PixelDiT2的解法是分块注意力加表示引导。它不会对全图做全局注意力而是在局部窗口内做注意力同时用表示学习提取的全局语义特征来补偿局部注意力的视野局限。这个设计思路和Swin Transformer有点像但目的不同——Swin是为了降低计算量PixelDiT2是为了在降低计算量的同时不丢全局语义。2.2 表示编码器选型为什么不用CLIP而用DINOv2表示编码器的选择是个关键决策。市面上常见的表示模型有CLIP、DINOv2、MAE等。PixelDiT2这类方案通常会选DINOv2或者类似的自监督表示模型而不是CLIP。原因在于CLIP的表示是图文对齐的它的语义空间是被文本监督过的偏向于这张图整体是什么类别。而DINOv2是纯视觉自监督的它的表示保留了更多的空间结构和局部语义信息。对于像素级生成任务你需要的是这个位置的像素应该属于什么语义区域而不是整张图是什么类别。DINOv2的patch级特征图更适合做这种空间约束。具体来说DINOv2对一张图会输出一个特征图每个patch对应一个特征向量。PixelDiT2在去噪的每一步会把当前噪声图像过一个冻结的DINOv2编码器拿到特征图然后通过一个轻量的投影网络把这个特征注入到DiT的每一层。注入方式通常是cross-attention或者adaLN自适应层归一化。注意表示编码器在训练时通常是冻结的不参与梯度更新。这样做一是省显存二是防止表示空间被生成任务带偏。如果你自己复现千万别手贱去微调编码器很容易把表示空间搞崩。2.3 DiT主干的改造从潜空间到像素空间的适配DiT原版是为潜空间设计的patch size通常是2或4因为潜空间本身已经降采样过了。搬到像素空间patch size的选择变得很关键。patch size太小比如1token数爆炸计算量扛不住。patch size太大比如8每个token覆盖的像素太多细节建模能力下降。PixelDiT2这类方案一般会选patch size为2或4配合多尺度注意力或者层级结构来平衡。另一个改造点是位置编码。像素空间的位置编码需要更精细因为像素级任务对空间位置敏感。常用的做法是可学习的绝对位置编码加上相对位置偏置或者用RoPE旋转位置编码的二维版本。RoPE的好处是外推能力强训练时用256×256推理时用512×512也能work不会因为位置编码没见过的长度而崩掉。还有一个细节是时间步嵌入。扩散模型需要知道当前去噪到哪一步了时间步嵌入的质量直接影响生成效果。PixelDiT2通常会用傅里叶特征加MLP的方式做时间步嵌入然后通过adaLN注入到每一层。这个和DiT原版一致没什么特别的。2.4 表示约束的注入方式对比表示约束怎么注入到DiT里有几种常见方案各有优劣。注入方式实现难度效果计算开销适用场景Cross-Attention中好中表示特征维度与DiT隐藏维度接近时adaLN低中低表示特征需要压缩成全局向量时拼接后自注意力低中高token数不多时自适应门控高好中需要动态调节表示强度时Cross-Attention是最直观的把表示特征作为key和valueDiT的token作为query做一次注意力。这样每个像素token都能查询表示特征找到自己对应的语义。缺点是计算量增加因为多了一次注意力。adaLN是把表示特征压成一个向量然后生成缩放和平移参数作用在DiT的归一化层上。这种方式计算开销小但表示信息被压缩得太狠空间结构丢失严重。PixelDiT2这类方案我推测会用Cross-Attention为主、adaLN为辅的混合方式。Cross-Attention负责注入空间语义adaLN负责注入全局风格。这样既有空间精度又有全局一致性。3. 实操复现从零搭建PixelDiT2的关键步骤3.1 环境准备与依赖安装复现PixelDiT2硬件门槛不低。像素空间扩散对显存的需求比潜空间高一个量级。我的建议是至少准备一张24GB显存的卡比如3090或4090如果要训256×256以上分辨率最好上A100 40GB或80GB。软件环境方面PyTorch 2.0以上是必须的因为要用到scaled_dot_product_attention和torch.compile。CUDA版本建议11.8或12.1配合对应的cuDNN。其他依赖包括pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install timm einops accelerate transformers pip install xformers # 可选但强烈建议装省显存DINOv2的权重可以从timm里直接加载不用自己去下import timm rep_encoder timm.create_model(vit_base_patch14_dinov2.lvd142m, pretrainedTrue, num_classes0) rep_encoder.eval() for p in rep_encoder.parameters(): p.requires_grad False提示DINOv2的输入分辨率是固定的比如518×518如果你的训练分辨率不是这个需要做插值或者resize。插值位置编码的时候注意用bicubic别用nearest否则表示质量会掉。3.2 数据准备ImageNet的预处理细节ImageNet是这类工作的标准测试集。1.28M训练图50K验证图1000类。预处理流程看起来简单但细节决定成败。标准流程是随机裁剪到目标分辨率比如256×256随机水平翻转然后归一化到[-1, 1]。但像素空间扩散对数据增强更敏感因为模型直接看像素任何增强都会改变像素分布。我的经验是不要用ColorJitter、RandomErasing这类会改变像素统计的增强。随机裁剪和翻转就够了。如果你要做类别条件生成标签用one-hot或者可学习的类别嵌入都行后者效果通常更好。数据加载用WebDataset或者FFCV能大幅加速。ImageNet这种规模用普通DataLoaderIO会成为瓶颈。FFCV能把数据加载速度提升5到10倍对于像素空间扩散这种训练慢的任务省下来的时间很可观。# FFCV的典型配置 from ffcv.fields import IntField, RGBImageField from ffcv.writer import DatasetWriter writer DatasetWriter(imagenet_train.beton, { image: RGBImageField(write_modejpg, max_resolution256), label: IntField() }, num_workers16)3.3 模型搭建DiT主干与表示注入的代码骨架下面是一个简化的PixelDiT2模型骨架重点展示表示注入的部分。import torch import torch.nn as nn from einops import rearrange class RepresentationInjector(nn.Module): def __init__(self, rep_dim, hidden_dim, num_heads8): super().__init__() self.proj nn.Linear(rep_dim, hidden_dim) self.cross_attn nn.MultiheadAttention(hidden_dim, num_heads, batch_firstTrue) self.norm nn.LayerNorm(hidden_dim) self.gate nn.Parameter(torch.zeros(1)) # 可学习的门控初始为0 def forward(self, x, rep_feat): # x: [B, N, D] DiT的token # rep_feat: [B, M, rep_dim] 表示编码器输出 rep self.proj(rep_feat) rep self.norm(rep) attn_out, _ self.cross_attn(x, rep, rep) # 门控机制让模型自己决定用多少表示信息 return x self.gate.tanh() * attn_out class PixelDiTBlock(nn.Module): def __init__(self, hidden_dim, num_heads, rep_dim): super().__init__() self.norm1 nn.LayerNorm(hidden_dim) self.self_attn nn.MultiheadAttention(hidden_dim, num_heads, batch_firstTrue) self.norm2 nn.LayerNorm(hidden_dim) self.mlp nn.Sequential( nn.Linear(hidden_dim, hidden_dim * 4), nn.GELU(), nn.Linear(hidden_dim * 4, hidden_dim) ) self.rep_injector RepresentationInjector(rep_dim, hidden_dim, num_heads) self.adaLN_modulation nn.Sequential( nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6) ) def forward(self, x, t_emb, rep_feat): # adaLN调制 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp \ self.adaLN_modulation(t_emb).chunk(6, dim-1) # 自注意力 h self.norm1(x) * (1 scale_msa.unsqueeze(1)) shift_msa.unsqueeze(1) h, _ self.self_attn(h, h, h) x x gate_msa.unsqueeze(1) * h # 表示注入 x self.rep_injector(x, rep_feat) # MLP h self.norm2(x) * (1 scale_mlp.unsqueeze(1)) shift_mlp.unsqueeze(1) h self.mlp(h) x x gate_mlp.unsqueeze(1) * h return x这个骨架里有两个设计点值得说。第一是门控机制self.gate初始化为0经过tanh后也是0意味着训练初期表示注入不起作用模型先学会纯像素去噪然后逐渐打开门控让表示信息进来。这个warm-up策略能防止表示信息在训练初期干扰模型学习基础去噪能力。第二是adaLN的6路输出这是DiT原版的设计分别控制自注意力和MLP的shift、scale、gate。这个设计比简单的LayerNorm效果好很多因为时间步信息能更精细地调制每一层。3.4 训练配置学习率、batch size与EMA像素空间扩散的训练比潜空间慢因为每一步都要处理更多token。我的经验配置是分辨率256×256起步稳定后再上512×512Batch size单卡24GB显存256×256分辨率batch size大概能到32到64学习率1e-4到2e-4用cosine schedulewarmup 5000步优化器AdamWweight decay 0.01EMAdecay 0.9999这个对生成质量影响很大千万别省训练步数ImageNet 256×256大概需要500K到1M步才能收敛注意像素空间扩散的loss曲线比潜空间抖动更大这是正常的。不要看到loss spike就以为训练崩了先看看是不是学习率太大或者batch size太小。如果spike之后能恢复就继续跑。EMA指数移动平均是扩散模型训练的标配。它维护一份模型参数的滑动平均推理时用EMA参数而不是训练参数。实测下来EMA能把FID降低10%到20%代价只是多一份模型参数的显存。这个投资回报率太高了没有理由不用。class EMA: def __init__(self, model, decay0.9999): self.model model self.decay decay self.shadow {k: v.clone().detach() for k, v in model.state_dict().items()} def update(self): for k, v in self.model.state_dict().items(): if v.dtype.is_floating_point: self.shadow[k].mul_(self.decay).add_(v, alpha1 - self.decay) else: self.shadow[k].copy_(v)3.5 推理采样DDIM还是DPM-Solver训练完之后采样器的选择直接影响生成速度和质量的平衡。DDIM是经典选择50到100步能出不错的结果。DPM-Solver更快20到30步就能达到DDIM 100步的质量。PixelDiT2这类像素空间模型我建议用DPM-Solver因为像素空间每一步的计算量比潜空间大减少步数的收益更明显。配置上用Karras sigma scheduleorder2大概25步就能出图。from diffusers import DPMSolverMultistepScheduler scheduler DPMSolverMultistepScheduler( num_train_timesteps1000, beta_schedulelinear, algorithm_typedpmsolver, solver_order2, use_karras_sigmasTrue ) scheduler.set_timesteps(25)如果你要做类别条件生成还需要classifier-free guidance。guidance scale一般设1.5到3.0太低没效果太高会过饱和。像素空间模型对guidance scale比潜空间敏感建议从1.5开始试。4. 常见问题与排查技巧实录4.1 训练不收敛或loss震荡怎么办这是像素空间扩散最常见的问题。排查顺序如下第一检查数据归一化。像素空间模型对输入范围极其敏感。如果你用[0, 1]而不是[-1, 1]loss会大很多收敛也慢。确认你的归一化是x x * 2 - 1。第二检查时间步采样。扩散模型训练时时间步t是从[0, T]均匀采样的。但如果你的模型在某个时间段loss特别大可以考虑用重要性采样让模型多训练那些难的时间步。第三检查表示注入的门控。如果门控打开得太早表示信息会干扰基础去噪能力的学习。确认你的gate初始化是0并且有warm-up。第四降低学习率。像素空间扩散的梯度比潜空间大学习率需要相应降低。如果1e-4震荡试试5e-5。4.2 生成结果模糊或细节丢失像素空间扩散理论上不该模糊如果模糊了通常是这几个原因表示约束太强门控值太大模型过度依赖表示信息忽略了像素级细节。试着降低门控的初始值或者加一个上限。patch size太大如果patch size是8每个token覆盖64个像素细节建模能力肯定不够。降到4或2试试。采样步数太少DPM-Solver虽然快但步数太少时高频细节会丢。试试把步数从25加到50。EMA decay太小如果EMA decay是0.999而不是0.9999滑动平均的窗口太短模型参数抖动大。调到0.9999或0.99999。4.3 显存不够用的优化手段像素空间扩散显存吃紧是常态。除了换卡还有这些手段优化手段显存节省速度影响实现难度梯度检查点30-50%慢20-30%低混合精度训练30-40%快10-20%低xformers注意力20-30%快10-20%低分块注意力40-60%慢10-20%中梯度累积线性节省慢低梯度检查点是最简单有效的torch.utils.checkpoint.checkpoint包一下就行。混合精度用torch.cuda.amp注意扩散模型的loss在fp16下容易溢出建议用bf16。xformers的memory_efficient_attention直接替换掉nn.MultiheadAttention省显存还提速。提示梯度累积虽然省显存但会改变batch norm的统计如果你用了BN。扩散模型一般用LayerNorm所以问题不大。但要注意学习率需要按累积步数缩放。4.4 表示编码器的常见坑表示编码器虽然冻结了但用起来还是有坑。第一个坑是输入分辨率不匹配。DINOv2训练时用的是固定分辨率如果你的输入分辨率差太多表示质量会下降。解决办法是resize到编码器的原生分辨率或者插值位置编码。第二个坑是表示特征的维度。DINOv2 base是768维large是1024维。如果你的DiT隐藏维度是512需要投影。投影层用线性层就行别搞太复杂。第三个坑是表示特征的归一化。DINOv2的输出没有归一化直接拿来做cross-attention的key和value数值范围可能很大。建议先做LayerNorm再投影。第四个坑是推理时的表示计算。训练时表示编码器只跑一次对干净图像但推理时每一步去噪都要跑一次表示编码器因为当前噪声图像在变。这会显著增加推理时间。一个优化是每隔几步才更新一次表示或者用上一次的表示做近似。4.5 常见问题速查表现象可能原因排查方法解决方案loss不下降学习率太大/数据归一化错误打印梯度范数/检查输入范围降学习率/修正归一化生成全黑或全白时间步嵌入错误/采样器配置错误检查t_emb的数值范围修正嵌入/换采样器生成结果重复表示约束太强/guidance太高降低门控值/降低guidance调整门控/guidance scale训练后期loss反弹过拟合/EMA decay太小看验证集loss加数据增强/调EMA推理速度慢表示编码器每步都跑profile推理流程缓存表示/隔步更新显存OOMbatch size太大/分辨率太高看显存峰值梯度检查点/降batch5. 这套方案的实际价值与适用边界5.1 什么时候该用PixelDiT2这类方案不是所有场景都值得上像素空间扩散。我的判断标准是当潜空间压缩损失成为瓶颈时才考虑像素空间。具体来说这几类场景值得试文字生成尤其是小字、精细纹理生成织物、毛发、金属拉丝、医学图像生成X光、病理切片、遥感图像生成需要保留地物细节。这些场景的共同点是高频信息重要潜空间的8倍压缩会丢关键细节。反过来如果你只是生成512×512的风景图、人像图潜空间扩散已经够用没必要上像素空间。像素空间的训练成本是潜空间的3到5倍推理成本也高投入产出比不划算。5.2 表示约束的边界什么时候会帮倒忙表示约束不是万能的。如果表示编码器的语义空间和你的生成任务不匹配约束会帮倒忙。比如你用自然图像训练的DINOv2去约束医学图像生成表示编码器可能把病灶区域编码成异常导致生成时回避这些区域反而生成不出病灶。解决办法是用领域数据微调表示编码器或者换一个在领域数据上预训练的表示模型。但微调表示编码器有风险容易把表示空间搞崩。折中方案是冻结编码器但在投影层上加一个领域适配的LoRA只微调很少的参数。另一个边界是表示约束的强度。约束太强生成结果会趋向于表示编码器见过的模式多样性下降。约束太弱又起不到引导作用。这个平衡需要根据具体任务调没有万能参数。5.3 后续可以扩展的方向从PixelDiT2这个思路出发有几个方向值得探索。第一是多尺度表示约束。现在通常只用最后一层的表示特征但浅层特征包含更多细节信息深层特征包含更多语义信息。把多尺度特征都注入进去可能能同时提升细节和语义一致性。第二是表示约束与蒸馏结合。用一个大模型比如SDXL生成图像然后用PixelDiT2去蒸馏同时用表示约束保证蒸馏后的模型不丢语义。这样能用小模型达到大模型的生成质量。第三是视频生成。视频是像素空间扩散的天然场景因为视频的时序一致性需要像素级约束。把PixelDiT2扩展到视频用表示约束保证帧间一致性是个有意思的方向。第四是可控生成。表示编码器提取的特征可以作为控制信号比如用分割图的表示特征去约束生成实现像素级的可控生成。这比ControlNet的注入方式更直接。我在实际跑这类方案时的体会是像素空间扩散的门槛主要在工程上不在算法上。算法思路是清晰的——用表示约束补偿像素空间的语义缺失。但工程上要处理的细节很多显存、速度、数值稳定性、表示编码器的适配。这些细节没有捷径只能一个个踩坑填坑。如果你准备上手建议先从256×256分辨率、小规模数据开始把流程跑通再上大规模。别一上来就怼512×512全量ImageNet大概率会卡在显存或者训练不收敛上白白浪费时间。