完整復現Sora，Open-Sora最新技術報告發布，提供完整訓練代碼、權重及數據處理工具 精華
文章鏈接：https://arxiv.org/pdf/2412.20404
項目鏈接：https://github.com/hpcaitech/Open-Sora
總結速覽
解決的問題
人工視覺智能，特別是生成和模擬我們所見世界的能力，相較于語言能力的突破仍然滯后。現有視頻生成模型在高保真視頻內容生成、靈活視頻合成及長時間視頻生成方面面臨諸多挑戰。
提出的方案
本文引入 Open-Sora，一個開源的視頻生成模型，支持文本生成圖像、文本生成視頻以及圖像生成視頻等多種視覺生成任務。通過空間-時間擴散Transformer (Spatial-Temporal Diffusion Transformer, STDiT) 框架，將空間與時間的注意力機制解耦，同時采用高度壓縮的3D自編碼器以壓縮表示，加速訓練過程。此外，提供完整的訓練代碼、模型權重及數據處理工具，推動社區發展。
應用的技術
- STDiT框架：高效的擴散視頻生成框架，解耦空間和時間注意力。
- 3D自編碼器：實現表示的高度壓縮，加速訓練。
- 定制化訓練策略：優化生成效率和效果。
達到的效果
- 支持生成最長15秒、分辨率最高720p的視頻，并適配任意寬高比。
- 在文本生成視頻、圖像生成視頻任務中實現可控的運動動態生成。
- 開源以來取得顯著成果，模型版本持續更新（v1.0到v1.2），當前最新版本（v1.2）已實現完整復現 OpenAI Sora 的技術，并支持生成多分辨率視頻。
數據
數據來源
所使用的數據集全部開源，以確保模型訓練的完全可復現性。總計生成了 30M 個視頻片段，時長從 2秒到16秒 不等，總時長達 80k小時。
- Webvid-10M：包含10M個來自庫存視頻網站的視頻-文本對。視頻為低分辨率并帶有水印。
- Panda-70M：一個大規模數據集，包含70M個視頻-字幕對。使用了其中20M高質量子集進行訓練。
- HD-VG-130M：由130M個文本-視頻對組成，字幕通過BLIP-2生成。發現其場景和文本質量相對較差。
- MiraData：一個高質量數據集，包含77k個長視頻，主要來源于游戲和城市探索。
- Vript：一個密集標注的數據集，包含400k個視頻。
- Inter4K：一個包含1k個4K分辨率視頻片段的數據集。
此外，還從 Pexels、Pixabay 和 Mixkit 獲取了免費授權的視頻。這些網站上的大部分視頻質量較高，對這些優秀平臺及其貢獻者表示由衷的感謝。
圖像數據集與視頻一起訓練，總計包含約 3M 張圖像：
- LAION：一個大規模開放數據集，使用了美學評分大于6.5的子集。
- Unsplash-lite：包含25k張自然主題的Unsplash照片，覆蓋了廣泛的使用場景和上下文。
數據預處理
高質量數據對于訓練優秀的生成模型至關重要。為此建立了一條完整的數據處理pipeline，可將原始視頻無縫轉換為高質量的視頻-文本對。pipeline如圖2所示。
數據處理步驟：
- 場景檢測與視頻剪輯: 使用PySceneCut檢測場景并將視頻剪輯為多個片段。
- 高質量視頻過濾: 主要遵循SVD 數據預處理pipeline
- 美學評分 (Aesthetic Score)：衡量視頻幀的美學偏好。使用來自LAION的評分器，并對每段視頻抽樣的三幀計算平均分。
- 光流評分 (Optical Flow Score)：衡量視頻的動態變化程度，用于篩選低運動量的視頻。此任務使用UniMatch模型完成。
- 文本場景過濾：部分視頻包含大量密集文字場景（如新聞廣播、廣告），不適合訓練。使用OCR（光學字符識別）檢測視頻中的文本，含有過多文字的視頻將被移除。OCR 使用DBNet++模型，由MMOCR實現。
- 視頻字幕生成: 為視頻提供高質量的字幕
- 使用GPT-4V 和 PLLaVA 生成字幕。前者通過 API 提供服務，后者為開源模型，可在本地部署。盡管生成結果存在一定程度的虛構現象，但足以滿足訓練文本生成視頻模型的需求。
- 在實踐中，使用預訓練的PLLaVA 13B 模型，并從每段視頻中選取 4幀 進行字幕生成，采用 2×2 空間池化 形狀。
- 字幕生成模型難以提供關于相機運動的信息，通過光流檢測相機運動并將其信息附加到字幕中。
最后階段所用視頻數據的統計信息如下圖4所示。展示了視頻時長和分辨率的基本統計數據，以及美學評分和光流評分的分布。此外，還從視頻字幕中提取了與對象和動作相關的標簽，并統計了其頻率。
模型架構
本文的視頻生成框架遵循 Sora 報告 的設計。視頻首先通過視頻壓縮網絡（即 3D 自編碼器）進行壓縮，文本通過文本編碼器編碼。然后，一個類似 DiT 的 Transformer 處理視頻和文本的隱空間變量。
3D 自編碼器
在 Open-Sora 1.0 和 1.1 中，使用了 Stability-AI 的 2D VAE（84M 參數），該模型以 8×8 的比例對視頻進行空間壓縮。為減少時間維度，通過每三幀提取一幀進行下采樣。然而，這種方法由于生成的幀率（FPS）降低，導致時間流暢性較差。
為解決這一限制，Open-Sora 1.2 引入了一個受 OpenAI 的 Sora 啟發的視頻壓縮網絡，在時間維度上實現了 4 倍壓縮。這消除了幀提取的需求，使視頻能夠以原始幀率生成。
由于訓練 3D VAE 的計算需求較高，嘗試利用 2D VAE 中嵌入的知識。通過 2D VAE 壓縮后，觀察到時間上相鄰的特征高度相關。基于這一觀察，開發了一個簡單但有效的視頻壓縮網絡，先以 8×8 的比例進行空間壓縮，然后在時間上以 4 倍壓縮。網絡架構如下圖5所示。
- 2D VAE 初始化：使用SDXL 的預訓練 VAE。
- 3D VAE 架構：采用Magvit-v2 的 VAE結構，總參數量為300M。結合2D VAE后，視頻壓縮網絡的總參數量為384M。
- 訓練細節：3D VAE 使用來自Pexels和Pixabay的視頻進行訓練，總共訓練了1.2M 步，本地批量大小為1。訓練數據主要為分辨率256×256的17 幀短視頻片段。
- 優化：在 3D VAE 中使用了因果卷積以提高圖像重建精度。
訓練過程
訓練過程分為三個階段：
- 階段 1（0–380k 步）：使用8 張 GPU進行訓練，凍結2D VAE的權重。目標包括重建2D VAE壓縮的特征，并應用身份損失（identity loss）來對齊3D VAE和2D VAE的特征。身份損失加速了收斂速度，并提高了初始圖像重建質量。
- 階段 2（380k–640k 步）：移除了身份損失，3D VAE 通過訓練進一步優化其時間維度的理解能力。
- 階段 3（640k–1.2M 步）：發現僅重建2D VAE特征不足以進一步提升模型性能，因此損失函數被替換為直接重建原始視頻。在此階段，使用24 張 GPU，并通過隨機化視頻長度（最長34 幀）結合適當的零填充實現混合視頻長度訓練，從而提高了模型對不同視頻時長的魯棒性。
在前兩個階段，數據集由 80% 視頻數據 和 20% 圖像數據 組成。對于視頻訓練，使用 17 幀視頻片段；圖像數據通過零填充與視頻輸入格式匹配。然而，這種方法導致了非標準長度視頻的模糊問題，而第三階段的混合長度訓練有效解決了這一問題。
堆疊的 VAE 架構在推理時占用的內存極少，因為輸入已經經過壓縮。為了提高效率，輸入視頻被分割為 17 幀的片段。與另一個開源 3D VAE 相比，本文的模型在顯著降低計算成本的同時，性能與其相當。
架構
本文的模型架構基于 PixArt，這是一種圖像擴散 Transformer。文本通過 T5 文本編碼器 進行編碼，視頻和文本隱空間變量之間采用 交叉注意力 機制。
為了實現高效的視頻生成，本文引入了一種空間-時間注意力機制，即 空間-時間擴散 Transformer（STDiT），靈感來源于 Latte，替代了對所有tokens的完全注意力機制。具體來說：
- 空間自注意力在每一幀內應用；
- 時間注意力在相同空間位置的幀之間應用。
為了專注于視頻生成，設計了基于強大的預訓練圖像生成模型的架構。模型以 PixArt-α 初始化，這是一種 T5 條件 DiT 結構，針對高質量和高效圖像生成進行了優化。為新引入的時間注意力設計的投影層初始值設為 0，以保持模型在訓練開始時的原始圖像生成能力。
時間注意力的加入使參數數量從 580M 增加到 1.1B。
條件控制
盡管文本到視頻生成具有高度的多樣性，但某些應用場景需要更精確的控制。為了實現這一點，通過引入圖像和視頻輸入的遮罩策略，擴展了模型的功能，支持圖像到圖像和視頻到視頻生成。如下圖 7 所示。
遮罩策略
在該方法中：
- 被指定為條件控制的幀會被取消遮罩。
- 在前向傳播過程中：
- 取消遮罩的幀被賦予時間步為 0。
- 其他幀保持其擴散時間步。
然而，將該策略直接應用于預訓練模型通常會產生次優結果，因為擴散模型未經過處理單個樣本中混合時間步的訓練。
受到 UL2 的啟發，通過隨機遮罩策略解決了這一問題。具體來說，幀會以隨機模式取消遮罩，例如：
- 第一幀
- 前k幀
- 最后一幀
- 后k幀
- 前后k幀的組合
- 或完全隨機的幀。
在 Open-Sora 1.0 的基礎上，對 50% 的訓練樣本應用mask策略進行實驗，發現模型在 10k 步之后能夠有效學習圖像和視頻的條件控制能力，同時對文本到視頻性能的影響最小。而較低的mask概率（例如 30%）會降低條件控制的效果。因此，從頭開始使用此mask策略預訓練模型。
為了進一步增強模型的控制能力，在字幕中附加分數，并將其作為附加條件輸入。這些分數包括美學分數、運動分數和相機運動描述。例如，一個美學分數為 5.5、運動分數為 10，并檢測到相機運動為“向左平移”的視頻，其字幕格式為：
[原始字幕] 美學分數: 5.5，運動分數: 10，相機運動: 向左平移。
在推理過程中，這些分數也可以調整以影響視頻生成。對于相機運動條件控制，手動標注了 13,000 個高置信度的片段。
這種方法使模型能夠對條件輸入有更細致的理解，從而提高了其在各種任務中生成高質量、上下文相關視頻的能力。
訓練策略
多分辨率和多長寬比
正如 Sora 的報告中所指出的，使用視頻的原始分辨率、長寬比和長度進行訓練能夠提高采樣的靈活性，并增強畫面構圖能力。為實現這一目標，評估了三種方法：
- NaViT
- 通過遮罩支持同一批次內的動態尺寸，效率損失最小。
- 實現復雜，可能無法充分利用如 Flash Attention [8] 之類的優化內核。
- 填充（FiT）
- 通過將較小分辨率填充到最大尺寸，支持同一批次內的動態尺寸。
- 實現簡單，但對分辨率變化較大的情況下，內存使用效率較低。
- 分桶（Bucket）（SDXL, PixArt）
- 通過將樣本分組到預定義的“桶”中支持跨批次的動態尺寸。
- 每個批次內分辨率、幀數和長寬比固定。分桶方法避免了遮罩或填充的復雜性，同時在統一尺寸的輸入上可以利用優化操作。
- 然而，這種方法的靈活性受到預定義尺寸集合的限制。
為簡單高效起見，采用了基于分桶的方法。預定義了一組固定的分辨率、長寬比和幀長度，并據此將樣本分配到相應的桶中。每個桶由三元組 分辨率幀數長寬比 定義，以覆蓋大多數常見的視頻格式。在每個訓練周期前，數據集會被重新打亂，樣本被分配到能夠容納其分辨率和幀長度的最大桶中。
為了進一步優化計算資源，為每個桶引入了兩個附加屬性：
- 保留概率（Probability of Keeping in the Bucket）
高分辨率視頻會根據概率被下采樣為較低分辨率，從而有效降低計算成本。 - 批次大小（Batch Size）
根據每個桶的情況調整批次大小，以平衡 GPU 負載，確保資源利用效率。
通過對這些參數進行微調，我們在桶之間實現了樣本的均衡分布，在保持高質量視頻生成的同時提升了總體訓練效率。
這種基于分桶的策略在實現簡單性與計算效率之間提供了一個實用的折中，支持具有多樣化分辨率和長寬比的視頻訓練。
模型適配
從 PixArt-Σ 2K checkpoint 開始，該模型使用 DDPM 和 SDXL VAE，在更高分辨率下進行訓練。通過在較小數據集上微調，該模型被高效地適配到視頻生成任務中。適配過程包括多個連續階段，全部在 8 個 H100 GPU 上完成：
- 多分辨率圖像生成
- 訓練模型以支持從 144p 到 2K 的分辨率，共 20k 步。
- 引入 QK 正則化
- 為提高穩定性，添加 QK-norm，訓練 18k 步。
- 轉向修正流
- 從離散時間的 DDPM 過渡到連續時間的修正流，共 10k 步。
- 強化修正流訓練
- 加入 logit-norm 采樣和分辨率感知時間步采樣，訓練 33k 步。
- 更小的 AdamW epsilon
- 新 VAE 和 FPS 條件控制
- 用 Open-Sora 的 VAE 替換原有 VAE，添加 FPS 條件控制到時間步條件控制中，訓練 25k 步。
- 在修正流訓練中，對每個通道進行歸一化被證明是關鍵步驟。
- 時間注意力模塊
- 添加零初始化的時間注意力模塊，初始僅在圖像上訓練 3k 步。
- 時間模塊的mask策略
- 使用mask策略專注于視頻的時間注意力模塊，訓練 38k 步。
適配后模型的優勢
完成適配后，模型不僅保留了生成高質量圖像的能力，還在視頻生成方面獲得了多重優勢：
- 加速訓練和推理
- 修正流將視頻采樣步驟從 100 減少到 30，顯著降低推理時間。
- 增強穩定性
- QK-norm 支持更激進的優化，提升訓練效率。
- 高效的時間維度壓縮
- 新 VAE 將時間維度壓縮至原來的四分之一，減少計算成本。
- 分辨率靈活性
- 模型能夠生成多種分辨率的視頻，從 144p 到 2K，支持多樣化應用場景。
這種全面的適配不僅提升了模型的視頻生成能力，還確保了訓練的高效性和可擴展性，為開源基于擴散的視頻生成樹立了新的標準。
多階段訓練
為了在有限的計算預算內優化性能，我們根據數據質量精心組織訓練數據，并將訓練過程分為三個階段。模型在一個 12×8 GPU 的設置上進行了大約兩周的訓練，完成了約 70k 步。
在第一階段，模型在 Webvid-10M 數據集（包含 40,000 小時的視頻）上進行了 30k 步訓練（2 個 epoch）。該數據集主要包含分辨率低于 360p 且帶有水印的視頻，非常適合初步訓練。我們專注于 240p 和 360p 分辨率的視頻，長度為 2 到 16 秒。訓練使用了原始數據集的字幕。
在第二階段，使用 Panda-70M 數據集 進行了訓練。由于該數據集的質量不均，使用了官方的 30M 子集，并篩選出只有美學分數高于 4.5 的視頻，最終得到 20M 的子集（41,000 小時）。訓練主要集中在 360p 和 480p 分辨率的視頻上，共進行了 23k 步訓練，約為 0.5 個 epoch。雖然這一階段的訓練沒有完全完成，但它為模型的廣泛應用提供了足夠的改進。
最后階段包括從各種來源精選的約 2M 高質量視頻片段，總計 5,000 小時。這些視頻來自 MiraData 和 Vript，并由 GPT 生成了字幕，其他來源則使用 PLLaVA 標注。該階段主要集中在更高分辨率（720p 和 1080p）的訓練，以提升模型處理更大分辨率的能力。在訓練過程中，采用了 25% 的mask比率，訓練共進行了 15k 步（約 2 個 epoch）。
為了進行驗證，從 Pixabay 中隨機抽取了 1k 個視頻來評估模型的性能。評估損失值針對不同長度（2s、4s、8s、16s）和不同分辨率（144p、240p、360p、480p、720p）的圖像和視頻進行了計算。每個配置的損失值在 10 個等距的時間步上進行了平均。
還在訓練過程中跟蹤了 VBench 分數。VBench 是一個用于評估短視頻生成的自動化基準工具。使用 240p 2 秒的視頻計算分數，為模型的進展提供了額外的驗證。評估損失和 VBench 分數都確認了模型在整個訓練過程中不斷取得進展。訓練過程中的 VBench 分數和驗證損失見下圖 9。
下表 2 顯示了各個模型的 VBench 分數，證明了 Open-Sora 在開源視頻生成模型中達到了行業領先水平。
結論
Open-Sora 代表了開源視頻生成領域的一次重大進步，提供了一個全面的框架，包括數據處理、訓練代碼和模型權重。通過成功地重現了 Sora 報告中的關鍵技術，并實現了最長 16 秒的視頻生成，分辨率可達 720p，且支持可控的運動動態，Open-Sora 使得先進的視頻生成技術得以普及。該項目不僅促進了社區的合作，還為該領域未來的發展奠定了基礎。
本文轉自AI生成未來 ，作者：AI生成未來
