TokenTrim: 面向自回归长视频生成的推理时Token剪枝方法 TokenTrim: Inference-Time Token Pruning for Autoregressive Long Video Generation
通过推理时剪枝不稳定token消除长视频生成中的时序漂移
前置知识
自回归视频生成
自回归视频生成是一种将长视频合成分解为顺序chunk的范式。模型先生成一个短视频片段,然后以前序片段的latent表征作为条件,迭代地生成后续片段。这种机制通过因果自注意力(causal self-attention)和KV缓存(key-value cache)实现跨chunk的时序依赖建模,使得模型能够生成远超单次生成能力的长视频序列。Self Forcing和Rolling Forcing是两种代表性的自回归推理策略。
本文的核心问题——时序漂移(temporal drift)——正是在自回归推理过程中产生的,理解这一范式是理解论文动机和技术方案的基础。
时序漂移(Temporal Drift)
时序漂移是自回归视频生成中的核心失败模式,指生成过程中误差逐渐累积和放大的现象。在自回归推理中,早期帧中的微小瑕疵(如颜色偏差、结构畸变)会被缓存到KV cache中,在后续生成步骤中被反复引用,导致这些误差线性或指数级放大,最终表现为物体身份改变、结构退化、运动失真等严重问题。这种现象也被称为暴露偏差(exposure bias),源于训练时使用真实历史与推理时使用自生成历史之间的分布偏移。
时序漂移是本文要解决的核心问题,论文的整个方法设计都围绕着如何在推理时检测和消除不稳定的latent token来阻断误差传播。
KV缓存(Key-Value Cache)
KV缓存是Transformer推理中的关键优化机制。在自回归视频生成中,模型对每个chunk进行自注意力计算时,会将之前chunk的注意力key和value投影存储起来,这样当前chunk的query可以直接与缓存的K、V进行注意力计算,而无需重新计算历史chunk的投影。形式上,注意力计算为softmax(QK^T/sqrt(d))V,其中K = [K_curr; K_cache],V = [V_curr; V_cache]。
TokenTrim的核心操作就是对KV缓存进行剪枝,通过移除不稳定的token位置来阻止误差传播,因此理解KV缓存的结构和作用是理解本文方法的关键。
Diffusion Transformer (DiT)
Diffusion Transformer是当前主流的视频生成模型架构,将Transformer架构与扩散模型结合。在视频DiT中,输入视频被表示为一组latent token(通过patchify操作将视频latent划分为时空patch),然后在去噪网络中反复应用自注意力机制。DiT在空间(帧内)和时间(跨帧)维度上进行注意力计算,是实现高质量视频生成的骨干网络。Wan2.1-1.3B是本文实验使用的基础视频扩散模型。
TokenTrim在DiT架构的latent空间中操作,通过比较相邻chunk的latent表征来检测不稳定的token,理解DiT的token结构有助于理解方法的技术细节。
FlowMo运动稳定初始化
FlowMo是一种推理时的运动一致性优化方法,通过计算生成帧之间的patch级时间方差,并应用梯度更新来最小化不连贯的运动轨迹。它通过引入运动感知方差引导(motion-aware variance guidance),在去噪过程中鼓励时间一致的运动模式,减少初始阶段的瑕疵。在TokenTrim中,FlowMo仅用于生成第一个batch,作为稳定的latent锚点。
FlowMo是TokenTrim的重要组成部分,提供运动稳定的初始化,与token剪枝机制互补,共同提升长视频生成质量。
研究动机
自回归视频生成是当前长视频合成的主流范式,通过将生成过程分解为顺序的chunk来扩展视频长度。然而,这种方法存在严重的时序漂移问题:在长序列生成中,早期帧中的微小瑕疵会通过KV缓存被反复引用和放大,导致物体身份改变、结构退化和运动失真。例如,使用Rolling Forcing生成的皮卡丘角色会出现肢体缺失或重复,颜色漂移和纹理退化在整个视频中逐渐加剧。现有方法如Self Forcing(通过随机梯度截断训练)和Rolling Forcing(通过滚动窗口联合去噪和注意力锚定机制)虽然有所改进,但主要依赖训练时的模型对齐,计算成本高昂且无法完全解决推理时的误差累积问题。更重要的是,这些方法缺乏评估缓存latent token可信度的机制,导致一旦某个区域的latent状态被破坏,就会持续影响后续生成。
本文的目标是本文的目标是提出一种纯推理时的方法来解决自回归长视频生成中的时序漂移问题。具体而言,TokenTrim旨在:(1)在推理时自动检测不稳定的latent token,这些token的表征与前序chunk存在显著偏差;(2)通过从KV缓存中移除这些不稳定的token来阻断误差传播;(3)在不修改模型架构、训练过程或需要额外训练的情况下,显著提升长序列视频的时间一致性和视觉质量;(4)保持极低的推理开销,使方法具有实际应用价值。
与已有工作不同的是,现有方法主要从两个角度应对时序漂移:训练时方法(如Self Forcing、Rolling Forcing、LongLive)通过重新训练或微调模型来对齐自回归任务,但计算成本高昂且无法完全消除推理时的误差累积;推理时方法(如FreeInit、FreeLong、FlowMo)主要通过频谱分析或运动优化来改善生成质量,但没有直接解决KV缓存中不稳定token的传播问题。TokenTrim的独特切入角度是:将时序漂移视为一个推理时的信息传播问题,通过直接在latent空间中识别和移除不稳定的token,将KV缓存从被动的历史日志转变为主动的自纠正记忆机制。这种方法不需要修改模型参数,可以与现有的自回归推理策略和推理时正交化技术无缝结合。
核心方法
TokenTrim的核心直觉是:时序漂移的根源不在于模型能力不足,而在于推理时不受控制地复用被破坏的latent conditioning token。当自回归生成过程中某个区域的latent状态被破坏后,它会在后续步骤中被反复引用并主导注意力计算,从而将误差向前传播。TokenTrim通过在latent空间中检测和移除这些不稳定的token来阻断误差传播。技术路线包括五个关键步骤:(1)使用FlowMo进行运动稳定的初始化,生成高质量的第一batch作为锚点;(2)对每个生成的batch构造latent summary,通过跨帧平均得到紧凑的latent表征;(3)计算每个spatial token的漂移分数,量化相邻chunk之间的latent偏差;(4)通过自适应阈值判断当前batch是否存在异常漂移;(5)当漂移超过阈值时,从KV缓存中剪除高漂移token并重新生成当前batch。
TokenTrim与已有方法的本质区别在于其主动的、token级的误差干预机制。现有方法要么被动地依赖模型自身的恢复能力(如Self Forcing),要么在频域或运动层面进行间接优化(如FlowMo、FreeLong),要么通过修改训练过程来增强模型的鲁棒性(如Rolling Forcing的注意力锚定)。TokenTrim则直接在推理时操作于KV缓存的token级别:通过比较相邻chunk的latent summary来检测每个spatial token的稳定性,然后移除不稳定的token。这种设计的核心创新是将漂移检测与硬剪枝相结合——不仅识别问题token,还从缓存中完全移除它们,而不是简单地降低权重或应用软约束。此外,自适应阈值机制(基于运行统计的均值加标准差)使得剪枝决策能够根据历史漂移动态调整,避免了固定阈值的局限性。
方法步骤详情
TokenTrim的单步执行流程如下:首先,给定前序batch X_{t-1}和候选当前batch X_t,分别将每帧编码为spatial latent token Z^{(f)}_{t-1} = E(x^{(f)}_{t-1}),然后通过跨帧平均构造latent summary Z_{t-1} = (1/F) * sum(Z^{(f)}_{t-1}),其中N是spatial token数量,D是latent维度。接着,计算每个spatial token的漂移分数d_i = ||Z_t(i) - Z_{t-1}(i)||_2。将漂移分数排序后,选取top p*N个最高漂移的token构成不稳定集合S_t,并计算漂移严重度D_t = (1/|S_t|) * sum(d_i)。然后,维护已接受batch的运行统计(均值mu_t和标准差sigma_t),当D_t > mu_t + lambda*sigma_t时触发剪枝。剪枝时,构造空间掩码m_t(i) = indicator[i not in S_t],对KV缓存进行过滤得到剪枝后的K和V,然后使用剪枝后的缓存重新生成当前batch。在前T_warm=2步的预热期内,禁用剪枝仅累积统计。每次最多执行R=1次重新生成。
技术新颖性
TokenTrim的技术新颖性体现在多个层面。首先,这是首个在推理时通过token级剪枝来解决自回归视频生成时序漂移的工作,开辟了一个新的研究方向。其次,漂移检测机制的设计独具匠心:通过比较latent summary(跨帧平均的紧凑表征)而非逐帧比较,既保证了token结构的一致性,又有效过滤了单帧噪声。硬剪枝策略(完全移除而非软加权)确保了不稳定的token无法继续影响后续生成,这比软约束更彻底。自适应阈值机制基于运行统计而非固定阈值,使得系统能够根据具体生成序列的漂移特性动态调整剪枝强度。此外,FlowMo仅用于第一batch的运动稳定初始化,与token剪枝形成互补效应——FlowMo处理早期初始化问题,TokenTrim处理后续累积问题。整个方法不修改模型参数,不引入额外训练,仅增加约8%的推理开销(相比之下FlowMo引入118%的开销),体现了极高的效率。
实验结果
TokenTrim在VBench基准上进行了全面评估,结果表明该方法在多个维度上显著优于基线方法。在Rolling Forcing上,TokenTrim将Final Score从75.12%提升至79.67%(+4.55%),其中Quality Score提升5.58%至87.30%,Semantic Score提升3.53%至72.05%。在Self Forcing上,改进更为显著,Final Score从75.93%提升至81.84%(+5.91%),Quality Score提升6.90%至89.79%,Semantic Score提升4.91%至73.89%。值得注意的是,TokenTrim大幅超越了推理时方法FlowMo:在Self Forcing上,FlowMo仅将Final Score提升至76.05%(+0.12%),而TokenTrim达到81.84%(+5.91%)。推理效率方面,TokenTrim仅增加8%的运行时间,而FlowMo导致2.18倍的减速。人类偏好研究进一步验证了这些结果:在VideoJAM-bench上,TokenTrim在时序漂移维度获得41.7%的偏好率(对比基线15.2%),运动连贯性30.3%(对比10.2%),美学质量40.3%(对比12.3%)。直接对比FlowMo时,TokenTrim在所有维度上都获得更高偏好,时序漂移维度达到49.7%对比10.2%。消融研究表明,10%的剪枝比例是最优选择,5%剪枝导致1.32%的性能下降,20%剪枝导致11.87%的严重退化。FlowMo初始化的贡献为2.34%的Final Score提升。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 长视频文本生成 | VBench Final Score | 79.67% (Rolling Forcing), 81.84% (Self Forcing) | 75.12% (Rolling Forcing), 75.93% (Self Forcing) | +4.55% (Rolling Forcing), +5.91% (Self Forcing) |
| 长视频文本生成 | VBench Semantic Score | 72.05% (Rolling Forcing), 73.89% (Self Forcing) | 68.52% (Rolling Forcing), 68.98% (Self Forcing) | +3.53% (Rolling Forcing), +4.91% (Self Forcing) |
| 长视频文本生成 | VBench Quality Score | 87.30% (Rolling Forcing), 89.79% (Self Forcing) | 81.72% (Rolling Forcing), 82.89% (Self Forcing) | +5.58% (Rolling Forcing), +6.90% (Self Forcing) |
| 长视频文本生成 | 推理时间开销 | 1.08x 基线时间 | 1.0x | 仅增加8%,远低于FlowMo的2.18x |
局限与改进
论文承认的主要局限性包括:首先,由于TokenTrim纯推理时操作且不修改模型参数,其改进效果受限于底层视频扩散骨干网络的能力和训练数据。当基础模型持续无法正确表示某个物体、保持身份或产生稳定运动时,TokenTrim主要能衰减误差传播而非完全纠正生成。其次,当前实现使用固定的硬剪枝预算(每步剪枝固定比例或top-k token),虽然轻量且开销可忽略,但单一的全局剪枝设置在不同prompt、内容类型和rollout长度下可能不是最优的。从独立观察来看,还有几个值得关注的局限:(1)方法依赖于latent summary的跨帧平均,这可能丢失细粒度的时间信息,特别是对于快速运动的场景;(2)漂移检测仅比较相邻chunk,可能无法捕捉跨多个chunk的累积漂移模式;(3)剪枝后重新生成的机制是单次的(R=1),对于特别困难的序列可能不够充分;(4)方法在Wan2.1-1.3B这样的小模型上验证,对于更大规模模型的效果尚未验证。
独立分析的弱点
TokenTrim的主要弱点包括以下几个方面。首先,漂移检测机制基于相邻chunk的latent summary比较,这种设计假设漂移是局部的、可检测的,但对于跨多个chunk的渐进式漂移可能不够敏感。改进方向可以考虑引入多尺度的漂移检测,不仅比较相邻chunk,还比较间隔更远的chunk来捕捉长期漂移模式。其次,固定的剪枝比例(p=0.1)对所有生成序列和所有步骤都相同,缺乏对具体内容的适应性。例如,包含复杂细节的场景可能需要更保守的剪枝,而简单背景的场景可能可以承受更激进的剪枝。改进方向是设计基于内容的自适应剪枝策略,根据当前chunk的复杂度和历史漂移模式动态调整剪枝比例。第三,hard剪枝策略虽然彻底,但可能导致信息丢失,特别是当被剪枝的token中包含部分有用信息时。可以探索混合策略,对高度不稳定的token进行硬剪枝,对中等不稳定的token进行软加权。最后,方法仅在30秒视频上进行评估,对于更长视频(如几分钟到十几分钟)的效果尚不明确。
未来方向
作者提出的重要未来方向是使剪枝自适应化:动态选择每步的剪枝率和剪枝结构,利用漂移统计或不确定性估计,使得困难序列获得更强的抑制而稳定序列保留更丰富的上下文。基于TokenTrim的成果,还可以延伸出多个研究方向:(1)将token级剪枝思想推广到其他模态的自回归生成,如长文本生成中的KV缓存管理;(2)探索预测性剪枝策略,通过预测未来chunk的漂移风险来提前进行预防性剪枝;(3)将TokenTrim与训练时方法(如LongLive的KV重新缓存机制)结合,在训练和推理两个层面协同优化;(4)研究剪枝策略的理论基础,建立漂移传播的数学模型,为剪枝决策提供理论指导;(5)开发跨模型通用的剪枝策略,使得在小模型上学到的剪枝策略能够迁移到大模型上。
复现评估
论文的复现条件相对良好。实验基于公开的Wan2.1-1.3B文本到视频模型,这是社区广泛使用的开源模型。TokenTrim本身是纯推理时方法,不需要额外训练,只需在现有推理流程中添加漂移检测和KV缓存剪枝模块。主要的复现挑战包括:(1)需要实现Rolling Forcing和Self Forcing两种自回归推理策略,这些方法本身也有一定的实现复杂度;(2)需要集成FlowMo用于第一batch的初始化;(3)VBench评估需要下载相应的评估数据和运行评估脚本。算力要求方面,实验在单张NVIDIA H100 GPU上进行,生成30秒16FPS 832x480分辨率的视频,这对于研究级实验是合理的。论文提供了详细的超参数设置(p=0.1, lambda=2.0, T_warm=2, R=1),便于复现。整体而言,具备视频生成领域经验的研究者应该能够复现本文结果。
论文图表