2025年11月27日星期四

再见VAE!英伟达PixelDiT硬刚SD/FLUX:破局像素生成,端到端效果比肩隐空间模型

点击下方卡片,关注"AI生成未来"

👇扫码免费加入AI知识星球,如您有工作需要分享,欢迎联系:aigc_to_future

图片

作者:Yongsheng Yu等

解读:AI生成未来
图片

文章链接:https://arxiv.org/pdf/2511.20645 

亮点直击

  • PixelDiT,一种单阶段、完全基于Transformer的像素空间扩散模型,无需单独的自编码器即可进行端到端训练。
  • 证明了高效的像素建模是实现实用像素空间扩散的关键因素,并提出了一种双层级DiT架构,该架构将全局语义学习与像素级纹理细节解耦。
  • 引入了像素级AdaLN调制机制和像素token压缩机制,共同实现了密集的每像素token建模。
  • PixelDiT在类别条件图像生成和文本到图像生成中均实现了高图像质量,显著优于现有的像素空间生成模型,并接近最先进的隐空间扩散模型。
图片

总结速览

解决的问题

  • 两阶段流程的缺陷:传统的隐空间扩散模型依赖于预训练的自编码器,这导致:
    • 有损重建:自编码器的重建过程会丢失高频细节,限制生成样本的保真度。
    • 误差累积:扩散过程和自编码器重建的误差会累积。
    • 联合优化困难:两阶段 pipeline 阻碍了端到端的联合优化。
  • 像素空间建模的挑战:直接在像素空间进行扩散面临计算效率与生成质量的权衡:
    • 计算开销大:对密集的像素级token进行全局注意力计算,复杂度高,成本高昂。
    • 细节丢失:为降低计算量而采用大尺寸图像块会削弱像素级建模,导致纹理细节生成不佳。

提出的方案

  • PixelDiT模型:一个单阶段、端到端的完全基于Transformer的扩散模型。
  • 核心设计:采用双层级架构来解耦图像语义和像素细节的学习:
    • 块层级DiT:使用大尺寸图像块,在较短的token序列上执行远程注意力,以捕获全局语义和布局。
    • 像素层级DiT:进行密集的像素级token建模,以细化局部纹理细节。

应用的技术

  • 像素级AdaLN调制:利用来自块层级的语义token对每个像素token进行条件调制,使像素级更新与全局上下文对齐。
  • 像素token压缩机制:在执行全局注意力之前压缩像素token,之后再进行解压缩。这使得像素级token建模得以实现,同时保持了全局注意力的计算效率。

达到的效果

  • 图像生成质量
    • 在ImageNet 256×256上达到FID 1.61,大幅超越了之前的像素空间生成模型。
  • 可扩展性
    • 成功扩展至文本到图像生成,并在1024×1024像素空间直接进行预训练。
    • 在GenEval上达到0.74,在DPG-bench上达到83.5,性能接近最佳的隐空间扩散模型。
  • 优势体现
    • 避免了VAE伪影:由于直接在像素空间操作,在图像编辑任务中能更好地保留内容细节。
    • 高效与细节并存:双层级设计结合高效像素建模,实现了高训练效率和快速收敛,同时保留了精细细节。
图片
图片
图片
图片

方法

PixelDiT,这是一种基于 Transformer 的扩散模型,直接在像素空间执行去噪。本工作的目标是在保持潜空间(latent space)方法所具备的收敛行为和样本质量的同时,使像素 Token 建模在计算上更加高效。

双层级 DiT 架构

如图 2 所示,本文采用了一种双层级 Transformer 组织结构,将语义学习集中在粗粒度的Patch 级路径(patch-level pathway)上,并在像素级路径(pixel-level pathway)中利用专门的 Pixel Transformer (PiT) 模块进行细节精修。这种组织方式允许大部分语义推理在低分辨率网格上进行,从而减轻了像素级路径的负担并加速学习,这与文献 [11, 28, 29] 中的观察一致。

图片

Patch 级架构:设输入图像为 。本文构建非重叠的  patch tokens ,其中  为 token 数量,并将它们投影到隐藏层维度 

遵循文献 [7],本文通过将 LayerNorm 替换为 RMSNorm 并在所有注意力层应用 2D RoPE 来增强 DiT 模块。Patch 级路径由  个增强的 DiT 模块组成;对于第  个模块,有:

其中 AdaLN 调制参数由全局条件向量  生成,然后广播到  个 patch tokens 上。这种全局到 Patch 的广播(global-to-patch broadcasting)对所有 patch tokens 应用相同的逐特征(per-feature)AdaLN 参数(即在 patch 级别是 token 无关的),这与随后在像素级路径中使用的逐像素 AdaLN 形成对比。

经过  个模块后,得到语义 tokens 。本着文献 [11, 28] 的设计精神,本文将像素级路径的条件信号定义为 ,其中  是时间步嵌入。这些 tokens 通过逐像素 AdaLN 为 PiT 模块提供语义上下文。

像素级架构:像素级 DiT 由  层 PiT Blocks 组成。它接收像素 tokens 和 Patch 级 DiT 的输出  作为输入,以执行像素 token 建模并生成最终结果。每个 PiT 模块的详细信息如下所述。

设计要点。Patch 级路径仅处理 patch tokens 以捕获全局语义。通过将细节精修委托给像素级路径,本文可以采用较大的 patch 尺寸 ,这缩短了序列长度并加速了推理,同时保留了逐像素的保真度。此外,像素级路径在缩减的隐藏维度 (例如 )下运行,确保密集的逐像素计算保持高效。

Pixel Transformer 模块

每个 PiT 模块包含两个核心组件。首先,逐像素 AdaLN(pixel-wise AdaLN)实现了单个像素层面的密集条件化,使逐像素更新与全局上下文对齐。其次,像素 Token 紧凑化(pixel token compaction)机制减少了像素 tokens 之间的冗余,使全局注意力能够在可控的序列长度上运行。

逐像素 AdaLN 调制。在像素级路径中,每个图像通过线性层被嵌入为"每像素一个 token"的形式:

为了与 Patch 级语义 tokens 对齐,本文将其重塑为  个序列,每个序列包含  个像素 tokens,即 。对于每个 patch,本文形成一个总结全局上下文的语义条件 token 

图片

如图 3(B) 所示,一种直接的 patch 级调制会为 patch 内的所有  个像素重复相同的参数。然而,这无法捕获密集的逐像素变化。相反,本文通过线性投影  将  扩展为  组 AdaLN 参数,从而为每个像素分配独立的调制:

并将  的最后一个维度划分为六组,每组大小为 ,得到 。这些调制参数是可学习的,并且如图 3(C) 所示在每个像素处都是不同的。它们通过逐像素 AdaLN 应用于 ,实现了像素特定的更新;相比之下,Patch 级 AdaLN 将同一组参数广播给 patch 内的所有像素,因此无法捕获此类空间变化。

像素 Token 紧凑化。在像素级路径中,直接对所有  个像素 tokens 进行注意力计算在计算上是不可行的。因此,本文在全局注意力之前将每个 patch 内的  个像素 tokens 压缩为一个紧凑的 patch token,随后将注意力后的表示扩展回像素。这将注意力序列长度从  减少到 ,即减少了  倍;当  时,这产生了 256 倍的缩减,同时通过逐像素 AdaLN 和可学习的扩展保留了逐像素更新。

本文通过一个可学习的"展平"操作来实例化紧凑化算子:一个联合混合空间和通道维度的线性映射 ,并配对一个扩展算子 。这种"压缩–注意力–扩展"流水线保持了全局注意力的高效性。与 VAE 中的有损瓶颈不同,该机制仅为了注意力操作而瞬间压缩表示。至关重要的是,这种紧凑化操作纯粹是为了减少自注意力的计算开销;它不会损害细粒度细节,因为高频信息通过残差连接和有效地绕过像素 token 瓶颈的可学习扩展层得以保留。

用于文生图的 PixelDiT

本文通过多模态 DiT (MM-DiT) 模块扩展了 Patch 级路径,该模块融合了文本和图像语义,而像素级路径保持不变。在每个 MM-DiT 模块中,图像和文本 tokens 形成两个流,具有独立的 QKV 投影。

文本嵌入  由冻结的 Gemma-2 编码器生成。遵循文献 [36],本文在用户提示前添加简洁的系统提示,然后将序列输入文本编码器。生成的 token 嵌入被投影到模型宽度,并用作 MM-DiT 中的文本流。

经验表明,来自 Patch 级路径的语义 tokens 足以将文本意图传达给像素更新。因此,像素级路径在架构上与类别条件模型相同:它对像素 tokens 进行操作,仅通过语义 tokens 和时间步进行条件化。文本 tokens 不会被直接路由到像素流。

训练目标

本文在像素空间采用 Rectified Flow 公式,并使用其速度匹配(velocity-matching)损失训练模型:

遵循文献 [31],本文包含一个对齐目标,鼓励中层 Patch 路径 tokens 与冻结的 DINOv2 编码器的特征一致。总体目标函数为 。类别条件模型和文本条件模型均使用相同的公式。

实验

实验围绕 PixelDiT 在类条件与文本到图像两大任务上的有效性、可扩展性与推断效率展开。整体思路是先在 ImageNet 256×256 的类条件生成上建立像素空间 DiT 的基线与上限,再扩展至 1024² 文本到图像场景,以验证双层级架构在高分辨率与复杂语义条件下的稳定性和质量。

图片

设置与规模化

本工作实例化了三种模型规模(B/L/XL),在 ImageNet-1K 进行类条件训练,默认采用 PixelDiT-XL。训练细节强调了稳定优化与收敛速度:bfloat16 混合精度、AdamW、EMA、高梯度裁剪门限的阶段性设定,以及基于 Rectified Flow 的训练范式。文本到图像方面,采用 Gemma-2 作为冻结文本编码器,并在patch级路径引入 MM-DiT 融合;预训练于 512×512,随后在 1024² 继续微调,数据规模约 2600 万对,覆盖多种纵横比。推断时使用 FlowDPM-Solver(Rectified Flow 形式的 DPMSolver++ 变体),类条件默认 100 步,文本到图像默认 25 步,以平衡质量与时延。

类条件生成(ImageNet 256×256)

在标准 50K 采样评估上,本工作报告了 gFID、sFID、IS、Precision–Recall。与像素生成系的代表方法(如 PixelFlow-XL、PixNerd-XL、JiT-G 等)相比,PixelDiT-XL 以显著更低的训练周期开销达到更好的或相当的 gFID,并在长训练(320 epoch)时取得 1.61 的 gFID 与 4.68 的 sFID,IS 为 292.7,Recall 达到 0.64,显示出在像素空间中逼近甚至改写既有上限的潜力。相较于隐空间的 DiT 系列(如 DiT-XL、SiT-XL、MaskDiT、LightningDiT、REPA/RAE 等),本工作虽不依赖 VAE,但质量与多样性指标已能与强潜在基线竞争,尤其在更长训练与更优 CFG 区间设置时进一步逼近最佳。

图片
图片
图片
图片

文本到图像(512×512 与 1024²)

在 GenEval 与 DPG-Bench 上,本工作重点评测文本对齐、计数、颜色/位置属性与组合关系的稳健性。PixelDiT-T2I 在 512×512 达到 GenEval Overall 0.78,1024² 达到 0.74;DPG-Bench 上分别为 83.7 与 83.5。与同为像素空间的近期系统相比,PixelDiT-T2I 的综合得分更高或更均衡;对比多款主流隐空间扩散系统(如 SDXL、Hunyuan-DiT、Playground 等),在 1024² 的综合表现接近甚至在部分维度上可比,同时参数规模更为紧凑。这表明双层级 DiT 在高分辨率的语义一致性和构图控制方面具有竞争力,缩小了像素空间与隐空间在大模型上的差距。

图片

采样策略与超参敏感性

采样器方面,FlowDPM-Solver 在无 CFG 的 100 步对比中(与 Euler/Heun)综合了较好的 gFID/sFID 与 IS/精确度/召回率权衡,因此成为默认选择。步数上,随着训练成熟度提升(如 400K、1.6M 步),增加采样步数带来的收益更明显;在类条件上 100 步可拿到最佳指标,而 50 步之后的收益递减。CFG 的刻度与生效区间对质量–多样性平衡影响显著:较早期(80 epoch)更偏向较强的引导且全程生效(如 3.25,[0.10,1.00]),更长训练(320 epoch)则偏向较温和引导且截断区间(如 2.75,[0.10,0.90]),得到最低 gFID 与较高召回的综合最优。

消融实验

核心组件的贡献

表 4 量化了不同模型变体中各像素建模组件的贡献。注意,表 4 中的标签 A–C 对应图 3 中的设计示意图。具体而言,本工作使用一个 30 层、经过  patch 化处理且直接在像素空间执行去噪的 DiT 作为基线模型(标记为"Vanilla DiT/16")。该基线模型仅在 patch token 上操作,没有专门的像素级路径,将每个  patch 视为一个高维向量。其在 80 epoch 时取得了 9.84 的 gFID。

图片

若引入双层级架构但不使用像素 token 紧凑化(pixel token compaction),会导致全局注意力计算量随像素数量呈二次方增长,从而引发显存溢出(OOM)。加入像素 token 紧凑化解决了这一瓶颈,它将全局注意力的序列长度从  个像素缩短为  个 patch,在同样的 80 epoch 预算下将质量显著提升至 3.50 gFID。

引入逐像素 AdaLN(pixel-wise AdaLN)进一步将逐像素更新与 patch 级路径产生的语义上下文对齐,使 gFID 在 80 epoch 时提升至 2.36,并在 320 epoch 时达到 1.61。

模型变体 A、B 和 C 之间的比较证明了每个提议组件的重要性。更重要的是,完整 PixelDiT 模型 C 与 Vanilla DiT/16 A 的对比表明,像素级 token 建模在像素生成模型中起着关键作用。若没有像素建模,即视觉内容仅在 patch 级别进行学习,模型将难以学习精细细节,视觉质量也会显著下降。

像素 Token 紧凑化分析

Token 紧凑化对于实现像素空间训练的可行性至关重要。对  个像素 token 进行全局注意力会产生  的显存占用和  的 FLOPs,即便在  分辨率下也会产生数十亿个注意力条目,如表 5 中该变体报告的 82,247 GFLOPs 所示。利用像素 token 紧凑化将像素分组为  的 patch,可将序列长度减少至 ,从而产生  倍的注意力开销缩减。

为了分析像素级路径中注意力的作用,本文包含了一个"无像素路径注意力(No Pixel-Pathway Attention)"的消融实验,该实验移除了注意力操作,仅在像素级保留逐像素 AdaLN 和 MLP。如表 5 所示,尽管该变体减少了 GFLOPs,但在不同的训练迭代次数下(例如从 80 到 160 epoch),其表现始终逊色于完整的 PixelDiT 模型,gFID 和 IS 均有明显的性能下降(degradation)。这表明紧凑的全局注意力对于将局部更新与全局上下文对齐是必要的。

模型规模与 Patch 大小的影响

本文研究了 Patch 大小  对不同规模模型性能的影响:PixelDiT-B、PixelDiT-L 和 PixelDiT-XL。对于所有评估,本文使用相同的 CFG 引导比例 3.25,区间为 。本文在 ImageNet 256×256 上评估了 4、8、16 和 32 的 Patch 大小;图 5(a) 可视化了相应的收敛行为。

对于 Base 模型,将  减小到  和  显著加速了收敛:在 200K 次迭代时,gFID 从 48.5 (B/32) 降至 15.1 (B/16) 和 6.7 (B/4),且 B/4 最终在 500K 次迭代时达到 3.4 gFID。更大的模型遵循类似的趋势,但极小 patch 带来的收益随着模型规模的增加而递减。对于 PixelDiT-L,使用  而非  仅适度改善了 gFID(在 300K 迭代时从 2.72 降至 2.15),而对于 PixelDiT-XL,( 和  之间的差距进一步缩小... 注:原文此处截断)。

图片

结论

本文重新审视了像素空间的扩散建模,并证明通过适当的架构设计,像素空间扩散Transformer能够实现高保真度和高效率,而无需依赖预训练的自编码器。PixelDiT将像素建模分解为双层级Transformer设计,引入像素级AdaLN和像素token压缩技术,在保持注意力计算可承受的同时,将全局语义与逐像素token学习解耦。在类别条件图像生成和文本到图像生成任务上的实验表明,该设计显著缩小了隐空间方法与像素空间方法之间的性能差距,并在高分辨率下实现了强劲性能。

尽管由于原始数据维度较高,像素空间扩散相比隐空间方法需要更高的计算成本,但本文的工作有效缩小了这一效率差距。总体而言,PixelDiT证明了实用像素空间扩散的主要障碍并非表示空间本身,而是缺乏高效的像素建模架构。

参考文献

[1] PixelDiT: Pixel Diffusion Transformers for Image Generation

技术交流社区免费开放

这是一个高质量AIGC技术社群。

涉及 内容成/理解(图像、视频、语音、文本、3D/4D等)、大模型、具身智能、自动驾驶、深度学习及传统视觉等多个不同方向。这个社群更加适合记录和积累,方便回溯和复盘。愿景是联结数十万AIGC开发者、研究者和爱好者,解决从理论到实战中遇到的具体问题。倡导深度讨论,确保每个提问都能得到认真对待。

图片
欢迎扫码加入
图片



技术交流

加入「AI生成未来社区」群聊,一起交流讨论,涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向,备注不同方向邀请入群!可添加小助手备注方向加群!

图片

没有评论:

发表评论

别傻了!养生赛道哪有不用AI获取流量的,不过要变现,还是要看谈单转化

今天有个好久没联系的伙伴问我:"指挥官,有没有做过AI赋能大健康的项目? 今天有个好久没联系的伙伴问我:"指挥官,有没有做过AI赋能大健康的项目?" 我听完笑了,翻看聊天记录,他是24年初就加了我微信,将近两年了,还在AI的门外犹犹豫豫。 先甩两...