多补丁全局到局部Transformer架构用于高效流匹配和扩散模型 MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
提出分层多尺度补丁transformer架构，将计算成本降低50%同时保持高质量生成
前置知识
扩散模型与流匹配
扩散模型通过逐步添加噪声来学习数据分布，训练过程是从干净数据逐渐添加噪声，生成过程是从纯噪声逐步去噪得到目标样本。流匹配是更通用的框架，通过学习从噪声分布到数据分布的速度场，使用ODE/SDE方程进行建模。训练目标是学习速度预测模型fθ，使得预测速度接近目标速度(n-z)，其中n是噪声，z是干净潜变量。这类模型的训练损失函数为LFM = Σ||fθ(zt, t, c) - (n-z)||²₂。
本文基于流匹配训练框架，所有改进都是在流匹配的框架下进行。理解流匹配的连续动力学性质是理解FNO时间嵌入创新的前提。
DiT架构与自注意力
DiT（Diffusion Transformer）将图像分割成固定大小的patch，转换为token序列，然后通过标准的Transformer架构处理。每个Transformer块包含多头自注意力机制和MLP层，通过LayerNorm和残差连接。自注意力的计算复杂度是O(n²)，其中n是token数量，这成为计算瓶颈。DiT采用AdaIN（Adaptive Instance Normalization）将时间步和类别条件注入到每个块中，通过缩放和偏移特征图来调节网络行为。
本文的主要改进针对DiT架构的等距设计问题，所有实验都是和DiT/SiT基线对比。需要理解DiT的计算瓶颈和条件注入方式才能理解MPDiT的改进动机。
潜变量扩散模型与VAE
潜变量扩散模型不直接在像素空间操作，而是通过训练好的VAE（变分自编码器）将图像压缩到低维潜变量空间。编码器E将图像x映射到潜变量z=E(x)，维度通常从H×W×3压缩到h×w×d，其中h=H/r，w=W/r，r是压缩比（如8）。解码器D将潜变量重建为图像。在ImageNet-256设置中，VAE输出的潜变量形状是(32, 32, 4)，对应1024个token。潜变量空间的压缩大大减少了计算量，是高效扩散训练的基础。
本文的所有实验都在潜变量空间进行，MPDiT的patch size设计基于(32, 32, 4)的潜变量形状。需要理解VAE压缩机制才能理解为什么可以用64个token表示256个token的信息。
研究动机
现有的DiT架构采用等距设计，每个Transformer块都处理相同数量的token，导致计算成本高昂。以ImageNet-256为例，标准DiT-B/2在潜变量空间有256个token，自注意力复杂度为O(256²)。即使使用更大patch size（DiT-XL/4）减少到64个token，虽然降低了计算量，但会损失局部细节，生成质量下降。MaskDiT尝试通过掩码建模减少token数量，但在75%掩码率下FID高达100，而DiT-XL/4只需40 FID，且不需要额外的解码器层。这表明简单减少token数量会同时损害全局和局部建模能力。
本文的目标是本文目标是设计一种分层transformer架构，在早期阶段使用大patch捕捉全局上下文，在后期阶段使用小patch细化局部细节，从而在大幅降低计算成本（目标减少50% GFLOPs）的同时保持或提升生成质量。此外，还希望通过改进时间嵌入和类别嵌入设计来加速训练收敛，使得在有限训练预算下达到更好的性能。
与已有工作不同的是，现有方法主要在注意力层级别应用全局-局部注意力（如Swin Attention、Focal Attention），这些方法需要在每个块内执行额外的支持操作（reshape、pooling、窗口划分），实际效率提升有限。本文的独特之处是在网络架构级别应用全局-局部思想：早期块完全在大patch token上运行，后期块完全在小patch token上运行，而不是在每个块内混合不同尺度的注意力。这种设计避免了重复的支持操作，实现了更好的效率-性能平衡。此外，本文还重新审视了被忽视的时间嵌入和类别嵌入模块，提出了FNO时间嵌入和多token类别嵌入，这些改进带来的性能提升（约10 FID点）甚至超过架构改进本身。
核心方法
MPDiT采用全局到局部的分层设计。首先，输入的潜变量通过PatchEmbedₚ₌₄转换为64个大patch token，这些token经过N-k个Transformer块处理，捕获全局上下文信息。然后，上采样模块将64个token扩展为256个小patch token，通过线性投影+pixel-unshuffle操作实现4倍空间扩展。最后，k个Transformer块在这些细粒度token上运行，细化局部细节。除了架构创新，MPDiT还改进了条件注入机制：使用FNO（Fourier Neural Operator）编码时间步，通过频域卷积学习平滑的时间依赖；使用多个可学习的token表示类别信息，而不是单个token，提供更丰富的语义表示。
核心创新点在于架构级别的全局-局部分离。早期块（N-k个）只用64个token（25%的原始数量）建模全局结构，后期块（k=4-6个）用256个token建模局部细节。这种设计基于一个关键洞察：全局结构可以用更少的token表示，而局部细节只需要在最后阶段细化。另一个关键创新是FNO时间嵌入：传统方法用正弦编码+MLP，本文将时间步t与1D网格结合，构建一个1D时间信号，然后通过3个MixedFNO块（谱卷积+局部卷积）处理，在频域和时间域同时建模，捕捉流匹配的连续动力学特性。多token类别嵌入也很关键：每个类别用m=16个token表示，这些token被添加到图像token序列前面，通过注意力机制与图像token交互，比单个token提供更强的监督信号。
方法步骤详情
MPDiT的完整流程包括四个步骤。步骤1：大patch嵌入。给定潜变量zt∈ℝ³²×³²×⁴，使用PatchEmbedₚ₌₄将其转换为64个token（8×8空间网格），同时保留原始潜变量的PatchEmbedₚ₌₂结果（256个token）用于后续skip connection。步骤2：全局上下文建模。64个token与m=16个类别token拼接，经过N-k个Transformer块处理（对于MPDiT-B，N=12，k=6；对于MPDiT-XL，N=28，k=6）。这些块使用共享的FNO时间嵌入，每块包含自注意力和MLP。步骤3：上采样。分离类别token和图像token，图像token通过Linear+PixelUnshuffle从64扩展到256个token，然后与原始256个skip connection token相加，保留高频细节。经过GELU后与类别token重新拼接，通过一个轻量级Linear层重新对齐类别和图像token的关系。步骤4：局部细化。最后k个Transformer块在这些256个token上运行，生成速度预测，训练损失为||fθ(zt, t, c) - (n-z)||₂。
技术新颖性
MPDiT的技术新颖性体现在三个方面。一是架构层面的层次化设计，避免了传统全局-局部注意力在每个块内的重复操作，第一次将多尺度思想应用到网络层级别。二是FNO时间嵌入的创新性应用，首次将Neural Operator的频域建模引入到扩散模型的时间编码中，更好地匹配流匹配的ODE动力学本质。三是多token类别嵌入，将单个密集向量替换为多个分布式token，这在理论上提供了更强的表达能力，实验证明加速了训练收敛。这三个改进相互独立但协同作用，共同带来了约10 FID点的性能提升和50%的计算减少。
实验结果
在ImageNet-256上，MPDiT-XL在240个训练epoch内达到非cfg FID 7.36和cfg FID 2.05（cfg scale=1.4），而SiT-XL需要1400个epoch才能达到FID 9.35。MPDiT的训练效率是SiT的约11.36倍（240/1400=17.2%的epoch数）。计算量方面，MPDiT-XL的GFLOPs是59.3，相比DiT-XL的118.66减少了50.1%。消融实验显示每个组件的贡献：共享AdaIN减少30%参数（130M→90M）但FID只增加0.4；多token类别嵌入（m=16）降低FID约7点；FNO时间嵌入进一步降低FID约4点；完整的MPDiT架构（k=6）将GFLOPs从23降到16.6，FID几乎不变。关于k的选择，k=6在B和XL配置下都是最佳，与等距DiT相比FID下降小于1点，但计算节省显著。在ImageNet-512上，MPDiT-XL在120个epoch内达到FID 2.47，计算量228.4 GFLOPs仅为基线的43.5%，训练总计算量仅为8.7%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet-256类条件生成（80 epochs） | FID↓ | 24.74 (MPDiT-B) | 43.47 (DiT-B/2) | 43.1%提升 |
| ImageNet-256类条件生成（240 epochs, cfg） | FID↓ | 2.05 (MPDiT-XL) | 9.35 (SiT-XL/2, 1400 epochs) | 78.1%提升，训练快11.36倍 |
| ImageNet-512类条件生成（120 epochs, cfg） | FID↓ | 2.47 (MPDiT-XL) | 3.04 (DiT-XL/2, 600 epochs) | 18.8%提升，训练快5倍 |
| 计算效率 | GFLOPs↓ | 59.3 (MPDiT-XL) | 118.66 (DiT-XL/2) | 50.0%减少 |
| 参数效率 | Params(M) | 101.2 (MPDiT-B with shared AdaIN) | 130.0 (DiT-B/2) | 22.2%减少 |
局限与改进
作者承认的主要局限性是架构扩展到大规模设置（如文本到图像模型SDv3、Flux，或文本到视频模型如Sora）仍是开放问题。虽然这些任务看起来有前景，但需要大量计算资源。从观察中，MPDiT在超高分辨率（512及以上）虽然表现优异，但论文中的实验主要在类条件生成场景下，对于更复杂的文本条件生成是否适用需要验证。另外，FNO时间嵌入在实现上比传统方法复杂，需要FFT操作，可能在某些硬件上优化困难。多token类别嵌入虽然有效，但对于有1000个类别的ImageNet，需要学习1000×16×D的嵌入矩阵，当类别数更大时（如开放词汇生成）可能不可行。论文中k值（最后细化块数量）的选择主要基于实验验证，缺乏理论指导，可能在其他数据集上需要重新调参。
独立分析的弱点
MPDiT的架构设计假设全局结构可以用大patch充分表示，这对自然图像有效，但对于纹理细节或高频内容丰富的图像（如医学图像、遥感图像），早期的大patch可能会丢失重要信息，导致后期细化无法恢复。改进方向可以是自适应patch size，根据内容动态调整。上采样模块使用简单的Linear+PixelUnshuffle，没有学习复杂的空间变换，可能在某些场景下信息丢失严重。可以尝试引入可学习的上采样（如子像素卷积、转置卷积）或引入额外的细化路径。FNO时间嵌入使用固定的1D网格（linspace(-1, 1, 32)），网格大小和模式数需要手工调参，可能不适合所有训练场景。可以尝试自适应网格或学习grid参数。多token类别嵌入在类别数很大时内存占用高，对于开放词汇生成需要用文本编码器替代，此时如何设计多token文本嵌入是一个开放问题。MPDiT的层次化设计增加了架构复杂度，上采样块和skip connection增加了实现难度，可能影响部署和优化。
未来方向
作者提出的未来方向是扩展到大规模文本到图像和文本到视频模型，这需要解决如何将多patch思想与文本条件结合的问题。可能的方案是在早期块使用粗糙的文本-图像交互，后期块使用细粒度交互。另一个方向是探索更深层的层次（如p∈{16, 8, 4, 2}）以适应更高分辨率（1024及以上）。FNO时间嵌入可以进一步推广到其他连续条件的建模，如相机参数、光照条件等。多token嵌入的思想可以扩展到其他条件模态，如深度图、边缘图等。与蒸馏技术结合是一个有前景的方向，可以用MPDiT作为教师模型，蒸馏出更小的学生模型，实现部署时的效率提升。另一个有趣的方向是研究MPDiT与其他架构改进的结合，如线性注意力、状态空间模型，探索能否进一步降低计算成本。
复现评估
论文声称代码已开源（https://github.com/quandao10/MPDiT），这是积极信号。实验使用ImageNet标准数据集，训练细节详细：单节点8×A100 40GB GPU，学习率2×10⁻⁴，batch size 1024（256分辨率），EMA decay 0.9999。论文提供了模型配置的完整参数（Table 2）：MPDiT-B有12个块，k=6，模型维度768；MPDiT-XL有28个块，k=6，模型维度1152。消融实验充分，提供了每个组件的独立贡献分析。FNO时间嵌入的完整PyTorch实现（Algorithm 1和2）在补充材料中，包括SpectralConv1d的详细代码，这大大降低了复现难度。论文还提供了大量定性结果（Figure 4-13），声明是非精选的结果，这增加了可信度。潜在的复现挑战是需要8张A100 GPU才能复现大batch size训练，但可以通过梯度累积在更少GPU上实现。总体而言，复现难度中等，开源代码和详细实现细节使其他研究者能够验证和扩展。
论文图表
