實(shí)時(shí) AI 視頻生成來了！

本周三，新加坡國立大學(xué)尤洋團(tuán)隊(duì)提出了業(yè)內(nèi)第一種可以實(shí)時(shí)輸出的，基于 DiT 的視頻生成方法。

該技術(shù)名為 Pyramid Attention Broadcast （PAB）。通過減少冗余注意力計(jì)算，PAB 實(shí)現(xiàn)了高達(dá) 21.6 FPS 的幀率和 10.6 倍的加速，同時(shí)不會(huì)犧牲包括 Open-Sora、Open-Sora-Plan 和 Latte 在內(nèi)的流行基于 DiT 的視頻生成模型的質(zhì)量。值得注意的是，作為一種不需要訓(xùn)練的方法，PAB 可以為任何未來基于 DiT 的視頻生成模型提供加速，讓其具備實(shí)時(shí)生成的能力。

自今年起，OpenAI 的 Sora 和其他基于 DiT 的視頻生成模型引起了 AI 領(lǐng)域的又一波浪潮。然而與圖像生成相比，人們對于視頻生成的關(guān)注點(diǎn)基本都在于質(zhì)量，很少有研究專注于探索如何加速 DiT 模型推理。加速視頻生成模型的推理對于生成式 AI 應(yīng)用來說已經(jīng)是當(dāng)務(wù)之急。

PAB 方法的出現(xiàn)，為我們打開了一條路。

原始方法與 PAB 視頻生成速度的比較。作者在 Open-Sora 上測試了 5 個(gè) 4s（192 幀）480p 分辨率的視頻。

GitHub 鏈接：https://github.com/NUS-HPC-AI-Lab/OpenDiT?tab=readme-ov-file#pyramid-attention-broadcast-pab-blogdoc

金字塔式注意力廣播

近期，Sora 和其他基于 DiT 的視頻生成模型引起了廣泛關(guān)注。然而，與圖像生成相比，很少有研究專注于加速基于 DiT 的視頻生成模型的推理。此外，生成單個(gè)視頻的推理成本可能很高。

圖 1：當(dāng)前擴(kuò)散步驟和先前擴(kuò)散步驟之間的注意力輸出差異，使用均方誤差 (MSE) 對差異進(jìn)行量化。

實(shí)現(xiàn)

這項(xiàng)研究揭示了視頻擴(kuò)散 transformer 中注意力機(jī)制的兩個(gè)關(guān)鍵觀察結(jié)果：

首先，不同時(shí)間步驟的注意力差異呈現(xiàn)出 U 形模式，在最初和最后 15% 的步驟中發(fā)生顯著變化，而中間 70% 的步驟則非常穩(wěn)定，差異很小。

其次，在穩(wěn)定的中間段內(nèi)，注意力類型之間存在差異：空間注意力變化最大，涉及邊緣、紋理等高頻元素；時(shí)間注意力表現(xiàn)出與視頻中的運(yùn)動(dòng)和動(dòng)態(tài)相關(guān)的中頻變化；跨模態(tài)注意力是最穩(wěn)定的，將文本與視頻內(nèi)容聯(lián)系起來，類似于反映文本語義的低頻信號(hào)。

基于此，研究團(tuán)隊(duì)提出金字塔式注意力廣播來減少不必要的注意力計(jì)算。在中間部分，注意力表現(xiàn)出微小的差異，該研究將一個(gè)擴(kuò)散步驟的注意力輸出廣播到幾個(gè)后續(xù)步驟，從而顯著降低計(jì)算成本。

此外，為了更有效的計(jì)算和最小的質(zhì)量損失，作者根據(jù)不同注意力的穩(wěn)定性和差異性設(shè)置了不同的廣播范圍。即使沒有后期訓(xùn)練，這種簡單而有效的策略也能實(shí)現(xiàn)高達(dá) 35% 的加速，同時(shí)生成內(nèi)容的質(zhì)量損失可以忽略不計(jì)。

圖 2：該研究提出了金字塔式注意力廣播，其中根據(jù)注意力差異為三個(gè)注意力設(shè)置不同的廣播范圍。注意力變化越小，廣播范圍越廣。在運(yùn)行時(shí)，該方法將注意力結(jié)果廣播到接下來的幾個(gè)步驟，以避免冗余的注意力計(jì)算。x_t 指的是時(shí)間步 t 的特征。

并行

下圖 3 為本文方法與原始動(dòng)態(tài)序列并行（Dynamic Sequence Paralle, DSP）之間的比較。當(dāng)時(shí)間注意力得到傳播時(shí)，則可以避免所有通信。

為了進(jìn)一步提升視頻生成速度，本文基于 DSP 來改進(jìn)序列并行。序列并行將視頻分割為跨多個(gè) GPU 的不同部分，從而減少了每個(gè) GPU 的工作負(fù)載并降低了生成延遲。不過，DSP 引入了大量的通信開銷，需要為時(shí)間注意力準(zhǔn)備兩個(gè) All to All 通信。

通過在 PAB 中傳播時(shí)間注意力，本文不再需要對時(shí)間注意力進(jìn)行計(jì)算，由此減少了通信。相應(yīng)地，通信開銷大幅降低了 50% 以上，使得實(shí)時(shí)視頻生成可以進(jìn)行更高效的分布式推理。

評估結(jié)果

加速

下圖為不同模型在 8 塊英偉達(dá) H100 GPU 上生成單個(gè)視頻時(shí)，測量得到的 PAB 總延遲。當(dāng)使用單塊 GPU 時(shí)，作者實(shí)現(xiàn)了 1.26 至 1.32 倍的加速，并在不同的調(diào)度器中保持穩(wěn)定。

當(dāng)擴(kuò)展到多塊 GPU 時(shí)，本文方法實(shí)現(xiàn)了 10.6 倍的加速，并得益于高效的序列并行改進(jìn)實(shí)現(xiàn)了與 GPU 數(shù)量之間的近線性擴(kuò)展。

定性結(jié)果

以下三個(gè)視頻分別為 Open-Sora、Open-Sora-Plan 和 Latte 三個(gè)不同的模型使用原始方法與本文方法的效果對比。可以看到，本文方法在不同的 GPU 數(shù)量下均實(shí)現(xiàn)了不同程度的 FPS 加速。

定量結(jié)果

下表為 Open-Sora、Open-Sora-Plan 和 Latte 三個(gè)模型的 LPIPS（學(xué)習(xí)感知圖像塊相似度）和 SSIM（結(jié)構(gòu)相似度）指標(biāo)結(jié)果。

更多技術(shù)細(xì)節(jié)和評估結(jié)果可以查看即將推出的論文。

項(xiàng)目地址：https://oahzxl.github.io/PAB/

本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/tidE-qSM3nZ8kUfjNcpMZA??