视频大模型无损提速:删除多余token,训练时间减少30%

  克雷西发自凹非寺

  量子位公众号 QbitAI

  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