用于高效、可重构与可变分辨率扩散建模的可堆叠且可跳过的乐高积木学习

扩散模型在生成逼真图像方面表现出色，但其在训练和采样阶段均伴随着显著的计算开销。尽管已有多种技术致力于缓解这些计算挑战，但一个较少被深入探索的问题是：如何设计一种高效且可适应的网络主干结构，以支持迭代优化过程。当前常用的架构如U-Net和视觉Transformer（Vision Transformer）通常依赖于计算资源密集的深层网络，且缺乏灵活性，难以在不同分辨率下生成图像，或在采样时使用比训练阶段更小的网络结构。本研究提出了一种名为LEGO bricks的新架构，其核心思想是无缝融合局部特征增强（Local-feature Enrichment）与全局内容协调（Global-content Orchestration）。这些LEGO砖块可堆叠组合，构建出一种可在测试阶段动态重构的扩散模型主干网络。通过选择性跳过部分砖块，该架构能够有效降低采样成本，并生成高于训练数据分辨率的图像。每个LEGO砖块利用多层感知机（MLP）增强局部区域特征，并通过Transformer模块进行特征变换，同时在整个网络中保持一致的全分辨率图像表示。实验结果表明，LEGO bricks不仅提升了训练效率、加速了模型收敛，还支持可变分辨率图像生成，同时维持了强大的生成性能。更重要的是，相较于其他方法，LEGO显著减少了采样时间，展现出作为扩散模型关键增强组件的巨大潜力。相关代码与项目主页已公开，访问地址为：https://jegzheng.github.io/LEGODiffusion。