跳转至

MobileWan: Closing the Quality Gap for Mobile Video Diffusion

会议: NeurIPS 2026
arXiv: 2607.06173
代码: https://qualcomm-ai-research.github.io/MobileWan
领域: 视频理解
关键词: 视频扩散模型、移动端部署、循环蒸馏、注意力头剪枝、步骤蒸馏

一句话总结

MobileWan 通过可学习高噪声偏置注意力头剪枝、因果线性混合注意力循环蒸馏、少步采样蒸馏及扩展时序感受野 VAE 解码器,首次将 50 亿参数服务端级视频扩散模型(Wan2.2-5B)部署于商用移动端 NPU,以 20 秒端到端延迟生成 5 秒 480×832 高质量视频。

研究背景与动机

基于扩散 Transformer(DiT)的视频生成模型通过在时空 Token 序列上扩展参数规模至数十亿(如 Wan2.2、CogVideoX、HunyuanVideo),显著提升了画面真实度与时序动态连贯性。然而,现有面向端侧设备的移动视频扩散模型受限于移动芯片严格的峰值内存、自注意力二次方计算复杂度、庞大的 VAE 解码器开销以及严苛的端到端时延预算,通常只能被迫采用 0.4B 至 1.8B 的极小骨干网络(如 Neodragon、SnapGen-V、MobileVD、MVDiT)。这种显著压缩模型容量的妥协方案不可避免地导致了画面细节模糊、运动僵硬以及物理规律失真,使端侧与服务端级视频生成之间存在难以逾越的质量鸿沟。

在移动端部署大尺寸 DiT 面临着全方位的系统性瓶颈:首先,传统自注意力的激活值随帧数与分辨率呈二次方激增,导致 5B 规模模型在长视频生成时不可避免地触发内存溢出(OOM);其次,直接采用粗粒度块剪枝(Block Pruning)会导致 Wan2.2 模型的生成质量灾难性骤降,而简单的纯线性注意力又存在严重的表征能力缺陷;最后,原生 3D 因果卷积 VAE 解码器内存开销巨大,现有轻量化 2D 替代解码器又因时序感受野极度受限而引入剧烈的画面闪烁与跳帧。

本文的核心洞察在于:高质量的端侧视频生成并不意味着必须妥协采用小型骨干模型,服务端级 5B 视频扩散模型完全可以通过精细化结构压缩与循环架构重构在内存受限的移动硬件上高效运行。核心 idea:以 Wan2.2-5B 为基座,结合高噪声偏置的可学习注意力头剪枝、因果线性与 Softmax 混合的块级循环蒸馏(RNN 化)、解耦分布匹配采样步数蒸馏(D-DMD/DMD2)以及长感受野轻量 VAE 解码,实现首个在商用高通骁龙 8 Gen 5 NPU 原生运行的 5B 级端侧视频扩散系统。

方法详解

核心机制

MobileWan 解决移动端部署 5B 视频扩散模型的核心机制包含四大协同维度: 1. 细粒度自注意力头门控剪枝:相比于破坏网络深层拓扑的粗粒度整块剪枝,模型引入可学习标量门控,以流动匹配目标与稀疏度正则化驱动注意力头选择,并在训练中对扩散时间步施加高噪声偏置采样,优先保全决定全局结构与宏观运动的关键注意力头。 2. 因果混合注意力与推理期 RNN 重构:将视频时空 Token 切分为局部时序分块(Chunk),块内采用保留全量表达能力的局部 Softmax 注意力(附带相邻帧重叠缓冲区以消除块边界伪影),历史块的上下文信息则通过可学习多项式特征映射投影至因果线性注意力空间。在推理期将线性注意力累积量转化为固定维度的递归状态向量与矩阵累加器,使得自注意力计算开销从序列长度的 \(\mathcal{O}(N^2)\) 骤降为常数级内存 \(\mathcal{O}(1)\) 和线性计算量 \(\mathcal{O}(N)\)。 3. 少步数流形匹配蒸馏:通过 DMD/D-DMD/DMD2 将数十步的流匹配去噪过程压缩至 2~3 个函数评估步(NFEs),并通过解耦无分类器引导(CFG)增强项与独立噪声采样,根除步数蒸馏常见的画面过度饱和与动态退化。 4. 长历史依赖因果解码:重构轻量 2D 卷积解码器的内存块结构,将因果回看窗口由单帧前向依赖扩展至回看前 4 帧(\(F_{t-4}\)),以极低的延迟代价显著抑制视频抖动与闪烁。

整体框架与流程

整个系统将 Wan2.2-5B 先经 81 帧 480×832(16 FPS)目标分辨率微调,随后依次实施可学习注意力头剪枝、逐块循环注意力蒸馏、少步采样蒸馏及长窗口因果 VAE 调优,并在高通骁龙 8 Gen 5 NPU 上采用 W8A8 混合精度量化部署。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["文本提示词 + 潜空间初态<br/>81×480×832 目标分辨率"] --> B["文本编码器与分辨率微调特征"]
    B --> C["高噪声偏置门控头剪枝 DiT<br/>剪枝 23% 自注意力头 (5.0B→4.7B)"]
    C --> D["因果混合自注意力循环蒸馏<br/>块内 Softmax + 跨块因果线性累加器 (RNN)"]
    D --> E["少步分布匹配蒸馏 (DMD2 / D-DMD)<br/>去噪步数压缩至 2~3 NFEs"]
    E --> F["长因果视窗 VAE 解码器<br/>MemBlock 回看 F_{t-4} 消除频闪"]
    F --> G["端侧 NPU 输出<br/>5秒 16 FPS 视频 (20s 端到端延迟)"]

关键设计

1. 高噪声偏置的可学习注意力头剪枝:克服块剪枝的崩塌敏感性

在 5B 稠密 DiT 中,整块剪枝极易破坏层间复杂的高阶特征流动(移除 3~5 块即引起 VBench 评分雪崩)。MobileWan 转为剪除开销巨大的多头自注意力(MHSA)中的冗余注意力头。针对每个自注意力头引入连续松弛门控标量: $\(g_{b,i} = \operatorname{sigmoid}\left(\frac{\alpha_{b,i}}{\eta}\right)\)$ 其中 \(\alpha_{b,i}\) 为可学习参数(初始化为 \(5.0\),使初始门控接近 1),\(\eta\) 为温度系数。模型在流匹配损失 \(\mathcal{L}_{\mathrm{FM}}\) 基础上加入 \(L_1\) 稀疏惩罚项: $\(\mathcal{L}_{\mathrm{pruning}} = \mathcal{L}_{\mathrm{FM}} + \lambda \sum_{b,i} g_{b,i}\)$ 优化采用三阶段退火策略:第一阶段在 \(\eta=1.0\) 下自由探索软门控;第二阶段将 \(\eta\) 指数退火至 \(0.1\),促使门控趋于二值化;第三阶段将门控彻底二值化并冻结,固定稀疏结构继续微调。关键洞察在于高噪声偏置学习策略:前两阶段的时间步采样采用偏向高噪声的分布 \(\sigma \sim \operatorname{LogitNormal}(1.5, 1)\)。由于高噪声阶段主导着视频的全局轮廓、主干构图与宏观运动流,该偏置迫使模型保留对于时空主干表征不可或缺的核心注意力头;二值化冻结后再切回标准分布 \(\sigma \sim \operatorname{LogitNormal}(0, 1)\)。相比传统贡献度启发式排序剪枝,该可学习方法在相同稀疏度下显著避免了性能退化。

2. 混合注意力与推理期 RNN 重构:消除二次方显存瓶颈

视频序列展开后的 Token 总数庞大(\(N = T \times H \times W\)),全注意力矩阵占用 15.7 GB 显存。MobileWan 将视频序列切分为时间步长为 \(T_c=4\) 的分块 \(X_t \in \mathbb{R}^{N' \times D}\)(\(N'=T_c HW\)),并在时间轴上设置重叠跨度 \(T_o=2\)。对于当前第 \(t\) 个分块,注意力接收域被显式划分为两部分:当前及局部重叠窗口 \(\mathcal{T}_t^S\) 由 Softmax 注意力承载,保证高频局部空间和短程动态的精细建模;更早的历史时间步 \(\mathcal{T}_t^L = \{j \mid j < \max(t N' - T_o HW, 0)\}\) 则完全由因果线性注意力处理: $\(\hat{y}_t = \frac{\sum_{j \in \mathcal{T}_t^{\mathrm{S}}} \exp\left(Q_t k_j^{\top}/\sqrt{D} - c_t\right) v_j + \phi_q(Q_t) s_t}{\sum_{j \in \mathcal{T}_t^{\mathrm{S}}} \exp\left(Q_t k_j^{\top}/\sqrt{D} - c_t\right) + \phi_q(Q_t) z_t}\)$ 其中 \(\phi_q, \phi_k\) 为轻量级多项式核映射(2 层 MLP,3 阶多项式展开)。在采样推理期,历史信息被压缩为常数大小的矩阵累加器 \(s_t \in \mathbb{R}^{D' \times D}\) 与归一化向量累加器 \(z_t \in \mathbb{R}^{D' \times 1}\),并按自回归递归式随块递推更新: $\(s_{t+1} = s_t + \sum_{j \in \mathcal{T}'^{\mathrm{S}}_t} \phi_k(k_j) v_j^{\top}, \qquad z_{t+1} = z_t + \sum_{j \in \mathcal{T}'^{\mathrm{S}}_t} \phi_k(k_j)\)$ 该机制将所有 30 个 Transformer Block 全部转化为恒定显存开销的 RNN 图,将自注意力权重显存从 15.7 GB 剧降至 0.1 GB,使整网能在 9.9 GB 移动端 RAM 内一次性完整装载。

3. 解耦分布匹配采样蒸馏:兼顾极速生成与动态色彩

针对流匹配模型原本需要数十步迭代的问题,MobileWan 探索了 DMD 及其演进方案。标准 DMD 将干净视频预测推入加噪空间后利用教师流模型与假分数流模型(Fake-score model)的梯差驱动单步/少步生成器更新,但常规 3 步蒸馏后视频极易出现色彩过饱和与动态幅度减小(动态度评分从 68.33 跌至 45.83)。MobileWan 进一步解耦其核心梯度为分布匹配项(DM)与无分类器引导增强项(CA): $\(\nabla_{\theta} L_{\mathrm{dmd}}^{\mathrm{core}} = \nabla_{\theta} L_{\mathrm{DM}}^{\mathrm{core}} + (c-1) \cdot \nabla_{\theta} L_{\mathrm{CA}}^{\mathrm{core}}\)$ 其中 CA 项充当类似分数的驱动主力,DM 项负责正则化。作者发现将 CA 的噪声采样水平 \(\tau_{\mathrm{CA}}\) 与当前输入噪声解耦并独立采样,可有效纠正色调漂移;进一步引入判别器特征对抗损失的 DMD2 方案更在保留真实色彩的同时大幅回升运动动态度(达到 74.72),使模型仅需 2~3 个去噪步即可达到 84.03 的优异 VBench 综合分。

4. 扩展时序因果感受野 VAE 解码:根除逐帧 2D 卷积跳闪

Wan2.2 原生 VAE 基于因果 3D 卷积构建,在手机端解码直接 OOM;开源高效替代架构 LightX2V 虽采用结构化 2D 卷积大幅降耗(推理仅 72.6 ms),但其因果记忆块(MemBlock)仅以单帧历史 \(F_{t-1}\) 为条件,有限的时序感受野导致重建视频产生刺眼的块状伪影和跳帧。MobileWan 冻结 Wan2.2 编码器,扩展解码器 MemBlock 内部的因果回看窗口至前 4 帧(即条件依赖覆盖至 \(F_{t-4}\)),并在 80K 合成视频上结合 \(L_2\) 重建损失与感知损失进行端到端蒸馏微调。该架构仅引入适度参数(167.86M),在手机端 NPU 上以 501 ms 的低延迟运行,且将重建 PSNR 从 LightX2V 的 25.01 dB 大幅提至 28.85 dB。

理论分析与设计权衡

在移动端硬件(如高通骁龙 8 Gen 5)上,计算瓶颈通常并非单纯由浮点运算量(FLOPs)决定,而是受制于紧耦合内存(TCM)与系统 RAM 之间的数据搬运流量(Spill-Fill 开销)。 - 剪枝粒度权衡:直接剔除 Transformer 块虽然能线形减少参数和 FLOPs,但 Wan2.2 浅层与深层特征耦合极高,块剪枝引发不可逆的重构损伤。可学习头剪枝在降低 23% 注意力头时仅伴随 0.03% 的 VBench 波动,但单纯头剪枝仍然无法摆脱因上下文序列过长导致的显存溢出。 - 全 RNN 化与内存溢出根治:传统自注意力在 81 帧展开下需要存储数百兆的中间注意力图。将全部 30 个块转换为因果混合循环架构后,上下文长度在推理时被折叠进固定尺寸的累加器 \((s_t, z_t)\) 中。由此,模型彻底消除了多块拆分载入(Split-based execution)所需的反复内存搬运开销,将 NPU 溢出转储流量(Spill/Fill)从基准的 313.2 GB 骤降至 87.6 GB,使全模型单次前向从不可运行降至 6.6 秒,证明了内存常数化重构对移动计算平台体系结构的决定性赋能。

实验关键数据

主实验:综合性能对比

在 VBench 官方评测套件上,MobileWan 与业界主流服务端视频大模型及移动端专用模型进行了系统对比:

模型分类 模型名称 参数量 VBench 总分 ↑ 质量得分 (Quality) ↑ 语义真实度 (Semantic) ↑
服务端模型 (≤5B) Open-Sora Plan V1.3 ~5B 77.23 80.14 65.62
CogVideoX 5B 5B 81.91 83.05 77.33
CogVideoX1.5 5B 5B 82.01 82.72 79.17
LTX-Video ~2B 80.00 82.30 70.79
SANA-Video ~2B 83.71 84.35 81.35
Attention Surgery ~5B 83.21 85.19 75.25
Wan2.2 5B (服务端基准) 5.0B 83.12 84.16 78.95
移动端模型 SnapGen-V <2B 81.14 83.47 71.84
Mobile Video DiT (Mobile) <2B 81.45 83.12 74.76
Neodragon (此前移动端 SOTA) <2B 81.61 83.68 73.36
S2DiT - AR <2B 83.26 85.63 73.79
本文方案 MobileWan (Ours) 4.7B 83.79 85.16 78.32

消融与组件级评测

1. 逐步集成优化组件的性能表现 (Table 5 & Table 8)

评测基于 81×480×832 分辨率,单步生成开销在骁龙 8 Gen 5 NPU 上实测:

优化阶段 TFLOPs 参数量 (B) MHSA 权重 (GB) 端侧 RAM (GB) 单步时延 (s) VBench 总分
Wan2.2 5B (微调基准) 98 5.0 15.7 OOM NA 83.12
+ 头剪枝 (23% 门控) 88 4.7 12.1 OOM NA 83.16
+ 循环蒸馏 (30 块全转) 73 4.8 0.1 9.9 6.6 82.75
+ 步骤蒸馏 (3 步采样) 73 4.8 0.1 9.9 ~20 (端到端) 84.03
+ 优化因果解码器 (端到端) 73 4.8 0.1 9.9 20.3 (整机) 83.79

2. VAE 解码器重建性能与时延对比 (Table 4)

在 DAVIS 数据集(33×512×512)及目标分辨率(81×480×832)实测对比:

解码器配置 参数量 (M) PSNR (dB) ↑ 移动端耗时 (ms) VBench 质量分 ↑
Wan 2.2 原生因果 3D VAE 555.04 32.69 OOM 84.16
LightX2V VAE 9.90 25.01 72.6 83.54
Tiny VAE 9.90 26.77 72.6 83.34
Ours (4帧长因果视窗) 167.86 28.85 501.0 83.99

3. 门控学习中的噪声分布偏置消融 (Table 1)

门控训练噪声分布 VBench 质量分 语义分 VBench 总分
低噪声偏置 (Low Noise Biased) 82.19 72.18 80.19
标准分布 (Standard) 82.61 74.12 80.91
高噪声偏置 (High Noise Biased) 83.64 76.40 82.19

关键发现

  • 打破参数上限:MobileWan 以 83.79 的 VBench 得分显著超越此前所有移动端专用模型(如 Neodragon 的 81.61),且在人类双盲盲测中,对 Neodragon 取得了 80% 的绝对胜率。
  • 高噪声对剪枝的重要性:在注意力门控学习过程中,偏向高噪声的训练分布比标准分布提升了 1.28 分的 VBench 总分,证实了早期去噪时间步对结构和运动建模的关键指引作用。
  • 循环机制的工程飞跃:将所有 30 个 Block 彻底转换为因果循环结构不仅节约了显存,更彻底免除了模型在 NPU 内部进行分块切分换入换出的系统死锁等待,将单步运行时间压缩至 6.6 秒。

亮点与洞察

  • 高噪声偏置门控剪枝:发现扩散模型各时间步对注意力头的敏感度非均匀分布,利用高噪声时间步直接约束结构性门控学习,兼具端到端自适应性与任务物理特异性。
  • 从自注意力到恒定显存 RNN 的无缝过渡:巧妙地将时空混合注意力重构为“局部 Softmax + 历史因果多项式核累加”,兼得 Softmax 的高频细节与线性注意力的时序无限长递归。
  • 软硬件系统级协同优化:结合算子重构、量化感知训练(QAT)与 NPU TCM 数据流图优化,跨越了学术模型与商用边缘算力芯片落地之间的断层。

局限与展望

  • 作者指出的局限:模型在中距离人脸等精细纹理上的生成质量仍有瑕疵;步骤蒸馏易引起部分色彩轻度失真或动态减弱;在极端场景下 RNN 的分块自回归仍偶发时序轻度断裂,结合快速 VAE 偶见频闪。
  • 端侧功耗与热设计约束:生成一个 5 秒视频需持续运行 NPU 约 20 秒,对移动设备的持续电池放电率与瞬时发热仍具一定挑战。
  • 改进方向:可进一步探索针对人脸区域的时空自适应超分补偿模块,以及更紧凑的一步式流形一致性蒸馏(Consistency Distillation)。

相关工作与启发

  • vs Neodragon: Neodragon 采用金字塔下采样与粗粒度块剪枝构建 <2B 骨干网,虽然速度快,但生成画面动态度与细节严重退化;MobileWan 保留 5B 服务端主体能力,通过细粒度头剪枝与 RNN 转化实现升维打击。
  • vs SANA-Video / ReHyAt: 继承并拓展了混合线性注意力的优势,更进一步实现了 100% 块的循环 RNN 转换与真机 NPU 的图优化落地。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次打破移动端只能跑小模型的固有认知,体系化实现 5B 视频 DiT 端侧运行。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 VBench 定量、消融实验、真机时延显存 Profiling 及大规模双盲人工评测。
  • 写作质量: ⭐⭐⭐⭐⭐ 结构清晰,核心设计公式与系统层级优化逻辑严密自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为未来端侧生成式视频大模型的产品化与芯片部署提供了极其重要的基准范式。