论文链接：https://arxiv.org/pdf/2508.21058
项目链接：https://primecai.github.io/moc/

亮点直击
提出了一种自适应上下文混合（Adaptive Mixture of Contexts，MoC）框架，该框架学习将每个查询路由到视频序列中最相关的片段，而不是依赖统一或静态的稀疏注意力机制或固定选择策略。
引入了因果路由 mask，以强制交互图为有向无环图，从而防止异常的循环闭合现象，提升在分钟级序列上的展开鲁棒性。
展示了将密集自注意力替换为本工作提出的自适应上下文混合（MoC）机制，能够将长视频生成转化为一种内部上下文检索过程。MoC 的高效性体现在其能够裁剪超过 85% 的 token 对，将注意力 FLOPs 成本最多减少7倍，并在分钟级场景（约 18 万个 token）中实现 2.2 倍的端到端生成速度提升。
MoC 是首个证明学习型稀疏上下文路由能够突破二次注意力机制实际瓶颈的工作，能够以接近短视频的成本提供分钟级的长上下文视频记忆，同时在保真度与一致性方面维持甚至超越密集基线。

总结速览

解决的问题

长视频生成中存在长上下文记忆难题，模型需在长时间范围内有效保留和检索关键事件，同时避免信息崩溃或漂移。
扩散式Transformer（DiTs）因自注意力机制二次计算成本过高，难以扩展到长上下文视频生成，面临内存与计算资源瓶颈，且在长序列上优化困难。

提出的方案

提出一种可学习的稀疏注意力路由模块——上下文混合机制（Mixture of Contexts, MoC），将长视频生成重新定义为内部信息检索任务。
MoC 动态地为每个查询选择少量信息丰富的上下文块（chunk），并强制引入两个锚点（字幕和局部窗口连接）以保障跨模态交互与局部保真度。
引入因果路由掩码，确保交互图为有向无环图，防止循环闭合现象，增强长序列生成的稳定性。
训练中逐步稀疏化路由策略，引导模型聚焦关键历史信息，实现计算资源的有效分配。

应用的技术

采用可学习的 top-k 路由机制，实现 token 到上下文块的动态稀疏选择。
集成Flash Attention内核，支持可变序列长度处理，提高计算吞吐量。
通过分块策略（按帧、镜头和字幕划分）组织多模态 token 流，增强内容对齐与检索效率。
使用因果路由掩码确保生成过程的无环性，提升长序列生成的鲁棒性。

达到的效果

实现近似线性的计算扩展，大幅降低注意力计算成本（FLOPs 最多减少7倍），生成长度达到约18万token的分钟级视频。
在分钟级时间尺度上保持身份、动作与场景的一致性，维持甚至超越密集基线的生成质量。
端到端生成速度提升2.2倍，token裁剪比例超过85%，显著提高了长视频生成的训练与推理效率。
首次验证学习型稀疏上下文路由在实际应用中的有效性，以接近短视频的成本实现长上下文视频记忆。

方法

为了在生成长视频时避免标准自注意力机制的二次计算成本，本文方法用一种自适应、内容对齐的上下文混合（Mixture of Contexts, MoC）层替代了 DiT主干中的密集注意力。从整体上看，MoC 具备以下功能：

将每个查询仅路由到最相关的上下文块；
将这些上下文块与自然的视频边界（如帧、镜头和字幕 token）对齐；
强制因果性，确保信息沿时间方向单向流动。

下面将详细介绍路由公式、交错文本生成视频中的分块与选择策略、计算效率。本工作方法的整体流程如下图 1 所示。

上下文混合

扩散 Transformer 中的原始注意力机制。 首先回顾扩散 Transformer（DiT）中常用的注意力模块，这是当前最先进视频生成模型的主干。一个注意力模块定义为：

其中、和分别表示查询、键和值特征，表示特征维度。注意，当本工作将视为一组独立向量时，公式 (1) 可写为：该公式以查询为单位进行计算。

通过 Top-k 选择实现动态路由。 在视频 DiT 中，对于一段 480p、时长 1 分钟的视频，其序列长度容易扩展至接近 200k。这使得自注意力的计算开销变得极其昂贵。

由于特征存在冗余，常见做法是将视频序列划分为多个块，使得一个查询 token 仅与这些块的一个子集进行交互。自回归视频生成方法，如[3, 4, 48] 通常按帧将上下文划分为块，其中查询仅关注最近的几个块，从而丢失了更远距离的上下文信息。

相反，本工作采用一种可学习的路由策略，使每个被路由到最相关的块。

其中表示一组路由后的索引，是查询所关注的所有上下文位置的索引。给定所有块的列表，对于每个，仅通过 top-k 操作选取少量块用于注意力计算。

其中表示连接所有 top-k 块的索引。查询与块序列之间的相关性由与的描述符的内积决定。本工作使用简单、高效且有效的均值池化操作作为描述符变换。

本工作认为这种均值池化操作对于视频生成任务来说是高度充分且具有表达力的。首先，均值池化的动机来自于多模态匹配中的既有实践，例如在 CLIP 中，通过对 token 嵌入进行平均来获得全局图像表示，以与文本嵌入计算相似度。这种方法能够有效捕捉主导语义特征，同时对局部变化具有鲁棒性，这一特性自然适用于视频块，其中空间和时间上相邻的 token 通常表示冗余或相关的视觉元素（例如静态背景或渐变运动）。

在本工作的可训练框架中，均值池化并非静态启发式，而是一种自适应机制：尽管 top-k 本身是不可微的，模型仍可通过对选中块的注意力机制间接学习。具体而言，如果某个被选中的块在注意力计算中证明是无关的，损失函数的梯度将通过其 key/value 反向传播，而这些 key/value 是均值池化描述符的来源。该过程会削弱无效表示，并促使查询/key 投影在训练迭代中产生更具判别性的相似度。

这一自我修正过程与硬路由 MoE 系统和稀疏注意力框架中的间接适应机制一致（例如在离散且不可微选择的情况下，由下游模块提供学习信号）。这种端到端可微性和无参数的路由器确保了看似简单的点积路由变得高度具有表达力，因为网络会塑造嵌入以强化稀疏注意力中的判别性特征，同时不引入额外参数或计算开销。对预训练模型的零样本应用进一步验证了其有效性。

上下文丢弃与注入。 为增强本工作的上下文混合（Mixture of Contexts，MoC）机制的鲁棒性，并缓解类似于混合专家系统（Mixture-of-Experts，MoE）中的“死专家”问题——即某些专家或路径在训练过程中被低估或忽略——本工作引入了两种互补的正则化技术：上下文丢弃（context drop off）和上下文注入（context drop in）。这些扰动在训练过程中以随机方式应用，用以模拟噪声或不完美的路由决策，促使模型学习更具弹性的表示，避免过度依赖特定块，同时确保上下文空间的均衡利用。

受到以下观察的启发：由于嵌入中的噪声或数据分布的演化，路由可能出现不准确性，context drop off 会随机移除每个查询 token 所选 top-k 块中的一部分。具体而言，对于给定的查询，在根据公式 3 计算出路由索引后，本工作从中采样一个丢弃概率，并从中随机屏蔽个块。这迫使模型即使在部分上下文偶尔不可用的情况下也能生成连贯输出，从而在学习的依赖关系中引入冗余，并防止因路由错误导致的灾难性失败。

相反，context drop in 会向所选集合中注入额外的块，以模拟过度包含的路由行为。对于每个查询，本工作从中随机采样个块，将其加入到选中的池中。该技术通过人为激活被低估的块来对抗“死路由”问题，确保梯度能够流经更广泛的上下文片段，并随着时间推移平衡路由分布。由于本工作的路由器是无参数的，仅依赖于均值池化的特征相似度，这些正则化技术不会干扰路由机制本身的学习。如果某个块确实重要，其相关性将通过注意力模块中的反向传播自然增强，因为模型会调整查询与 key 的投影，以放大有意义的相似性并弱化无关特征。

本质上，系统的端到端可微性意味着注意力过程本身就作为路由器的学习信号，使整个框架无需专用路由参数即可实现自我纠正与自适应。

注意力块划分与路由

内容对齐的块划分 在上下文混合（Mixture of Contexts）中，一个关键但常被忽视的设计维度是如何将庞大的 token 流划分为候选块。在长上下文的大语言模型（LLMs）中，这一决策相对简单：输入是一个同质的一维子词 token 序列，配有单一的 RoPE，因此将其切分为固定长度的窗口（如 MoBA 中所做）既能保持局部语义一致性，也能匹配单调的位置信度量。

相比之下，视频生成的 DiT 模型通常是多模态的，运行在异构的 3D+模态网格上：一个展平的顺序，交错排列了空间 patch、时间帧和文本 token，这些元素具有各自独立的 3D RoPE 因子。因此，相邻索引可能在时空上相距甚远，或跨越一个突兀的镜头切换，而一个静态背景 patch 可能在数百帧中重复出现在一个高度熵变的运动 token 附近。统一的窗口会模糊这些差异信号，污染用于公式 3 中的均值池化 key，迫使 top-k 选择器浪费在内部不一致的 key 上。

本工作改为沿内容感知的边界——帧、镜头和模态条带——对序列进行划分，使每个块在语义上是同质的，并在 3D 位置流形中具有几何上的局部性。这种对齐方式保留了公式 3 中均值池化 key 的判别力，带来更具信息量的 top-k 检索，并在不牺牲长距离一致性的前提下大幅降低二次计算开销。这种块划分策略不仅适用于现有的单镜头文本生成视频模型，也兼容现有的长视频生成方法，该方法直接在交错的文本-视频对构成的超长序列上计算注意力。

固定的跨模态选择 除了动态路由的视觉块外，本工作明确要求每个视觉查询 token 必须关注序列中的所有文本 token。尽管文本 token 通常只占所有 token 的不到 1%，但它们编码了最具语义信息的信号——指定了全局风格、角色身份和关键动作。计算开销可以忽略不计，但带来的收益却非常显著：将生成锚定在提示上显著减少了提示漂移错误，并防止在长视频生成过程中稀有属性词语的消失。

此外，这种强制的跨模态连接促进了文本与视觉嵌入的联合梯度传播，收紧它们的共享潜在空间，并在下游任务（如文本引导的视频编辑）中显著提升可编辑性。

固定的镜头内选择 长视频自然呈现出严格的层次结构：帧嵌套在镜头中，镜头嵌套在场景中。为了利用这一点，本工作在注意力机制中显式地强制镜头内连接，确保每个 token 始终关注其所属的镜头，从而捕捉物体轨迹、光照连续性以及其他可预测的线索。

这种设计使上下文混合（Mixture of Contexts, MoC）框架能够将其稀疏注意力预算分配给真正的长程依赖，而不是冗余地建模局部上下文。强制这种连接带来了若干好处：它防止了在场景切换处出现语义不连续的情况（此时相邻 token 可能变得无关）；它保证每个注意力矩阵至少包含一个良好条件的块；即使在激进的自适应剪枝下，它也提供了一条连续且内存高效的回退路径。

这一策略在微调预训练的视频生成模型时尤为有效，因为它从一开始就保留了每个镜头的保真度，并使模型能够在训练过程中逐步学习对齐更广泛的上下文信息。

稀疏 MoC 中的因果性 稀疏路由本质上在 token 交互图中引入了方向性，因为每个块只能选择有限的其他块进行注意力。然而，在缺乏显式排序约束的情况下，这一过程可能退化为病态的闭环。例如，在消融实验中，当每个块仅允许选择一个对等块时，本工作经常观察到这样的情况：块 5 路由到块 6，而块 6 同时又路由回块 5，形成一个孤立的双节点循环（见下图 2）。这类自环会使信息局部化，阻碍梯度传播，并在双向生成过程中表现为运动停滞或帧重复。

为了解决这一问题，本工作在路由阶段引入了因果掩码，限制每个块仅能关注序列中较早位置的 key；具体而言，任何满足的边都会在 top-k 选择之前被掩蔽掉。该约束将路由图转化为有向无环图（DAG），确保信息严格沿时间向前流动，并在结构上排除了闭环的可能性。

从经验上看，因果路由不仅消除了孤立的反馈对，还促进了更丰富的长程依赖，从而带来更平滑的时间动态和更稳定的训练过程。

计算效率

与 Flash-Attention 内核的结合。 处理内容对齐且块大小高度不均的情况远比像 MoBA 和 NSA 中那种平均划分的设置复杂得多。为了在保持效率的同时兼顾帧、镜头和模态结构，本工作实现了一种完全在 GPU 上运行的自适应注意力机制，并显式利用了视频 DiTs 中的结构线索。

本工作首先为展平后的 token 流打上帧、镜头和文本边界标签，并使用 torch.bucketize 和前缀和表（如 cu_seqlen、cu_shot 等）来派生内容对齐的可变长度块，这些块的起止索引与边界一致，从而确保每个块在语义上是同质的。边界信息还用于构建预路由掩码：在 top-k 稀疏化步骤之前插入强制连接（例如文本-视觉、镜头内自连接），以确保路由器不会在已经是强制连接的块上浪费预算。

对于每个保留的块，本工作通过 segment_reduce 的均值池化在线获取一个代表性 key，从而避免对整个块进行物化，即使块大小相差数个数量级，也能保持内存平坦。token 按照 head-major 顺序收集（通过 rearrange(..., 's x h d → h s x d')），使得随后的 gather 操作能够合并，异构的（query, key）对被打包成一次 Flash-Attention 的变长调用。

该设计生成了一个注意力内核，既能遵循视频特定的结构约束，又能在处理数百万 token 时保持内存与计算效率。由于所有操作都是 head 无关的，本工作可以充分利用张量并行性，并将计算在多个设备间分片。

节省的 FLOPs。 对于每个注意力头，设为序列长度或查询 token 数量，为内容对齐的块数，为每个查询 token 保留的 top-k 块数，为所选块的平均长度，为每个注意力头的维度。

每个块内的 key 做均值池化仅需次加法，几乎可以忽略。
路由阶段对每个查询–块对计算一次内积，代价为 FLOPs（乘以 2 是因为每个维度一次乘法一次加法）。
最后，在剪枝后的集合上进行精细注意力计算，对每个查询 token 最多涉及个 key，代价约为 FLOPs。

三项相加，总 FLOPs 为：

对于相同的和，一个普通的全注意力头的计算代价为：

它们的比值可简化为：

其比值随序列长度线性增长。例如，考虑一个常见的 VAE 压缩率（空间下采样 16×，时间下采样 4×），一个分辨率为 480P、帧率为 12fps、时长为 1 分钟的视频将变成一个大约包含 180k 个 token 的序列。

假设本工作使用，，，，则可计算出：

相比之下，在相同序列上进行稠密自注意力的计算代价为：

因此，自适应上下文混合（Mixture of Contexts）层将乘加操作减少了超过 7 倍。

实验

本工作主要在两个任务上进行实验：镜头级文本生成视频任务，以及包含多个镜头切换的场景级文本生成视频任务。

基础模型。 本工作基于长上下文视频生成器 LCT 构建模型，该模型支持单镜头和多镜头视频生成。LCT 采用了一个拥有 30 亿参数的 MMDiT 架构，在图像、单镜头视频和多镜头视频的原始分辨率和时长上进行训练。该模型的全自注意力从每个镜头的范围扩展到最多包含八个镜头的场景级上下文窗口（大约 8 秒，每个镜头约 22k 个 token），使用交错的三维 RoPE 为每个镜头提供独特的绝对坐标，同时保持文本与视频 token 的相对布局。本工作从预训练的 LCT 初始化模型权重，并将其注意力模块替换为本工作的 MoC，然后使用与 LCT 相同的训练方案进行微调。

基线方法。 对于单镜头视频生成，本工作与原生的 30 亿参数 MMDiT 视频生成模型进行比较，该模型是 LCT 和本工作工作的基础。本工作在 8 秒、分辨率为 320×192、帧率为 12FPS 的视频上进行测试，每个视频大约包含 6300 个 token。对于多镜头视频生成，本工作与本工作的基础模型 LCT 进行比较。在这些实验中，本工作测试包含 8 个镜头的序列，每个镜头是一个 8 秒、480p、12FPS 的视频，整个 64 秒的场景大约包含 180k 个 token。

评估指标。 对于短的单镜头视频生成和长的多镜头视频生成，本工作遵循以往工作，在广泛使用的 VBench 基准上进行评估。具体而言，主体一致性（Subject Consistency） 和 背景一致性（Background Consistency） 衡量视频中是否忠实保留了输入图像中的主要主体和背景；运动平滑度（Motion Smoothness） 评估运动的流畅性（是否存在抖动或突变）；动态程度（Dynamic Degree） 衡量视频中的运动程度（鼓励生成动态内容而非静态场景）。本工作还报告了单镜头视频的 美学质量（Aesthetic Quality） 和 图像质量（Image Quality），以量化每帧的视觉吸引力和技术质量。此外，本工作还报告了计算指标，如稀疏度、FLOPs 和相较于 Flash Attention 的推理加速比。

定量结果。 下表 1 和下表 2 展示了本工作提出的内容对齐上下文混合（Mixture of Contexts, MoC）模型与稠密注意力基线在短单镜头片段和长多镜头场景上的定量对比。对于短的单镜头视频（6k 个 token），尽管进行了激进的稀疏化，本工作的方法在所有 VBench 指标上均与稠密基线持平或更优。这表明，将计算资源引导至最相关的块不仅减少了 FLOPs，还使模型更好地保持了角色的一致性和场景的连贯性。然而，对于如此短的序列，索引收集和池化所带来的额外开销超过了计算节省，导致端到端流程变慢。

相反，对于长的多镜头视频（180k 个 token），MoC 显示出明显的计算优势。通过丢弃 85% 的上下文，本工作的方法实现了的加速。此外，它显著提升了模型的性能，特别是在运动多样性方面，动态程度（Dynamic-Degree） 从 0.46 提升至 0.56，同时保持了 运动平滑度（Motion-Smoothness）。尽管这种增加的运动预算导致外观保真度略有下降，但所有质量指标仍然保持在较高水平。总体而言，这些结果验证了本工作方法的核心理念：通过学习到的、结构感知的稀疏性，将计算从冗余帧重新分配到显著的视觉事件，从而在不牺牲（甚至在某些情况下提升）感知质量的前提下实现显著的效率提升。

定性结果。 下图 3 和下图 4 分别展示了单镜头和多镜头视频生成任务的定性对比。本文认为，这种均值池化操作非常适用于视频，因为在空间上接近的像素和在时间上相邻的帧往往描绘的是同一个物体或背景区域。在 DiT 的 patch 嵌入之后，这些 token 占据了一个非常狭窄的子空间：其第一主成分在实际中通常可以解释超过 90% 的局部方差。算术平均值正是居中数据的第一主成分估计器，因此简单的平均操作已经能够捕捉整个块的主导语义，同时丢弃高频噪声。

零样本实验支持这一观点 —— 将这种路由策略直接应用于一个预训练的视频生成模型，如本文补充材料中所示。尽管公式中的路由得分实际上只是查询向量与均值池化后的键向量之间的点积，但它并不是一个固定的启发式方法：被平均的键向量和用于评分的查询向量都是由训练中更新的权重生成的。梯度会通过均值池化操作和后续的 top- 掩码流向投影矩阵，使模型能够学习如何塑造每个块的池化键和每个查询，从而最有效地区分有用与无关的上下文。

在实践中，这使得表面上“简单”的均值 + top- 规则具有高度表达能力，而无需引入额外的路由参数或计算开销，因为网络会持续地调整其内部表示以加以利用。

MoC 实现基准测试。 在下图 5 中对自适应 MoC 的性能与全注意力机制（使用 Flash Attention 2 实现）进行了基准测试，其中本工作的方法在计算量（FLOPs）和延迟方面相对于镜头数量，或者换句话说，序列长度，保持近似线性。

除了稀疏性之外，实现这一高效性的关键在于三个设计决策：

(1) 使用实时的 segment_reduce 池化操作，避免在内存中生成可变长度的块；

(2) token 按照 head-major 顺序组织，以确保在 gather 操作中进行合并的内存访问；

(3) 整个路由 + 注意力计算被封装在一个 Flash Attention 的变长调用中，从而保留内核融合并最小化开销。

结论

自适应上下文混合（Adaptive Mixture of Contexts, MoC）表明，可学习的稀疏注意力路由可以作为一种强大的、数据驱动的记忆检索引擎发挥作用。本工作的工作可以说是首次展示，通过结合高效且可学习的稀疏路由机制扩展训练数据，模型能够发展出复杂的长期记忆方法。该方法以与短视频生成相当的成本实现分钟级记忆。

关键在于，这种能力的出现并不依赖于显式的启发式方法，例如 3D 先验或固定规则选择；模型完全通过数据学习哪些历史上下文是重要的。由于路由是可学习的，且推理过程中的实现非常快速，MoC 为下一代可扩展、可控且负责任的长视频生成模型提供了蓝图。它证明了去除二次注意力瓶颈不仅仅是提升效率，更是开启视频生成中涌现的长期记忆的直接路径。

局限性与未来工作。 目前，本工作的训练和测试设置与 LCT 完全相同。然而，MoC 在更长序列上节省计算的能力尚待探索。尽管本工作的方法已经以接近短视频的成本支持分钟级上下文，但当前的运行时实现依赖于通用的变长注意力机制和框架级的 gather 操作。

鉴于本文实现了的 FLOPs 节省，仍有大量进一步加速的空间，这可以通过软硬件协同设计实现，例如块稀疏、感知块结构的变长注意力机制、更高效的定制 CUDA/Triton 内核、融合的路由+注意力算子、持久化执行、以及改进的键/值布局或量化方法。

此外，值得在更广泛的应用场景中探索 MoC，例如视频世界模型数据集。