史上首個實時AI視頻生成技術:DiT通用,速度提升10.6倍 精華
實時 AI 視頻生成來了!
本周三,新加坡國立大學尤洋團隊提出了業內第一種可以實時輸出的,基于 DiT 的視頻生成方法。
該技術名為 Pyramid Attention Broadcast (PAB)。通過減少冗余注意力計算,PAB 實現了高達 21.6 FPS 的幀率和 10.6 倍的加速,同時不會犧牲包括 Open-Sora、Open-Sora-Plan 和 Latte 在內的流行基于 DiT 的視頻生成模型的質量。值得注意的是,作為一種不需要訓練的方法,PAB 可以為任何未來基于 DiT 的視頻生成模型提供加速,讓其具備實時生成的能力。
自今年起,OpenAI 的 Sora 和其他基于 DiT 的視頻生成模型引起了 AI 領域的又一波浪潮。然而與圖像生成相比,人們對于視頻生成的關注點基本都在于質量,很少有研究專注于探索如何加速 DiT 模型推理。加速視頻生成模型的推理對于生成式 AI 應用來說已經是當務之急。
PAB 方法的出現,為我們打開了一條路。
原始方法與 PAB 視頻生成速度的比較。作者在 Open-Sora 上測試了 5 個 4s(192 幀)480p 分辨率的視頻。
GitHub 鏈接:https://github.com/NUS-HPC-AI-Lab/OpenDiT?tab=readme-ov-file#pyramid-attention-broadcast-pab-blogdoc
金字塔式注意力廣播
近期,Sora 和其他基于 DiT 的視頻生成模型引起了廣泛關注。然而,與圖像生成相比,很少有研究專注于加速基于 DiT 的視頻生成模型的推理。此外,生成單個視頻的推理成本可能很高。
圖 1:當前擴散步驟和先前擴散步驟之間的注意力輸出差異,使用均方誤差 (MSE) 對差異進行量化。
實現
這項研究揭示了視頻擴散 transformer 中注意力機制的兩個關鍵觀察結果:
首先,不同時間步驟的注意力差異呈現出 U 形模式,在最初和最后 15% 的步驟中發生顯著變化,而中間 70% 的步驟則非常穩定,差異很小。
其次,在穩定的中間段內,注意力類型之間存在差異:空間注意力變化最大,涉及邊緣、紋理等高頻元素;時間注意力表現出與視頻中的運動和動態相關的中頻變化;跨模態注意力是最穩定的,將文本與視頻內容聯系起來,類似于反映文本語義的低頻信號。
基于此,研究團隊提出金字塔式注意力廣播來減少不必要的注意力計算。在中間部分,注意力表現出微小的差異,該研究將一個擴散步驟的注意力輸出廣播到幾個后續步驟,從而顯著降低計算成本。
此外,為了更有效的計算和最小的質量損失,作者根據不同注意力的穩定性和差異性設置了不同的廣播范圍。即使沒有后期訓練,這種簡單而有效的策略也能實現高達 35% 的加速,同時生成內容的質量損失可以忽略不計。
圖 2:該研究提出了金字塔式注意力廣播,其中根據注意力差異為三個注意力設置不同的廣播范圍。注意力變化越小,廣播范圍越廣。在運行時,該方法將注意力結果廣播到接下來的幾個步驟,以避免冗余的注意力計算。x_t 指的是時間步 t 的特征。
并行
下圖 3 為本文方法與原始動態序列并行(Dynamic Sequence Paralle, DSP)之間的比較。當時間注意力得到傳播時,則可以避免所有通信。
為了進一步提升視頻生成速度,本文基于 DSP 來改進序列并行。序列并行將視頻分割為跨多個 GPU 的不同部分,從而減少了每個 GPU 的工作負載并降低了生成延遲。不過,DSP 引入了大量的通信開銷,需要為時間注意力準備兩個 All to All 通信。
通過在 PAB 中傳播時間注意力,本文不再需要對時間注意力進行計算,由此減少了通信。相應地,通信開銷大幅降低了 50% 以上,使得實時視頻生成可以進行更高效的分布式推理。
評估結果
加速
下圖為不同模型在 8 塊英偉達 H100 GPU 上生成單個視頻時,測量得到的 PAB 總延遲。當使用單塊 GPU 時,作者實現了 1.26 至 1.32 倍的加速,并在不同的調度器中保持穩定。
當擴展到多塊 GPU 時,本文方法實現了 10.6 倍的加速,并得益于高效的序列并行改進實現了與 GPU 數量之間的近線性擴展。
定性結果
以下三個視頻分別為 Open-Sora、Open-Sora-Plan 和 Latte 三個不同的模型使用原始方法與本文方法的效果對比。可以看到,本文方法在不同的 GPU 數量下均實現了不同程度的 FPS 加速。
定量結果
下表為 Open-Sora、Open-Sora-Plan 和 Latte 三個模型的 LPIPS(學習感知圖像塊相似度)和 SSIM(結構相似度)指標結果。
更多技術細節和評估結果可以查看即將推出的論文。
項目地址:https://oahzxl.github.io/PAB/
本文轉自 機器之心 ,作者:機器之心
