VideoMLA:用于分钟级自回归视频扩散的低秩潜在KV缓存 VideoMLA: Low-Rank Latent KV Cache for Minute-Scale Autoregressive Video Diffusion
首个将MLA应用于视频扩散的工作,压缩KV缓存92.7%同时保持长时生成质量
前置知识
KV Cache
键值缓存是自注意力机制中用于存储先前计算结果的数据结构,包含键和值矩阵。在自回归生成中,缓存过去的token的K和V可以避免重复计算,显著提升推理效率。在视频生成中,由于序列长度很长(分钟级视频可能包含数百帧),KV缓存会占用大量内存,成为主要瓶颈。
本文的核心创新就是压缩KV缓存,理解KV缓存的结构和内存占用是理解VideoMLA的基础
MLA (Multi-Head Latent Attention)
多头潜在注意力是DeepSeek-V2提出的注意力机制变体,用共享的低秩潜在表示替代每个注意力头的独立键值。核心思想是通过降维投影将键值压缩到低维潜在空间,使用时再通过上投影恢复到各头的维度。这样可以大幅减少内存占用,同时保持表达能力。
VideoMLA是首个将MLA应用于视频扩散的工作,理解MLA的基本原理和与标准MHA的区别是读懂本文的关键
RoPE (Rotary Positional Embedding)
旋转位置编码是一种相对位置编码方法,通过将查询和键向量在复平面旋转来编码位置信息。3D-RoPE将其扩展到视频的三个维度(时间、高度、宽度),分别应用旋转编码。VideoMLA使用解耦的3D-RoPE,将位置信息单独存储和编码,与内容潜在表示分离。
VideoMLA的一个重要创新是使用head共享的解耦3D-RoPE键,这需要理解RoPE的工作原理
Causal Video Diffusion
因果视频扩散是用于流式生成长视频的方法,通过将双向教师模型蒸馏成因果学生模型,实现自回归逐帧或逐块生成。学生模型维护一个滚动的KV缓存,只关注过去的帧,这样可以生成长视频而不需要等待所有帧完成。缓存大小通常固定为滑动窗口,限制了可用的上下文长度。
本文针对的场景就是因果视频扩散的KV缓存问题,需要理解这个问题的背景和挑战
有效秩 (Effective Rank)
有效秩是衡量矩阵低秩程度的指标,定义为捕获特定比例(如99%)能量所需的最小秩。通过奇异值分解可以计算累积能量,有效秩反映了矩阵的低秩近似能力。在注意力机制中,如果WK和WV的联合矩阵有效秩较低,说明可以用低秩表示而损失较少信息,这是MLA有效的理论基础。
本文的一个重要发现是视频注意力不是低秩的,这与MLA在语言模型中的动机形成对比,理解有效秩的概念是理解这个发现的关键
研究动机
长序列因果视频扩散模型的KV缓存内存占用巨大,成为生成更长视频的主要瓶颈。以Wan-1.3B模型为例,每个缓存的token存储密集KV标量。当缓存21个潜在帧、每帧1560个token、30层transformer时,密集KV缓存包含约3.02B个标量,在bf16/fp16精度下约占6.0GB内存。为了控制内存,现有系统使用固定大小的滑动窗口缓存,但这只是限制了缓存的token数量,没有减少每个token、每层的per-head KV存储开销。随着模型扩展到更长序列,per-head KV缓存越来越成为决定性的性能瓶颈。
本文的目标是本文的目标是在保持分钟级视频生成质量的同时,大幅减少per-token、per-layer的KV缓存内存占用。具体来说,通过将MLA引入视频扩散领域,用共享的低秩潜在表示替代密集的per-head键值对,将每个token的缓存状态从密集存储压缩到潜在表示加上位置编码,实现高达92.7%的内存减少。同时要理解为什么MLA在视频扩散中有效,尽管预训练的注意力权重的谱结构并不满足低秩假设。
与已有工作不同的是,现有工作主要在固定大小的滑动窗口内进行创新,如改变哪些token占用窗口、如何编码位置、如何重新分配或压缩token,但这些方法都保留了密集的per-head KV布局本身。另一条工作线改变注意力计算本身,如使用线性注意力替代softmax注意力,或分离时序推理和帧级渲染以减少缓存层数。VideoMLA采用了不同的角度:它保持所有30层自注意力层缓存,但将每个token的缓存状态大幅减少,直接针对per-token、每层的KV布局本身。与语言模型中MLA通过预训练低秩结构来动机不同,本文发现视频注意力不是低秩的,提出了一个新的视角:MLA的瓶颈而非预训练谱决定有效秩。
核心方法
VideoMLA的核心思想是用共享的低秩潜在表示替代密集的per-head键值对,同时将位置信息解耦到单独的RoPE子空间。整体思路分为三个部分:压缩KV缓存构造、解耦3D-RoPE、训练时前向传播。首先,通过共享的KV下投影将视频潜在token压缩到低维潜在空间,这个潜在表示就是写入缓存的内容。位置信息不折叠到潜在中,而是通过解耦的3D-RoPE键单独存储。使用时,通过上投影从内容潜在恢复各头的键和值,再与旋转的位置键结合计算注意力。这样每个缓存的token只需要存储内容和位置编码而非密集的per-head KV状态,大幅减少内存占用。
VideoMLA的核心创新是首次将MLA应用于视频扩散,并提出head共享的解耦3D-RoPE。与标准MHA中每个头都有独立的键值不同,VideoMLA使用一个共享的内容潜在,通过上投影恢复各头的键值,实现跨head共享内容表示。同时,将位置信息与内容分离,使用一个head共享的3D-RoPE键,旋转后供所有头使用。这种设计的关键洞察是:内容表示可以跨head共享和压缩,而位置信息需要保持高保真度,因此用分离的RoPE通道处理。另一个关键洞察是MLA的有效性不依赖于预训练权重的低秩结构,而是通过架构瓶颈设定的秩预算,训练时模型会充分利用并适应这个预算。
方法步骤详情
VideoMLA的完整方法包含三个主要步骤。第一步是压缩KV缓存构造:对于每个视频潜在token,通过共享的KV下投影压缩到内容潜在空间,这是写入滚动缓存的内容对象。然后通过上投影恢复各头的键和值。第二步是解耦3D-RoPE:将每个头的维度分为NoPE部分和RoPE部分,NoPE部分从内容潜在重建且不使用旋转位置编码,RoPE部分形成解耦的3D-RoPE键。对于每个token,计算单个head共享的位置键,然后在组装注意力窗口时旋转。查询路径类似,从查询潜在得到内容查询和位置查询。第三步是训练时前向传播:对于查询token位置和缓存token位置,注意力头结合内容和位置贡献计算分数,然后在活动注意力窗口上做softmax后与值做加权求和得到每头输出,最后通过输出投影混合各头输出。
技术新颖性
VideoMLA的技术新颖性体现在多个方面。首先,这是首次将MLA应用于视频扩散的工作,将MLA从语言模型推广到视频生成领域,处理了视频特有的时序和空间结构。其次,提出了head共享的解耦3D-RoPE设计,与标准MLA中每头独立的RoPE不同,VideoMLA使用一个head共享的位置键,进一步减少内存占用。第三,通过实验发现了一个违反直觉的现象:预训练视频注意力不是低秩的,但MLA在这个压缩比例下仍保持高质量。这表明MLA的有效性不是来自恢复预训练权重的低秩结构,而是通过架构瓶颈设定的秩预算,训练时模型会充分利用这个预算并在其中适应。这个发现挑战了MLA的常见动机,提出了一个新的理解视角。第四,系统性地研究了潜在维度和RoPE维度的分配对视觉保真度和运动一致性的影响,识别出适合分钟级视界的分配策略。
实验结果
实验结果表明VideoMLA在保持长时生成质量的同时实现了显著的内存压缩。在VBench长时视频生成评估中,VideoMLA在30秒和60秒视界下都取得了最佳的动态度得分,分别为0.981和0.958,表明KV压缩不会抑制运动或导致静态生成。VideoMLA还取得了最佳的成像质量和运动平滑度,在60秒视界下达到最高总体得分0.859,显著优于Reward Forcing、Infinity-RoPE、LongLive和LongSANA等基线方法。在30秒视界下,VideoMLA也具有竞争力,在使用更小KV缓存内存的情况下取得了第二佳的总体得分。效率方面,在单块B200 GPU上,VideoMLA达到了23.96 FPS的吞吐量和3.38秒的延迟,是chunk-wise自回归模型中吞吐量最高、延迟最低的方法。CLIP-T得分为0.3278也是最佳,表明与文本的对齐性更好。定性结果显示VideoMLA在30秒滚动中保持主体身份、场景结构和视觉保真度,与代表性流式因果视频基线相比具有可比的质量,同时需要更快的推理和显著更低的内存。消融实验表明,在固定B200内存预算下,密集MHA在batch size为28时达到内存限制,而MLA将OOM悬崖推向更右侧。per-request内存斜率从MHA的6.26 GB/batch减少到MLA的0.57-1.43 GB/batch,实现77-91%的减少。因此,MLA在相同内存上限下支持4.6倍到至少11.4倍更大的非OOM批次,默认设置给出8.0倍批次余量。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 60秒长时视频生成 (VBench) | Overall Score | 0.859 | Reward Forcing (0.808) | 提升6.3% |
| 60秒长时视频生成 (VBench) | Dynamic Degree | 0.958 | Reward Forcing (0.676) | 提升41.7% |
| 60秒长时视频生成 (VBench) | Imaging Quality | 0.715 | Infinity-RoPE (0.638) | 提升12.1% |
| 单GPU吞吐量 (B200) | Throughput (FPS) | 23.96 | Self-Forcing (18.06) | 提升32.7% |
| 单GPU延迟 (B200) | Latency (seconds) | 3.38 | Self-Forcing (4.19) | 降低19.3% |
| 文本对齐 (CLIP-T) | Score | 0.3278 | LongSANA (0.2978) | 提升10.1% |
| 内存占用 (per-token per-layer) | Scalars | 224 (默认设置) | Dense MHA (3072) | 减少92.7% |
| 用户研究 | Overall | 3.17 | Reward Forcing (2.91) | 提升8.9% |
局限与改进
作者指出的局限性包括:VideoMLA减少了per-token KV缓存,但潜在预算不能任意缩小。小预算如较低维度可以提高内存余量但会丢失细粒度细节并降低质量,使得潜在维度成为一个质量-效率权衡参数。实验主要关注Wan2.1-T2V-1.3B和分钟级生成,更大的backbone、更高分辨率、更长视界和提示切换仍然是未来工作。此外,从观察到的局限性来看,VideoMLA在解耦位置编码时假设了时间、高度和宽度维度可以用相同频率的旋转编码处理,但不同维度可能需要不同的编码策略。内容维度和RoPE维度的分配是经验选择的,可能不是最优的。MLA的有效性依赖于架构瓶颈设定的秩预算,但如何自动确定最优的潜在维度仍不清楚。训练过程中MLA几乎完全占用了秩预算,这意味着可能存在过拟合或利用所有可用容量的倾向,这可能影响泛化性。最后,虽然VideoMLA减少了内存占用,但推理时的上投影会增加计算开销,在高分辨率场景下这个开销可能更加明显。
独立分析的弱点
VideoMLA的一个潜在弱点是在极低秩预算下质量下降明显。在较低维度时,虽然内存减少更多,但会丢失细粒度视觉细节,导致生成质量下降。这表明MLA在视频中的质量-效率权衡比在语言中更敏感。改进方向可以是自适应压缩策略,根据内容复杂度动态调整压缩率,或者使用层次化压缩,对重要帧或区域使用更小压缩率。另一个弱点是RoPE的固定频率分配,当前使用最高频带在时间、高度、宽度维度分配复数对,但不同视频内容可能需要不同的频率分布。改进方向可以是学习频率分配或使用自适应RoPE。第三个弱点是上投影的计算开销,在高分辨率和大量头的情况下,从潜在维度恢复到各头维度的上投影可能成为计算瓶颈。改进方向可以使用分组上投影或稀疏上投影来减少计算。第四个弱点是训练时秩预算的饱和使用,可能导致过拟合或泛化性问题。改进方向可以引入正则化鼓励稀疏使用或dropout来防止过度依赖秩预算。最后,VideoMLA主要针对静态prompt的长时生成,对prompt切换的支持可能有限。改进方向可以探索prompt感知的压缩策略或缓存复用机制。
未来方向
作者提出的未来工作方向包括:探索更大的backbone、更高分辨率、更长视界和提示切换场景。基于本文成果,可以延伸的研究方向包括:自适应秩预算选择,开发根据视频内容复杂度自动确定最优潜在维度的方法;分层压缩策略,对不同层或不同时间尺度使用不同的压缩率;跨模态MLA,探索将MLA应用于其他模态如图像生成、音频生成等;混合MLA架构,结合MLA和标准MHA,对关键层使用MHA而对其他层使用MLA;训练效率优化,研究如何更高效地训练MLA模型,如渐进式压缩或知识蒸馏;量化与MLA结合,探索在压缩KV缓存的同时使用量化进一步减少内存;分布式MLA缓存,研究如何在大规模分布式系统中部署MLA缓存;理论分析,进一步分析MLA在视频中的理论保证和收敛性质;增量学习,研究如何在新数据上微调MLA模型而不破坏秩预算的有效利用;评估指标扩展,开发更适合长时视频生成的评估指标,如运动多样性、场景一致性等。
复现评估
复现VideoMLA需要一定的资源和门槛。作者在8块B200 GPU上运行所有训练实验,使用bf16混合精度训练。训练遵循三阶段Causal Forcing流水线,包括Teacher Forcing、一致性蒸馏初始化到四步、和DMD,总batch size为128。学习率使用特定值用于不同阶段。数据集方面,一致性蒸馏阶段使用47680个视频,来自OpenVid-1M和合成片段。实现基于Wan-2.1 T2V-1.3B backbone,只替换自注意力层而保持其余架构不变。模型有30个transformer块,隐藏维度1536,12个头,每头维度128。默认设置使用特定的潜在维度和查询维度,头维度分为内容部分和RoPE部分。解耦3D-RoPE通道在时间、高度、宽度轴分配特定数量的复数对,使用最高频带。这给出每个token的缓存大小对应从密集标量KV缓存的13.7倍减少。虽然作者提供了详细实现细节,但代码是否开源未明确说明,需要联系作者获取实现细节。复现的难度主要在于需要8块B200的计算资源和大规模视频数据集,以及实现MLA和3D-RoPE的技术复杂性。
论文图表