Muon-Accelerated Attention Distillation for Real-Time Edge Synthesis via Optimized Latent Diffusion

📄 arXiv: 2504.08451v1 📥 PDF

作者: Weiye Chen, Qingen Zhu, Qian Long

分类: cs.CV

发布日期: 2025-04-11

💡 一句话要点

提出Muon-AD框架，加速边缘设备上潜在扩散模型的实时合成。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 边缘计算 扩散模型 注意力蒸馏 模型优化 实时合成

📋 核心要点

  1. 现有扩散模型在边缘设备部署面临计算和内存瓶颈，难以实现实时合成。
  2. Muon-AD通过Muon优化器和注意力蒸馏协同设计，加速收敛并降低内存占用。
  3. 实验表明，Muon-AD在保持合成质量的同时，显著提升了边缘设备的推理速度和效率。

📝 摘要（中文）

视觉合成领域的最新进展利用扩散模型和注意力机制实现了高保真度的艺术风格迁移和逼真的文本到图像生成。然而，由于计算和内存的限制，在边缘设备上进行实时部署仍然具有挑战性。我们提出了Muon-AD，一个协同设计的框架，它集成了Muon优化器和注意力蒸馏，用于实时边缘合成。通过正交参数更新和动态剪枝消除梯度冲突，Muon-AD实现了比Stable Diffusion-TensorRT快3.2倍的收敛速度，同时保持了合成质量（FID降低15%，SSIM提高4%）。我们的框架将Jetson Orin上的峰值内存降低到7GB，并通过混合精度量化和课程学习实现了24FPS的实时生成。在COCO-Stuff和ImageNet-Texture上的大量实验表明，Muon-AD实现了帕累托最优的效率-质量权衡。我们展示了分布式训练期间通信开销减少65%，以及边缘GPU上实时10秒/图像的生成。这些进步为在资源受限环境中普及高质量视觉合成铺平了道路。

🔬 方法详解

问题定义：论文旨在解决在资源受限的边缘设备上实时部署高质量视觉合成模型的问题。现有的扩散模型计算量大，内存需求高，难以在边缘设备上实现实时推理，限制了其应用范围。

核心思路：论文的核心思路是通过协同设计优化器和模型结构，在加速模型收敛的同时降低计算和内存需求。具体来说，利用Muon优化器加速训练，并结合注意力蒸馏减少模型复杂度，从而实现边缘设备的实时合成。

技术框架：Muon-AD框架主要包含以下几个模块：1) Muon优化器：通过正交参数更新和动态剪枝消除梯度冲突，加速模型训练。2) 注意力蒸馏：通过知识蒸馏将大型模型的知识迁移到小型模型，减少模型复杂度。3) 混合精度量化：降低模型参数的精度，减少内存占用和计算量。4) 课程学习：逐步增加训练难度，提高模型的泛化能力。整体流程是先使用Muon优化器和注意力蒸馏训练模型，然后进行混合精度量化，最后使用课程学习进行微调。

关键创新：论文的关键创新在于Muon优化器和注意力蒸馏的协同设计。Muon优化器通过正交参数更新和动态剪枝，有效避免了梯度冲突，加速了模型训练。注意力蒸馏则通过知识迁移，在不显著降低模型性能的前提下，大幅减少了模型复杂度。这种协同设计使得模型能够在边缘设备上实现实时推理。

关键设计：Muon优化器采用正交参数更新，确保不同参数组的更新方向相互正交，避免梯度冲突。动态剪枝则根据参数的重要性，在训练过程中逐步移除不重要的参数，降低模型复杂度。注意力蒸馏采用教师-学生框架，教师模型是大型的预训练扩散模型，学生模型是小型模型，通过最小化学生模型和教师模型的输出差异，实现知识迁移。混合精度量化采用INT8量化，将模型参数从FP32转换为INT8，减少内存占用和计算量。课程学习则从简单的图像开始训练，逐步增加图像的复杂度和难度，提高模型的泛化能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Muon-AD框架在边缘设备Jetson Orin上实现了24FPS的实时图像生成，峰值内存占用降低到7GB。与Stable Diffusion-TensorRT相比，收敛速度提高了3.2倍，FID降低了15%，SSIM提高了4%。在分布式训练中，通信开销减少了65%。

🎯 应用场景

该研究成果可广泛应用于智能安防、自动驾驶、AR/VR等领域。例如，在智能安防中，可以利用该技术在边缘设备上实时生成高质量的监控图像，提高监控效率。在自动驾驶中，可以用于实时生成逼真的场景图像，提高自动驾驶系统的安全性。在AR/VR中，可以用于实时生成高质量的虚拟现实内容，提升用户体验。

📄 摘要（原文）

Recent advances in visual synthesis have leveraged diffusion models and attention mechanisms to achieve high-fidelity artistic style transfer and photorealistic text-to-image generation. However, real-time deployment on edge devices remains challenging due to computational and memory constraints. We propose Muon-AD, a co-designed framework that integrates the Muon optimizer with attention distillation for real-time edge synthesis. By eliminating gradient conflicts through orthogonal parameter updates and dynamic pruning, Muon-AD achieves 3.2 times faster convergence compared to Stable Diffusion-TensorRT, while maintaining synthesis quality (15% lower FID, 4% higher SSIM). Our framework reduces peak memory to 7GB on Jetson Orin and enables 24FPS real-time generation through mixed-precision quantization and curriculum learning. Extensive experiments on COCO-Stuff and ImageNet-Texture demonstrate Muon-AD's Pareto-optimal efficiency-quality trade-offs. Here, we show a 65% reduction in communication overhead during distributed training and real-time 10s/image generation on edge GPUs. These advancements pave the way for democratizing high-quality visual synthesis in resource-constrained environments.