跳转至

Keep-or-Drop? Adaptive Tokenizer for Compact Video Representation

会议: ECCV 2026
论文: ECCV 原文
PDF: EventHosts PDF
领域: 图像生成
关键词: 视频生成, 自适应分词, 变分自编码器, 流匹配, 扩散模型

一句话总结

提出 KATok 自适应视频 VAE,通过 Gumbel-Softmax 与软注意力掩码联合学习动态 token 筛选,并结合级联掩码先验解决稀疏 token 生成中的时空错位,实现高保真度与数倍加速的紧凑视频生成。

研究背景与动机

隐空间扩散模型(LDMs)已成为高保真图像与视频合成的核心基石,其通过变分自编码器(VAEs)将高维像素压缩至紧凑的连续隐空间,在维持视觉质量的同时显著降低生成计算开销。然而,当前视频生成广泛采用的卷积或标准 Transformer VAE 往往对全视频序列采用固定网格的压缩率(如固定的空间与时序降采样比例)。这种非自适应表征机制使 token 数量严格随视频空间分辨率与时间长度线性激增,严重制约了高分辨率与长视频生成的扩展性。

视频数据在时空维度上天然存在极高的信息冗余,静态背景或平滑无纹理区域所需的信息容量远低于剧烈运动或复杂纹理区域,固定预算划分不可避免地在无信息区域浪费大量表达容量。近期提出的可变长度分词方法虽然允许用户指定压缩预算,但其缺乏内容自适应性,在下游任务中需依靠额外的人工预算设定或耗时的推理期搜索,无法端到端自发决定每个视频片段到底需要多少 token。更棘手的是,若将自适应丢弃后的稀疏 token 直接用于扩散模型训练,由于空间网格结构被打破,生成模型极易出现内容与空间坐标错位(content-position misalignment),导致物体形变与时序抖动。

针对上述瓶颈,KATok 探索端到端可微分的自适应 token 学习机制,使模型能够根据视频时空复杂度动态决定每个 token 的保留与剔除,并解决稀疏 token 在下游扩散生成中的位置对其问题。核心 idea:构建端到端可微分的自适应 Transformer VAE(KATok),通过 Gumbel-Softmax 门控与解码器软注意力掩码实现数据驱动的动态 token 压缩,并提出级联掩码先验生成架构,在彻底消除时空内容错位的同时实现数倍的生成加速与更高画质。

方法详解

整体框架

KATok 的整体流程包含自适应分词与下游稀疏扩散生成两部分。在分词阶段,输入时空视频经过 3D 卷积/线性切块得到 patch 序列,连同全局 Register Tokens 输入带有 3D RoPE 的 Transformer 编码器提取特征。轻量分类头预测每个 token 的保留概率,并在训练期通过 Gumbel-Softmax 松弛采样软掩码,同时通过软注意力偏置将掩码传导至解码器交叉注意力中;解码器通过非对称精细查询序列重构视频。在下游生成阶段,为了避免直接在稀疏 token 上做扩散导致的坐标错位,系统解耦位置选择与内容生成:先由轻量掩码先验模型预测保留网格坐标,主流匹配模型(Flow Matching DiT)基于显式空间位置条件生成稀疏隐变量内容,最后由 KATok 解码器输出完整视频。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入视频 (T×H×W×C)"] --> B["自适应保留-丢弃分词<br/>Gumbel-Softmax门控与软掩码"]
    B --> C["软注意力掩码解码<br/>非对称多网格查询重构"]
    C --> D["级联掩码先验生成<br/>轻量先验预测活跃位置"]
    D --> E["稀疏流匹配生成<br/>条件化坐标生成紧凑内容"]
    E --> F["输出高保真视频"]

关键设计

1. 自适应保留-丢弃分词与软注意力掩码:让模型自发学习时空冗余裁剪 针对传统 VAE 压缩比例固定、在平坦时空背景中过度分配参数的痛点,KATok 在编码器后接入轻量分类头预测保留/丢弃对数几率 \(\alpha_i \in \mathbb{R}^2\)。训练阶段使用 Gumbel-Softmax 松弛技术采样可微软掩码 \(\tilde{m}_i\),将连续隐变量调制为 \(z_i = \tilde{m}_i \cdot \hat{z}_i\)。为了确保被丢弃的 token 在解码重构时彻底切断信息流动,解码器注意力得分被施加软偏置修正: $\(A_{ij} \leftarrow A_{ij} + \log(\tilde{m}_j + \varepsilon)\)$ 配合稀疏正则项 \(\mathcal{L}_{\text{sparse}} = \mathbb{E}_X [\sum_{i=1}^N \tilde{m}_i(X)]\),其权重 \(\lambda_{\text{sparse}}\) 从初始松弛阶段逐渐退火递增,使得模型自发学会在运动复杂区域保留致密 token、在静态无纹理背景中激进丢弃,消除了人工制定 token 预算的弊端。

2. 非对称粗到细解码与全局 Register Tokens:小隐空间与高保真重建兼顾 为在极致压缩率下保留微小几何边缘与高频细节,KATok 采用编码与解码分辨率非对称的设计。编码器在较大的 \(16^2 \times 8\) 粗粒度 3D patch 上进行切块与稀疏采样,大幅缩短送入注意力机制与下游生成的隐变量序列长度;解码器则构建 \(8^2 \times 4\) 精细网格的可学习查询向量,通过与保留的隐变量交叉注意力实现超分辨率级重构。同时,编码器引入全局 Register Tokens 作为锚点,吸收全局时空语义;消融证实去除 Register Tokens 会使 token 压缩率下降并损害全局结构,二者配合实现了隐空间极致精简与像素高精度的平衡。

3. 级联掩码先验条件生成:彻底消除稀疏 token 的时空坐标错位 直接在经自适应剪枝后的稀疏连续 token 上训练流匹配扩散模型时,模型难以在缺乏网格几何支撑的无序点集上同时对齐语义内容与时空坐标,导致严重的物体轮廓崩塌。虽然解耦时间步的双分支联合生成(Joint Generation)能够一定程度缓解该问题,但对两套噪声调度的超参数极其敏感。本文提出级联先验方案(Cascaded Generation):设计一个仅 8.3M 参数的微型掩码先验模型先行预测时空 token 网格的二值活跃掩码,提取激活位置坐标后,作为确定性位置编码显式注入到 686M 参数的主流匹配生成器中。这种将“在哪里生成”与“生成什么细节”完全解耦的机制,在极高稀疏度下依然保证了完美的时空几何一致性。

损失函数 / 训练策略

分词器训练目标综合多项约束: $\(\mathcal{L} = \mathcal{L}_{\text{recon}} + \lambda_{\text{KL}}\mathcal{L}_{\text{KL}} + \lambda_{\text{sparse}}\mathcal{L}_{\text{sparse}} + \lambda_{\text{adv}}\mathcal{L}_{\text{adv}} + \lambda_{\text{align}}\mathcal{L}_{\text{align}}\)$ 其中重构损失采用基于 S3D 网络的时空 Video-LPIPS,增强跨帧动作连贯性;\(\mathcal{L}_{\text{align}}\) 基于冻结的 vJEPA-2 特征进行语义表征对齐,并注入均匀分布在 \([0, 0.2]\) 区间的隐变量高斯噪声以强化下游生成鲁棒性。训练采用分阶段策略:单分辨率预训练(\(256^2 \times 16\))、多分辨率微调与判别器感知对抗微调。

实验关键数据

主实验

在 Panda-70M 验证集(5,000 条视频)上评估视频重建质量与压缩表现,以及在 UCF-101、SkyTimelapse 和 Kinetics-600 上的视频生成性能(SiT-XL 骨干网络)。

表 1:Panda-70M 视频重建与压缩性能对比

模型 分辨率 平均 Token 数 通道数 压缩率 \(\uparrow\) PSNR \(\uparrow\) SSIM \(\uparrow\) LPIPS \(\downarrow\) rFVD \(\downarrow\)
Omni-VAE \(256^2 \times 17^*\) 5120.00 8 96.00 28.10 0.88 0.05 7.84
Omni-VAE \(512^2 \times 33^*\) 36864.00 8 96.00 24.07 0.80 0.06 16.85
ElasticTok-KL \(256^2 \times 16\) 3845.56 8 102.25 30.52 0.91 0.06 12.37
KATok (Ours) \(256^2 \times 16\) 366.24 64 134.21 31.24 0.94 0.04 5.12
KATok (Ours) \(512^2 \times 32\) 1554.24 64 253.00 33.23 0.95 0.05 6.40

*注:\(*\) 表示指标在前 16/32 帧上评测;ElasticTok 不支持 \(512^2\) 评估。

表 2:跨数据集生成质量与效率对比 (gFVD \(\downarrow\))

生成模型 SkyTimelapse UCF-101 Kinetics-600 训练吞吐加速 生成吞吐 (videos/s)
Omni-VAE + SiT-XL 23.28 100.00 206.58 基准 (1.0×) 4.91
Elastic-KL + SiT-XL 95.53 712.56 — — 4.24
Ours-Joint 21.36 73.16 193.78 — 5.01
Ours-Cascaded 23.19 61.53 160.84 6.9× 15.71

消融实验

在 100K 步训练下针对 VAE 核心架构组件及扩散生成策略进行控制变量消融。

表 3:KATok 分词器组件消融分析 (100K 步,最大 Token 上限 514)

配置 PSNR \(\uparrow\) LPIPS \(\downarrow\) SSIM \(\uparrow\) 活跃 Token 数 \(\downarrow\) 说明
Full model 30.85 0.07 0.93 365.57 完整基准配置
− 隐变量正则 (latent reg.) 31.26 0.07 0.94 361.70 重建略微上升但损害生成
− 非对称解码 (\(16^2 \times 8\)) 29.61 0.11 0.92 377.80 细粒度重构大幅退化
− Video-LPIPS (改用图像 LPIPS) 29.28 0.11 0.91 404.00 缺乏时序监督,冗余 token 增加
− Register Tokens 29.17 0.12 0.91 410.20 失去全局锚点,抗剪枝稳定性下降
− 软注意力掩码 19.00 0.51 0.54 2.00 训练崩溃,仅剩 register tokens
− Gumbel-Softmax 18.91 0.52 0.53 2.00 梯度断裂导致训练崩溃

表 4:扩散生成策略在 UCF-101 上的 gFVD 消融

实验组 方案配置 gFVD \(\downarrow\) 核心分析
隐变量正则影响 包含隐变量正则 101.23 提升特征空间平滑度,大幅增益生成
去除隐变量正则 161.23 特征边界过拟合重建,生成质量急剧恶化
生成架构选择 朴素扩散生成 (Naïve) 95.69 严重受困于稀疏 token 的时空坐标错位
联合预测 (Joint Gen.) 73.16 联合去噪带来对齐,但受限于超参敏感度
级联生成 (Cascaded Gen.) 61.53 先验解耦结构与内容,取得最优质量与稳定性

关键发现

  • 软注意力掩码与可微采样是命脉:去除软注意力掩码或 Gumbel-Softmax 会直接造成优化崩溃,模型坍缩至仅保留 2 个 Register Token,证明显式注意力截断与可微梯度的结合是自适应保留学习的绝对前提。
  • Token 数量高度依赖时空变化熵:统计分析显示活跃 token 数量与视频时间熵具有高达 0.865 的 Pearson 相关系数(空间熵仅为 0.618,时空联合熵达 0.872)。在纯白或静态场景下仅需 28 个 token 即可完美重建,验证了其对时空动态的高度敏感性。
  • 级联先验带来速度与精度的双重飞跃:得益于 token 序列由 5,120 缩减至约 366,生成模型训练收敛提速 6.9 倍,推理吞吐达到 15.71 视频/秒(提升 3.2 倍),且 FVD 相比全量稠密 token 基准 OmniTokenizer 降低了 38.5 点。

亮点与洞察

  • 软注意力掩码与 Gumbel-Softmax 的协同:在交叉注意力矩阵中直接注入 \(\log(\tilde{m}_j + \varepsilon)\) 作为负无穷偏置,以完全连续可微的方式模拟硬截断,使反向传播能直接根据最终重构质量惩罚不重要的 token。
  • 自然涌现的时空可控性:在推理时直接调整输入的高斯噪声 token 长度(如从 200 增至 400),无需任何微调即可平滑调控生成视频的运动幅度与视觉复杂性。
  • 解耦先验解决稀疏几何错位:将 3D 点云级稀疏生成分解为“轻量 8M 掩码先验定坐标 + 大模型填充语义内容”,为离散/稀疏结构在连续流匹配生成中的应用提供了极具通用性的范式。

局限与展望

  • 极端动态场景下的 token 膨胀:面对全画幅剧烈晃动、雨雪粒子等全场非线性动态视频,自适应 token 比例会快速上升,失去相对稠密模型的显著计算优势。
  • 级联生成的两阶段误差累积:轻量掩码先验模型若预测了不合理的空间分布,后续主生成器无法纠正骨架结构错误,未来可探索端到端双向校准反馈机制。
  • 跨模态扩展潜力:目前仅在视频模态上验证,该范式具备推广到高分辨率 3D 点云、多视角 4D 动态场景重构与生成的广泛潜力。

相关工作与启发

  • vs OmniTokenizer: OmniTokenizer 采用固定尺寸的网格对时空 patch 统一压缩,缺乏对内容冗余的自适应能力;KATok 在降低 90% 以上生成 token 的同时,依靠内容感知保留取得了更优的重建与生成 FVD。
  • vs ElasticTok: ElasticTok 采用截断式的可变长度设计,需要二分查找进行样本级 token 预算搜索;KATok 则是端到端内容自适应的无监督涌现,推理完全免除了搜索开销,训练速度快 46.7 倍。
  • vs 传统 Token 剪枝 (DynamicViT / ToMe): 传统剪枝技术大多局限于判别式模型的推理期加速,难以支撑生成式重建;KATok 实现了训练与推理一体的连续生成隐空间构建。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次在连续视频扩散生成中实现端到端数据驱动的自适应 token 剪枝与级联先验对齐。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 Panda-70M、UCF-101、SkyTimelapse、Kinetics-600 并在重建、生成、消融及熵相关性上做了详尽验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 架构与机制叙述清晰,图表逻辑扎实,实验设计严谨自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为高分辨率、长时程视频生成提供了一条极具效率潜力与可扩展性的紧凑隐空间新路线。