方法贡献

MOSAIK：像素扩散模型如何学会“看人下菜碟”

作者

快速导读：MOSAIK通过预测不同区域粗化造成的保真度损失，动态分配异构分块大小，在显著降低计算量的同时保持像素空间扩散模型的生成质量。

像素空间扩散模型（Pixel-space Diffusion Models）直接在图像空间进行去噪生成，避免了VAE压缩带来的高频细节损失，但代价是自注意力机制的二次复杂度导致计算成本极高。MOSAIK的核心洞察在于：不同图像区域对分块粗化的敏感度存在显著差异——人脸、文字等细节密集区域需要细粒度令牌，而天空、背景等平坦区域即使使用大分块也几乎不影响最终质量。

MOSAIK提出了一套完整的“损伤引导异构分块”框架，包含三阶段训练：先让模型学会处理不同统一分块大小，再训练其处理任意混合分块布局，最后训练一个轻量级预测器根据中间特征估计每个区域的粗化损伤。推理时，预测器在给定令牌预算下将精细分块分配给高损伤区域，粗分块分配给低损伤区域，实现内容感知的动态计算分配。

英文题目：MOSAIK: Multi-Patch Content-Aware Spatial Allocation of Image Tokens for Efficient Generation

论文出处：arXiv 每日论文精选 · arXiv:2608.05450

原始论文：PDF / 论文页面

这篇论文解决什么问题？

扩散模型的主流范式是通过VAE将图像压缩到潜空间进行生成，但VAE的重建过程不可避免地丢失高频细节，尤其在人脸、文字、纹理等区域表现明显。像素空间扩散模型直接在原始图像空间操作，从根本上避免了这一瓶颈，但代价是令牌数量大幅增加——以1024×1024分辨率为例，使用16×16分块会产生4096个令牌，自注意力的计算量随令牌数平方增长。

现有高效推理方法主要沿两个方向展开：一是时间维度的优化，如DDiT在不同去噪步使用不同统一分块大小，TeaCache和TaylorSeer通过缓存特征跳过冗余计算；二是减少去噪步数。这些方法都忽略了一个关键事实：在同一张图像内部，不同空间区域对计算精度的需求是不均匀的。统一的分块策略要么在简单区域浪费计算，要么在关键区域损失细节。

MOSAIK正是抓住了这一空间异质性。论文通过实验证明，在潜空间进行分块粗化造成的感知损伤远大于在像素空间进行同等程度的粗化（Figure 7, p.12），这意味着像素空间扩散模型天然更适合进行空间维度的动态令牌分配——因为粗化带来的代价更小，优化空间更大。

核心创新

  • 提出损伤引导的异构分块空间分配框架，首次在像素空间扩散模型中实现内容感知的动态令牌分配。
  • 设计三阶段训练策略，解耦异构布局的执行与分配策略的学习，无需端到端联合优化。
  • 引入基于中间特征的轻量级损伤预测器，以极小的计算开销（每步2.0ms）指导令牌路由。

方法概览

MOSAIK包含三阶段训练：1) 训练多分块生成器以匹配教师模型在不同统一分块大小下的输出；2) 训练生成器处理任意异构分块布局；3) 训练一个轻量级预测器，根据中间去噪特征估计每个区域粗化造成的“损伤”，并在推理时根据令牌预算动态分配精细分块给高损伤区域。

  • 三阶段自蒸馏训练框架（Figure 2, p.2）：整个训练基于冻结的PixelDiT教师模型，学生模型通过LoRA进行高效微调。第一阶段训练学生模型匹配教师在不同统一分块大小（16、32、64）下的输出，使其掌握多种分块粒度下的生成能力。
  • 第二阶段训练混合分块生成器：将图像划分为网格区域，随机分配不同的分块大小，形成异构分块布局。学生模型学习在单次前向传播中处理这种混合布局，输出与教师模型一致的结果。这一阶段的关键在于让模型适应任意分块组合，为后续的动态分配奠定基础。
  • 第三阶段训练损伤预测器：从去噪网络的中间层特征（h6）和池化后的提示嵌入（c_pool）出发，预测每个区域因粗化（从16到32、从32到64）造成的保真度损失。预测器结构极轻量，仅包含少量卷积层，推理开销约每步2.0ms。
  • 推理时的动态分配策略（Figure 3, p.3）：在每个去噪步，损伤预测器输出每个区域的粗化损伤分数。根据预设的令牌预算，将精细分块（16×16）优先分配给损伤最高的区域，中等分块（32×32）分配给中等损伤区域，粗分块（64×64）分配给低损伤区域。随着去噪推进，损伤分布逐渐变化，布局也随之动态调整。
  • 损伤分数的设计：预测器输出两个损伤分数，分别对应p64→p32和p32→p16的细化收益。实验表明，共享同一损伤分数指导两种细化转换虽非最优，但已足够有效（Table 7, p.13）。
  • 与特征缓存方法的结合：MOSAIK的空间优化与TeaCache、TaylorSeer的时间优化是正交的。实验证明两者可以直接叠加，在80% FLOPs缩减下，结合TeaCache后GenEval得分0.73，而单独使用TeaCache降至0.67（Table 2, p.7）。

逐图理解论文

直觉与失败案例

直觉与失败案例
Figure 4: Qualitative comparisons of MOSAIK against DDiT, TeaCache, and TaylorSeer at a 60% FLOPs reduction (∼73% token reduction). The base model (i.e., PixelDiT) with the highest FLOPs is included as a reference. Zoomed-in regions are provided to better illustrate fine-grained details.

想象一张肖像照：人脸需要精细刻画，但背景的天空即使模糊一些也无关紧要。然而现有方法要么全图统一分块，要么只在时间维度上调整——它们都忽略了一个事实：同一张图里，不同区域对计算精度的需求天差地别。DDiT在不同去噪步切换分块大小，但每一步内全图一视同仁；TeaCache缓存特征跳过计算，但缓存策略同样不看内容。结果就是，简单区域浪费了计算，关键区域反而细节不足。

研究空白与核心区分

研究空白与核心区分
Figure 7: Latent-to-Pixel Coarsening Damage Ratio. Each bar is the LPIPS damage from coarsening a VAE latent divided by the damage from coarsening native pixels at a matched factor. The dashed line marks the equal-damage reference at 1.0.

这里的关键区分在于：现有方法做的是“时间维度的统一调度”，而真正需要的是“空间维度的内容感知分配”。更关键的是，论文发现了一个被忽视的优势——在像素空间进行分块粗化，造成的感知损伤远小于在潜空间做同样的事。这意味着像素扩散模型天然更适合动态令牌分配，因为粗化的代价更小，优化的空间更大。

方法贡献

方法贡献
Figure 2: MOSAIK training pipeline. The model is trained in three stages via self-distillation from a frozen PixelDiT teacher. Stage 1 trains a model to render all uniform patch sizes (P ∈{16, 32, 64}). Stage 2 adapts the model to process arbitrary mixed-patch layouts in a single forward pass, yielding a mixed-patch generator. Stage 3 trains an offline layout predictor to regress a per-region coarsening damage map from intermediate activations (h6) and prompt embeddings (cpool).

三阶段训练流水线概览。第一阶段训练模型处理不同统一分块大小，第二阶段训练处理任意混合分块布局，第三阶段训练损伤预测器。三个阶段均通过自蒸馏从冻结的PixelDiT教师模型学习，学生模型使用LoRA进行高效微调。

核心证据

核心证据
Table 1: Quantitative comparison against step reduction, feature caching, and temporal patch scheduling. TR (%) is the token reduction relative to the patch-16 baseline (4096 tokens). Bold and underline mark the best and second-best score per column within each compute tier. †: Derived from schedule-averaged tokens.

主要定量对比表。在多个FLOPs缩减档位下，MOSAIK在GenEval和DPG-Bench上均优于DDiT、TeaCache、TaylorSeer等方法。特别值得注意的是，在70% FLOPs缩减下MOSAIK的GenEval得分与全计算基线持平。

结论与局限

结论与局限
Table 2: Composing the spatial and temporal axes. MO- SAIK integrates seamlessly with caching to enhance perfor- mance under aggressive compression.

MOSAIK证明了像素空间扩散模型的高效推理，关键在于“看人下菜碟”——根据内容损伤动态分配计算资源。这一空间维度的优化与TeaCache等时间维度方法完全互补，两者结合能实现更强的压缩。不过也要看到，像素解码器等固定开销限制了纯空间方法的压缩上限，而且损伤预测器目前绑定特定模型和分块尺寸，换一个架构就得重新训练。以上就是这篇论文的核心内容，感谢收听。

实验如何设计？

  • 基础模型：PixelDiT-XL/2，在1024×1024分辨率下生成，默认使用16×16分块（4096个令牌）。
  • 对比基线：DDiT（时间维度动态分块调度）、TeaCache和TaylorSeer（特征缓存）、统一分块大小（p16/p32/p64）、减少去噪步数。
  • 评估基准：GenEval（文本到图像对齐质量）和DPG-Bench（密集提示图生成质量），辅以DISTS和LPIPS感知距离指标。
  • 计算量控制：以整个去噪器的FLOPs缩减百分比为统一度量，而非仅动态部分的缩减，报告数据较为保守（p.10-11）。
  • 实际延迟测量：在消费级GPU上测量端到端生成时间，验证FLOPs缩减到实际加速的转化效率（Figure 6, p.10）。

关键结果与论文证据

  • 在70% FLOPs缩减下，MOSAIK的GenEval得分0.74，与全计算PixelDiT持平；DPG-Bench得分82.5，仅下降1.0分（Table 1, p.6）。这表明MOSAIK在节省大量计算的同时几乎无损质量。
  • 在60% FLOPs缩减下，MOSAIK的GenEval得分0.74，显著优于DDiT的0.65和TeaCache的0.71（Table 1, p.6）。空间维度的动态分配比时间维度的统一调度更有效。
  • 损伤引导布局相比随机布局，在最大压缩档位下GenEval得分提升约3.3个百分点（0.736 vs 0.703，Table 3, p.7），证明损伤预测器的分配策略是关键。
  • 令牌分配分析显示，在60% FLOPs缩减时，主体边界区域获得其面积占比2.17倍的令牌，主体内部获得1.80倍，而背景仅获得0.8倍（Figure 5, p.13）。预测器学会了在语义对象内部进行更细粒度的分配，而非简单的主体/背景二分。
  • 在消费级GPU上，MOSAIK在最大压缩档位下实现2.59倍加速（8.4秒 vs 21.8秒，Figure 6, p.10）。
  • 与TeaCache结合后，在80% FLOPs缩减下GenEval得分0.73，而单独使用TeaCache降至0.67（Table 2, p.7），验证了空间与时间优化的互补性。
  • 消融实验表明，使用损伤引导的分配策略在所有压缩档位下均优于反向损伤布局和随机布局（Table 7, p.13），且预测器对两种细化转换使用共享损伤分数虽非最优但足够鲁棒。
  • 感知距离测量显示，MOSAIK在同等FLOPs缩减下与全计算基线的DISTS和LPIPS距离均小于DDiT和统一分块方法（Table 6, p.12），进一步验证了其保真度优势。

阅读时需要注意

  • 固定计算开销（如像素解码器）限制了纯空间方法的FLOPs缩减理论上限约为75%，无法像时间维度方法那样实现更极端的压缩（p.10-11）。
  • 损伤预测器基于特定模型（PixelDiT）和分块大小（16/32/64）训练，泛化到其他架构或分块尺寸需要重新训练，限制了即插即用性（p.11）。
  • 当前方法仅针对图像生成任务设计，未扩展到视频等时空数据，而视频领域对计算效率的需求更为迫切（p.11）。
  • 损伤预测器使用共享损伤分数指导两种细化转换，虽经验证有效，但可能不是每种转换的最优解，存在进一步优化的空间（p.13）。

关联工作

  • PixelDiT（p.1）：作为基础像素空间扩散模型骨干，MOSAIK在其上进行扩展和效率优化，所有实验均基于PixelDiT-XL/2。
  • DDiT（p.2）：主要对比基线，代表时间维度的动态分块调度方法，在不同去噪步使用不同统一分块大小。MOSAIK在空间维度上对其实现了超越。
  • TeaCache / TaylorSeer（p.5）：特征缓存方法的代表，通过缓存和复用中间特征减少计算。MOSAIK与其对比并展示了互补性，可结合使用。
  • DynaPatch（p.3）：同期工作，在潜空间视频扩散模型中进行动态分块。MOSAIK在像素空间解决了其面临的VAE瓶颈问题，且像素空间的粗化损伤天然更小（Figure 7, p.12）。

发表评论