LinFusion: 1 GPU, 1 Minute, 16K Image
作者: Songhua Liu, Weihao Yu, Zhenxiong Tan, Xinchao Wang
分类: cs.CV, cs.LG
发布日期: 2024-09-03 (更新: 2024-10-17)
备注: Work in Progress. Codes are available at https://github.com/Huage001/LinFusion
🔗 代码/项目: GITHUB
💡 一句话要点
LinFusion：利用线性注意力机制，单GPU一分钟生成16K图像
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱七：动作重定向 (Motion Retargeting)
关键词: 线性注意力 扩散模型 高分辨率图像生成 知识蒸馏 低秩近似
📋 核心要点
- 现有扩散模型在高分辨率图像生成中面临时间和内存复杂度高的挑战，主要由于自注意力机制的二次方复杂度。
- LinFusion提出一种广义线性注意力范式，作为现有线性token混合器的低秩近似，降低计算复杂度。
- 通过从预训练Stable Diffusion模型中蒸馏知识，LinFusion在保持或超越原有性能的同时，显著降低了计算成本。
📝 摘要（中文）
现代扩散模型，特别是那些使用基于Transformer的UNet进行去噪的模型，严重依赖自注意力操作来管理复杂的空间关系，从而实现令人印象深刻的生成性能。然而，由于其时间和内存复杂度与空间token的数量呈二次方关系，这种现有范式在生成高分辨率视觉内容方面面临重大挑战。为了解决这个限制，本文旨在采用一种新的线性注意力机制作为替代方案。具体来说，我们从最近引入的具有线性复杂度的模型（例如，Mamba2、RWKV6、门控线性注意力等）开始探索，并确定了两个关键特征——注意力归一化和非因果推理——可以增强高分辨率视觉生成性能。在此基础上，我们引入了一种广义线性注意力范式，它可以作为各种流行的线性token混合器的低秩近似。为了节省训练成本并更好地利用预训练模型，我们初始化我们的模型并从预训练的StableDiffusion (SD) 中提炼知识。我们发现，经过适度训练后，这种被称为LinFusion的蒸馏模型可以达到与原始SD相当或优越的性能，同时显著降低时间和内存复杂度。在SD-v1.5、SD-v2.1和SD-XL上的大量实验表明，LinFusion能够实现令人满意且高效的零样本跨分辨率生成，可以在单个GPU上处理像16K这样的超分辨率图像。此外，它与预训练的SD组件和pipeline（如ControlNet、IP-Adapter、DemoFusion、DistriFusion等）高度兼容，无需进行适配。
🔬 方法详解
问题定义：现有基于Transformer的扩散模型在生成高分辨率图像时，自注意力机制的计算复杂度和内存占用呈平方级增长，限制了其在高分辨率图像生成上的应用。现有方法难以在单GPU上高效生成超高分辨率（如16K）图像，且训练成本高昂。
核心思路：LinFusion的核心思路是用线性注意力机制替代传统的自注意力机制，从而将计算复杂度从平方级降低到线性级。通过广义线性注意力范式，LinFusion可以近似多种线性token混合器，并保留其关键特性，如注意力归一化和非因果推理，以保证生成质量。
技术框架：LinFusion基于Stable Diffusion的UNet架构，将原有的自注意力模块替换为提出的线性注意力模块。整体流程包括：1) 初始化模型参数；2) 从预训练的Stable Diffusion模型中蒸馏知识；3) 使用少量数据进行微调。该框架可以无缝集成到现有的Stable Diffusion pipeline中，如ControlNet等。
关键创新：LinFusion的关键创新在于提出了广义线性注意力范式，它不仅降低了计算复杂度，还通过注意力归一化和非因果推理保持了生成质量。这种范式可以看作是对多种线性token混合器的低秩近似，具有通用性和灵活性。
关键设计：LinFusion的关键设计包括：1) 使用低秩近似来构建线性注意力模块；2) 通过注意力归一化来稳定训练过程；3) 利用非因果推理来增强生成效果；4) 通过知识蒸馏从预训练的Stable Diffusion模型中迁移知识，减少训练成本。具体参数设置和损失函数细节未知。
🖼️ 关键图片
📊 实验亮点
LinFusion在SD-v1.5、SD-v2.1和SD-XL上进行了广泛实验，结果表明，经过适度训练后，LinFusion的性能可以达到与原始SD相当或优越的水平，同时显著降低了时间和内存复杂度。LinFusion能够在单GPU上生成16K超高分辨率图像，并且与ControlNet、IP-Adapter等现有SD组件和pipeline高度兼容，无需额外适配。
🎯 应用场景
LinFusion在图像生成领域具有广泛的应用前景，尤其是在需要高分辨率图像的场景中，如游戏开发、电影制作、广告设计、虚拟现实等。其高效的生成能力使得在消费级硬件上生成超高分辨率图像成为可能，降低了生成成本，加速了内容创作流程。未来，LinFusion有望应用于更多需要实时或高效图像生成的领域。
📄 摘要（原文）
Modern diffusion models, particularly those utilizing a Transformer-based UNet for denoising, rely heavily on self-attention operations to manage complex spatial relationships, thus achieving impressive generation performance. However, this existing paradigm faces significant challenges in generating high-resolution visual content due to its quadratic time and memory complexity with respect to the number of spatial tokens. To address this limitation, we aim at a novel linear attention mechanism as an alternative in this paper. Specifically, we begin our exploration from recently introduced models with linear complexity, e.g., Mamba2, RWKV6, Gated Linear Attention, etc, and identify two key features--attention normalization and non-causal inference--that enhance high-resolution visual generation performance. Building on these insights, we introduce a generalized linear attention paradigm, which serves as a low-rank approximation of a wide spectrum of popular linear token mixers. To save the training cost and better leverage pre-trained models, we initialize our models and distill the knowledge from pre-trained StableDiffusion (SD). We find that the distilled model, termed LinFusion, achieves performance on par with or superior to the original SD after only modest training, while significantly reducing time and memory complexity. Extensive experiments on SD-v1.5, SD-v2.1, and SD-XL demonstrate that LinFusion enables satisfactory and efficient zero-shot cross-resolution generation, accommodating ultra-resolution images like 16K on a single GPU. Moreover, it is highly compatible with pre-trained SD components and pipelines, such as ControlNet, IP-Adapter, DemoFusion, DistriFusion, etc, requiring no adaptation efforts. Codes are available at https://github.com/Huage001/LinFusion.