論文鏈接：https://arxiv.org/pdf/2508.21058
項目鏈接：https://primecai.github.io/moc/

亮點直擊

• 提出MoC框架，將每個查詢路由到視頻序列中最相關的片段，而不是依賴統一或靜態的稀疏注意力機制或固定選擇策略。
• 引入了因果路由 mask，以強制交互圖為有向無環圖，提升在分鐘級序列上的展開魯棒性。
• 將密集自注意力替換為MoC機制，能夠將長視頻生成轉化為一種內部上下文檢索過程。MoC 能夠裁剪超過 85% 的 token 對，將注意力 FLOPs 成本最多減少7倍，并在分鐘級場景中實現 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）中常用的注意力模塊，這是當前最先進視頻生成模型的主干。一個注意力模塊定義為：

本工作認為這種均值池化操作對于視頻生成任務來說是高度充分且具有表達力的。首先，均值池化的動機來自于多模態匹配中的既有實踐，例如在 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）。這些擾動在訓練過程中以隨機方式應用，用以模擬噪聲或不完美的路由決策，促使模型學習更具彈性的表示，避免過度依賴特定塊，同時確保上下文空間的均衡利用。

本質上，系統的端到端可微性意味著注意力過程本身就作為路由器的學習信號，使整個框架無需專用路由參數即可實現自我糾正與自適應。

注意力塊劃分與路由

內容對齊的塊劃分  在上下文混合（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）。這類自環會使信息局部化，阻礙梯度傳播，并在雙向生成過程中表現為運動停滯或幀重復。

從經驗上看，因果路由不僅消除了孤立的反饋對，還促進了更豐富的長程依賴，從而帶來更平滑的時間動態和更穩定的訓練過程。

計算效率

與 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 為：

對于相同的L和d ，一個普通的全注意力頭的計算代價為：

它們的比值可簡化為：

其比值隨序列長度線性增長。例如，考慮一個常見的 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% 的上下文，本工作的方法實現了2.2x的加速。此外，它顯著提升了模型的性能，特別是在運動多樣性方面，動態程度（Dynamic-Degree） 從 0.46 提升至 0.56，同時保持了 運動平滑度（Motion-Smoothness）。盡管這種增加的運動預算導致外觀保真度略有下降，但所有質量指標仍然保持在較高水平。總體而言，這些結果驗證了本工作方法的核心理念：通過學習到的、結構感知的稀疏性，將計算從冗余幀重新分配到顯著的視覺事件，從而在不犧牲（甚至在某些情況下提升）感知質量的前提下實現顯著的效率提升。

定性結果。 下圖 3 和下圖 4 分別展示了單鏡頭和多鏡頭視頻生成任務的定性對比。本文認為，這種均值池化操作非常適用于視頻，因為在空間上接近的像素和在時間上相鄰的幀往往描繪的是同一個物體或背景區域。在 DiT 的 patch 嵌入之后，這些 token 占據了一個非常狹窄的子空間：其第一主成分在實際中通常可以解釋超過 90% 的局部方差。算術平均值正是居中數據的第一主成分估計器，因此簡單的平均操作已經能夠捕捉整個塊的主導語義，同時丟棄高頻噪聲。

零樣本實驗支持這一觀點 —— 將這種路由策略直接應用于一個預訓練的視頻生成模型，如本文補充材料中所示。盡管公式3中的路由得分實際上只是查詢向量與均值池化后的鍵向量之間的點積，但它并不是一個固定的啟發式方法：被平均的鍵向量和用于評分的查詢向量都是由訓練中更新的權重生成的。梯度會通過均值池化操作和后續的 top-k 掩碼流向投影矩陣，使模型能夠學習如何塑造每個塊的池化鍵和每個查詢，從而最有效地區分有用與無關的上下文。

在實踐中，這使得表面上“簡單”的均值 + top- k規則具有高度表達能力，而無需引入額外的路由參數或計算開銷，因為網絡會持續地調整其內部表示以加以利用。

MoC 實現基準測試。 在下圖 5 中對自適應 MoC 的性能與全注意力機制（使用 Flash Attention 2 實現）進行了基準測試，其中本工作的方法在計算量（FLOPs）和延遲方面相對于鏡頭數量，或者換句話說，序列長度L ，保持近似線性。

除了稀疏性之外，實現這一高效性的關鍵在于三個設計決策：

(1) 使用實時的 ??segment_reduce?? 池化操作，避免在內存中生成可變長度的塊；

(2) token 按照 head-major 順序組織，以確保在 gather 操作中進行合并的內存訪問；

(3) 整個路由 + 注意力計算被封裝在一個 Flash Attention 的變長調用中，從而保留內核融合并最小化開銷。

結論

自適應上下文混合（Adaptive Mixture of Contexts, MoC）表明，可學習的稀疏注意力路由可以作為一種強大的、數據驅動的記憶檢索引擎發揮作用。本工作的工作可以說是首次展示，通過結合高效且可學習的稀疏路由機制擴展訓練數據，模型能夠發展出復雜的長期記憶方法。該方法以與短視頻生成相當的成本實現分鐘級記憶。

關鍵在于，這種能力的出現并不依賴于顯式的啟發式方法，例如 3D 先驗或固定規則選擇；模型完全通過數據學習哪些歷史上下文是重要的。由于路由是可學習的，且推理過程中的實現非常快速，MoC 為下一代可擴展、可控且負責任的長視頻生成模型提供了藍圖。它證明了去除二次注意力瓶頸不僅僅是提升效率，更是開啟視頻生成中涌現的長期記憶的直接路徑。

局限性與未來工作。 目前，本工作的訓練和測試設置與 LCT 完全相同。然而，MoC 在更長序列上節省計算的能力尚待探索。盡管本工作的方法已經以接近短視頻的成本支持分鐘級上下文，但當前的運行時實現依賴于通用的變長注意力機制和框架級的 gather 操作。

鑒于本文實現了7x的 FLOPs 節省，仍有大量進一步加速的空間，這可以通過軟硬件協同設計實現，例如塊稀疏、感知塊結構的變長注意力機制、更高效的定制 CUDA/Triton 內核、融合的路由+注意力算子、持久化執行、以及改進的鍵/值布局或量化方法。

此外，值得在更廣泛的應用場景中探索 MoC，例如視頻世界模型數據集。

本文轉自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/58lCMBPOF3JEWOnkOhGPHQ??