新智元報(bào)道
編輯:桃子
【新智元導(dǎo)讀】AI視頻生成進(jìn)入了秒生極速時(shí)代!UCSD等機(jī)構(gòu)發(fā)布的FastWan系模型,在一張H200上,實(shí)現(xiàn)了5秒即生視頻。稀疏蒸餾,讓去噪時(shí)間大減,刷新SOTA。
單塊H200,5秒即生一個(gè)5秒視頻。
最近,UCSD、UC伯克利、MBZUAI三大機(jī)構(gòu)聯(lián)手,祭出FastWan系視頻生成模型。
論文地址:https://arxiv.org/pdf/2505.13389
它的核心采用了「稀疏蒸餾」全新的訓(xùn)練方案,實(shí)現(xiàn)了高效生成,讓視頻去噪速度實(shí)現(xiàn)70倍飆升。
基于FastVideo架構(gòu),F(xiàn)astWan2.1-1.3B在單張H200上,去噪時(shí)間僅1秒,5秒內(nèi)生成了480p的5秒視頻。
在一張RTX 4090上,則耗時(shí)21秒生成一個(gè)視頻,去噪時(shí)間2.8秒。
若僅計(jì)算DiT處理時(shí)間
升級(jí)版FastWan2.2-5B,在單張H200上僅用16秒即可生成720P的5秒視頻。
FastWan模型權(quán)重、訓(xùn)練方案和數(shù)據(jù)集全部開(kāi)源
如今,終于實(shí)現(xiàn)AI實(shí)時(shí)視頻的生成了。
稀疏蒸餾,AI視頻進(jìn)入極速模式
「稀疏蒸餾」究竟是什么,能夠讓模型如此快速地生成視頻?
一直以來(lái),視頻擴(kuò)散模型成為了AI視頻生成領(lǐng)域的主流,比如Sora采用了擴(kuò)散模型+Transformer架構(gòu)。
這些模型雖強(qiáng)大,卻長(zhǎng)期受困于兩大瓶頸:
1. 生成視頻時(shí),需要海量的去噪步驟
2. 處理長(zhǎng)序列時(shí)的注意力二次方計(jì)算成本,高分辨率視頻必然面臨此問(wèn)題。
就以Wan2.1-14B為例,模型需運(yùn)行50次擴(kuò)散步驟,生成5秒720P視頻需處理超8萬(wàn)token,其中注意力操作甚至吞噬85%以上的推理時(shí)間。
此時(shí)此刻,「稀疏蒸餾」就成為了大殺器。
作為FastWan的核心創(chuàng)新,它首次在統(tǒng)一框架中實(shí)現(xiàn)稀疏注意力與去噪步驟蒸餾的聯(lián)合訓(xùn)練。
其本質(zhì)是回答一個(gè)根本問(wèn)題:在應(yīng)用極端擴(kuò)散壓縮時(shí),如用3步替代50步,能否保留稀疏注意力的加速優(yōu)勢(shì)?
先前研究認(rèn)為并不可行,而最新論文則通過(guò)「視頻稀疏注意力」(VSA)改寫了答案。
傳統(tǒng)稀疏注意力,為何會(huì)在蒸餾中失效?
當(dāng)前,現(xiàn)有的方法如STA、SVG,依賴的是多步去噪中的冗余性,來(lái)修剪注意力圖,通常僅對(duì)后期去噪步驟稀疏化。
但當(dāng)蒸餾將50步壓縮至1-4步時(shí),其依賴的冗余性徹底消失。
實(shí)驗(yàn)證實(shí),傳統(tǒng)方案在少于10步的設(shè)置下性能急劇退化——盡管稀疏注意力本身能帶來(lái)3倍加速,蒸餾卻可實(shí)現(xiàn)20倍以上增益。
要使稀疏注意力真正具備生產(chǎn)價(jià)值,必須使其與蒸餾訓(xùn)練兼容。
視頻稀疏注意力(VSA)是動(dòng)態(tài)稀疏注意力核心算法,能夠自主識(shí)別序列中的關(guān)鍵token。
不同于依賴啟發(fā)式規(guī)則的方案,VSA可在訓(xùn)練過(guò)程中直接替代FlashAttention,通過(guò)數(shù)據(jù)驅(qū)動(dòng)的方式學(xué)習(xí)最優(yōu)稀疏模式,同時(shí)最大限度保持生成質(zhì)量。
在步驟蒸餾過(guò)程中,當(dāng)學(xué)生模型學(xué)習(xí)用更少步驟去噪時(shí),VSA無(wú)需依賴多步去噪的冗余性來(lái)修剪注意力圖,而是能動(dòng)態(tài)適應(yīng)新的稀疏模式。
這使得VSA成為,首個(gè)完全兼容蒸餾訓(xùn)練的稀疏注意力機(jī)制。甚至,他們甚至實(shí)現(xiàn)了VSA與蒸餾的同步訓(xùn)練!
據(jù)團(tuán)隊(duì)所知,這是稀疏注意力領(lǐng)域的重大突破。
三大組件,全適配
基于視頻稀疏注意力(VSA)技術(shù),團(tuán)隊(duì)創(chuàng)新性地提出了稀疏蒸餾方法。
這是一種將稀疏注意力訓(xùn)練與步驟蒸餾相結(jié)合的模型后訓(xùn)練技術(shù)。
它的核心思想,是讓一個(gè)「少步數(shù)+稀疏化」的學(xué)生模型學(xué)會(huì)匹配「完整步數(shù)+密集計(jì)算」教師模型的輸出分布。
如下圖所示,該技術(shù)的整體框架包含以下關(guān)鍵要素:
稀疏學(xué)生網(wǎng)絡(luò)(VSA驅(qū)動(dòng),可訓(xùn)練)
真實(shí)評(píng)分網(wǎng)絡(luò)(凍結(jié),全注意力)
偽評(píng)分網(wǎng)絡(luò)(可訓(xùn)練,全注意力)
這三個(gè)組件均基于Wan2.1模型初始化。
訓(xùn)練時(shí),經(jīng)過(guò)稀疏蒸餾的學(xué)生網(wǎng)絡(luò)接收帶噪聲視頻輸入,通過(guò)VSA執(zhí)行單步去噪生成輸出。
該輸出會(huì)被重新添加噪聲,隨后分別輸入到兩個(gè)全注意力評(píng)分網(wǎng)絡(luò)——它們各自執(zhí)行一次全注意力去噪。
兩個(gè)分支輸出的差異構(gòu)成分布匹配梯度,通過(guò)反向傳播優(yōu)化學(xué)生網(wǎng)絡(luò);同時(shí)偽評(píng)分網(wǎng)絡(luò)會(huì)根據(jù)學(xué)生輸出的擴(kuò)散損失進(jìn)行更新。
這種架構(gòu)的精妙之處在于:學(xué)生模型采用VSA保證計(jì)算效率,而兩個(gè)評(píng)分網(wǎng)絡(luò)保持全注意力,以確保訓(xùn)練監(jiān)督的高保真度。
這種架構(gòu)的精妙之處在于:這種設(shè)計(jì)實(shí)現(xiàn)了運(yùn)行時(shí)加速(學(xué)生模型)與蒸餾質(zhì)量(評(píng)分網(wǎng)絡(luò))的解耦,使得稀疏注意力能夠與激進(jìn)的步數(shù)縮減策略兼容。
更廣泛地說(shuō),由于稀疏注意力僅作用于學(xué)生模型,該方案可適配各類蒸餾方法,包括一致性蒸餾、漸進(jìn)式蒸餾或基于GAN的蒸餾損失等。
那么,F(xiàn)astWan如何實(shí)現(xiàn)蒸餾的呢?
高質(zhì)量數(shù)據(jù)對(duì)任何訓(xùn)練方案都至關(guān)重要,尤其是對(duì)擴(kuò)散模型而言。為此,研究人員選擇使用高質(zhì)量的Wan模型自主生成合成數(shù)據(jù)集。
具體而言,采用Wan2.1-T2V-14B生成60萬(wàn)條480P視頻和25萬(wàn)條720P視頻,通過(guò)Wan2.2-TI2V-5B生成3.2萬(wàn)條視頻。
采用DMD進(jìn)行稀疏蒸餾時(shí),需在GPU內(nèi)存中同時(shí)加載三個(gè)140億參數(shù)大模型:
·學(xué)生模型
·可訓(xùn)練偽分?jǐn)?shù)模型
·凍結(jié)真分?jǐn)?shù)模型
其中兩個(gè)模型(學(xué)生模型與偽分?jǐn)?shù)模型)需持續(xù)訓(xùn)練,既要存儲(chǔ)優(yōu)化器狀態(tài)又要保留梯度,加之長(zhǎng)序列長(zhǎng)度的特性,使得內(nèi)存效率成為關(guān)鍵挑戰(zhàn)。
為此,他們提出的關(guān)鍵解決方案是:
1. 通過(guò)FSDP2實(shí)現(xiàn)三模型的參數(shù)跨GPU分片,顯著降低內(nèi)存開(kāi)銷
2. 應(yīng)用激活檢查點(diǎn)技術(shù)緩解長(zhǎng)序列產(chǎn)生的高激活內(nèi)存
3. 精細(xì)控制蒸餾各階段(如更新學(xué)生模型/偽分?jǐn)?shù)模型時(shí))的梯度計(jì)算開(kāi)關(guān)
4. 引入梯度累積在有限顯存下提升有效批次規(guī)模
Wan2.1-T2V-1.3B的稀疏蒸餾在64張H200 GPU上運(yùn)行4000步,總計(jì)消耗768 GPU小時(shí)。
一張卡,秒生視頻
在Scaling實(shí)驗(yàn)中,研究團(tuán)隊(duì)預(yù)訓(xùn)練一個(gè)4.1億參數(shù)視頻DiT模型,潛在空間維度位(16, 32, 32)。
在保持87.5%稀疏度情況下,VSA取得的損失值與全注意力機(jī)制幾乎一致。
同時(shí),它將注意力計(jì)算的FLOPS降低8倍,端到端訓(xùn)練FLOPS減少2.53倍。
從6000萬(wàn)擴(kuò)展到14億參數(shù)規(guī)模,進(jìn)一步證實(shí)了VSA始終能比全注意力機(jī)制實(shí)現(xiàn)更優(yōu)的「帕累托前沿」。
為評(píng)估VSA的實(shí)際效果,團(tuán)隊(duì)在Wan-14B生成的視頻潛空間(16×28×52)合成數(shù)據(jù)上,對(duì)Wan-1.3B進(jìn)行了VSA微調(diào)。
如表2所示,采用VSA的模型在VBench評(píng)分上甚至超越了原始Wan-1.3B。
在極端稀疏條件下,與免訓(xùn)練的注意力稀疏方法SVG對(duì)比時(shí),VSA盡管稀疏度更高仍表現(xiàn)更優(yōu),驗(yàn)證了稀疏注意力訓(xùn)練的有效性。
實(shí)際應(yīng)用中,Wan-1.3B的DiT推理時(shí)間從全注意力模式的31秒降至VSA模式的18秒。
VSA精細(xì)塊稀疏內(nèi)核在長(zhǎng)序列場(chǎng)景下,更加接近理論極限,相較于FlashAttention-3實(shí)現(xiàn)了近7倍加速。
即使計(jì)入粗粒度階段計(jì)算開(kāi)銷,VSA仍保持6倍以上的加速優(yōu)勢(shì)。
相比之下,采用相同塊稀疏掩碼(64×64塊大?。┑腇lexAttention僅獲得2倍加速。
結(jié)果顯示,將VSA應(yīng)用于Wan-1.3B和Hunyuan模型時(shí)(圖4a),推理速度提升達(dá)2-3倍。
下圖5所示,研究團(tuán)隊(duì)還檢測(cè)了經(jīng)微調(diào)13億參數(shù)模型,在粗粒度階段生成的塊稀疏注意力,呈現(xiàn)高度動(dòng)態(tài)性。
最后,團(tuán)隊(duì)還對(duì)模型做了定性實(shí)驗(yàn),下圖展示了隨著訓(xùn)練推進(jìn),模型逐漸適應(yīng)稀疏注意力機(jī)制,最終恢復(fù)了生成連貫視頻的能力。
作者介紹
這個(gè)work來(lái)自于咱們的老朋友Hao AI Lab的FastVideo研究小組。
Peiyuan Zhang(張培源)
張培源是UCSD計(jì)算機(jī)科學(xué)與工程系一年級(jí)博士生,導(dǎo)師為Hao Zhang教授。
他曾在南洋理工大學(xué)擔(dān)任研究助理,師從Ziwei Liu。
他主要研究機(jī)器學(xué)習(xí)系統(tǒng)以及高效的視頻模型架構(gòu),是TinyLlama, lmms-eval, LongVA以及FastVideo的作者。
Yongqi Chen
Yongqi Chen是UCSD Hao AI Lab的研究實(shí)習(xí)生,導(dǎo)師是Hao Zhang教授。
他曾在密歇根大學(xué)(UMich)獲得機(jī)器人學(xué)碩士學(xué)位,此前于2023年以榮譽(yù)畢業(yè)生的身份畢業(yè)于浙江大學(xué)竺可楨學(xué)院,獲機(jī)器人工程學(xué)士學(xué)位。
現(xiàn)階段,他的研究方向聚焦高效視頻生成技術(shù),致力于實(shí)現(xiàn)實(shí)時(shí)交互式視頻生成。
Will Lin
Will Lin是UCSD計(jì)算機(jī)科學(xué)與工程系Hao AI Lab的博士生,導(dǎo)師是Hao Zhang教授。
他在德克薩斯大學(xué)奧斯汀分校獲得計(jì)算機(jī)科學(xué)學(xué)士學(xué)位。
目前在Anyscale實(shí)習(xí),是開(kāi)源項(xiàng)目vLLM的活躍貢獻(xiàn)者。
Haofeng Huang(黃浩峰)
黃浩峰是清華姚班本科生,導(dǎo)師為Jianfei Chen 和 Jun Zhu教授。
目前,他主要研究方向?yàn)楦咝C(jī)器學(xué)習(xí),重點(diǎn)關(guān)注注意力機(jī)制、量化加速,以及圖像/視頻壓縮技術(shù)。
參考資料:
https://x.com/haoailab/status/1952472986084372835
https://hao-ai-lab.github.io/blogs/fastvideo_post_training/
特別聲明:以上內(nèi)容(如有圖片或視頻亦包括在內(nèi))為自媒體平臺(tái)“網(wǎng)易號(hào)”用戶上傳并發(fā)布,本平臺(tái)僅提供信息存儲(chǔ)服務(wù)。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.