DiT 都能用，生成影片無質量損失，也不需要訓練。
實時 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/
參考連結：
https://oahzxl.github.io/PAB/