視頻大模型無損提速:刪除多余token,訓練時間減少30%,幀率越高效果越好 | NeurIPS
Don’t look twice!
把連續相同的圖像塊合并成一個token,就能讓Transformer的視頻處理速度大幅提升。
卡內基梅隆大學提出了視頻大模型加速方法Run-Length Tokenization(RLT),被NeurIPS 2024選為Spotlight論文。
在精度幾乎沒有損失的前提下,RLT可以讓模型訓練和推理速度雙雙提升。
一般情況下,利用RLT,Transformer視頻識別模型的訓練時間可縮短30%,推理階段提速率提升更是可達67%。
對于高幀率和長視頻,RLT的效果更加明顯,30fps視頻的訓練速度可提升1倍,長視頻訓練token減少80%。
相比于傳統的剪枝方法,RLT能用更小的精度損失實現更好的加速效果。
有人想到了電視劇中的評論,認為這項研究找到了在壓縮空間中進行搜索的方法。
DeepMind科學家Sander Dieleman則評價稱,這項研究是一種“非主流”(Off-the-grid)的創新方法,但比起其他復雜的非主流研究,又顯得非常簡潔。
重復圖像塊合為一個token
RLT的核心原理,是利用視頻中存在大量時間上重復的圖像塊這一特點,將重復的圖像塊合并為一個token表示。
這種情況下,還需要用一個位置編碼來表示這個token的長度,但RLT總體上減少了輸入的token數量。
要想完成重復token的修剪,首先要對視頻進行分塊。
具體來說,視頻在空間和時間維度上會被劃分成固定大小的圖像塊,每個圖像塊的大小為C×D_x×D_y×D_t,每個圖像塊都對應一個空間-時間位置。
(其中C是通道數,D_x和D_y是空間維度大小,D_t是時間維度大小。)
劃分完成之后,需要比較時間上相鄰的圖像塊,判斷它們是否相似,也就是是否需要合并。
對于時間位置相差1的兩個圖像塊P_1和P_2,取P_1的第一幀和P_2的最后一幀,計算它們的L1距離。
如果距離小于一個預設的閾值τ,就認為P_1和P_2是靜態重復的(閾值τ表示允許多大程度的相似性,設置與具體數據集無關)。
完成判別之后,重復的圖像塊會被移除。
對于一串連續的、兩兩之間都是靜態重復的圖像塊,RLT只保留第一個塊對應的token。
這一步是在patch embedding之前完成的,因此移除token不需要改動模型結構。
經過這一步,輸入的token數量從N_P降低到了N_P’(N_P’≤N_P)。
為了讓合并后的token仍然能夠反映完整的視頻信息,接下來要給每個token加上長度編碼。
對于一個保留下來的token,系統會計算它所代表的原始token的長度l_i,也就是它到下一個沒有被移除的token的距離。
長度信息l_i與token的空間-時間位置(x,y,t)一起,用一個可學習的長度編碼矩陣映射成一個d維的embedding向量,與patch embedding相加,作為輸入token的最終表示。
最后只需要將處理后的token序列輸入到視頻Transformer中,進行常規的訓練或推理過程。
不過需要注意的是,由于每個視頻樣本計算出的token數量N_P’不盡相同,樣本之間可能有較大差異。
然而標準的Transformer是按批次處理固定長度的序列的。
為了能在一個批次中處理多個長度不一的視頻,RLT采用了一種稱為“example packing”的方法,將這一批中所有視頻樣本的token序列首尾相連,拼成一個超長的序列,作為Transformer的輸入。
這樣的話,Transformer實際上是在處理一個批次大小為1、長度為所有樣本token數量之和的序列。
通過以上步驟,RLT能夠去除視頻中的許多冗余token,在幾乎不損失精度的情況下,大幅降低內存占用和計算量,加速視頻Transformer的訓練和推理。
訓練時長下降30%
在訓練階段,RLT對ViT-B和ViT-L兩種規模的模型都有很好的加速效果。
在Kinetics-400上,ViT-BRLT和ViT-L訓練時間分別從14.4小時和21.6小時,降低到10.2小時和15.4小時,降幅均接近30%左右,精度損失不超過0.1個百分點;
在SSv2上,兩者的訓練時間分別從10.1和15.2小時,降低到7.2和10.8小時,降幅也接近30%,精度同樣僅下降0.1個百分點。
相比之下,傳統的剪枝方法Token Merging在精度下降0.1-0.5個百分點的情況下,加速只有10-20%。
在推理階段,也不需要額外的訓練,就可以將RLT作為現成的tokenizer,達到很好的加速效果。
具體來說,RLT能在幾乎不犧牲精度的情況下(不超過0.5個百分點),將推理階段的計算量和延遲降低30-60%。
同樣在Kinetics-400和SSv2上,對于ViT-B和ViT-L,RLT都能帶來60%以上的推理加速。
對于更大的ViT-H,在Kinetics-400上,RLT也能實現45%的加速效果。
特別地,作者還針對高幀率和長時長視頻數據集進行了測試,發現RLT帶來的token下降幅度比在普通數據集中更高。
同時在高幀率數據集當中,RLT能夠在精度損失同樣低的情況下,實現更好的加速效果。
而且幀率越高效果也越明顯,對于30fps的視頻,加速可達100%。
論文地址:https://arxiv.org/abs/2411.05222
代碼:https://github.com/rccchoudhury/rlt