Keep-or-Drop? 面向紧凑视频表示的自适应 Tokenizer

TL;DR KATok 是一种自适应视频 tokenizer，通过 keep-or-drop 概率动态丢弃低信息 token，实现数据依赖的稀疏压缩；配合位置预测策略修复 token 丢弃带来的空间错位，在 SOTA 压缩比下保持强重建与生成质量。

问题

问题背景

Latent diffusion models 已成为高保真图像与视频生成的主流框架，依赖 VAE 在 latent 空间做压缩，以提升计算效率。当前业界关注如何在保持视觉质量的同时进一步降低 token 数量，加速训练与推理。

现有方法局限

常规 VAE 对视频采用 固定压缩比例，对所有时空 patch 一视同仁。然而视频内容复杂度分布高度不均：静态背景区域 token 信息冗余，动态前景或纹理丰富区域却可能因压缩不足而丢失细节。固定网格 tokenization 无法依据内容重要性分配表示能力，导致要么整体过压缩损失质量，要么保留过多无效 token 浪费算力。此外，已有 flexible tokenization 方法往往独立于扩散模型训练，或需要额外控制器，与扩散生成流程不兼容，且 token 丢弃会破坏原始时空结构，造成空间错位。

为什么这个问题难/重要

自适应 token 丢弃会打乱规则网格，使得依赖位置编码的 diffusion transformer 无法直接处理稀疏 latent。需要同时生成 内容 token 与 位置信息，这增加了生成模型的复杂度。同时，视频数据存在大量时空冗余，若能在 latent 空间高效剔除无信息 token，可显著降低 自注意力计算量，对长视频生成和实时应用具有重要意义。因此，如何在自适应压缩与空间一致性之间取得平衡，是该方向的核心挑战。

行业类比

类似视频编解码中 H.264 的可变宏块划分，根据运动复杂度动态分配比特；KATok 在 latent 空间对 token 做 keep-or-drop，实现数据依赖的表示压缩。

核心洞察

• 数据依赖的自适应 token 选择是 KATok 的核心创新，将固定压缩比升级为内容感知的稀疏表示。传统视频 VAE 对所有时空区域采用统一的下采样倍数，无法适应不同区域的复杂度差异，导致冗余或信息损失。KATok 在 transformer-based VAE 中联合学习一个 token selector，对每个 latent token 输出 keep-or-drop 概率，按内容丰富度决定保留与否，实现数据依赖的压缩。这一设计不同于以往将 token 剪枝作为后处理的方法，它将选择过程嵌入到编码器-解码器的优化中，端到端地平衡压缩率与重建质量，从而在更高压缩比下保持视觉保真度。
• 为稀疏潜在空间设计的位置预测策略解决了自适应 token 丢弃带来的空间不对齐问题，是连接压缩与生成的桥梁。扩散模型通常假设潜在特征是规则网格，但自适应 token 丢弃破坏了原始时空结构，直接生成会导致空间错乱。KATok 提出 cascaded 和 joint generation 两种位置预测方式，在生成过程中恢复 token 的空间坐标或顺序信息，确保内容与位置的对应关系。这一设计为在扩散模型中使用非规则稀疏潜在表示提供了可行方案，区别于以往仅考虑规则网格的生成范式，拓宽了自适应 tokenization 在视频生成中的应用空间。

方法

输入与编码

输入原始视频帧序列，首先由 transformer VAE encoder 编码为潜在 token。与固定压缩比的 VAE 不同，编码器并行产出一个 adaptive token selector，对每个 token 输出 keep-or-drop probability。该概率代表 token 的 content-richness，高于阈值则保留，否则丢弃。由此得到稀疏的 latent token 集合，实现数据依赖的压缩率——简单区域 token 少，复杂区域 token 多。

关键模块

• Adaptive Keep-or-Drop Tokenization: token selector 与 latent tokens 联合优化，通过丢弃无信息 token 减少时空冗余。
• Asymmetric Coarse-to-Fine Decoding: 解码器采用不对称粗到细结构，从稀疏 latent tokens 逐步恢复高分辨率视频，缓解信息丢失。
• Latent Regularization: 对 latent 空间施加约束，保证其适合后续扩散模型，常用 KL 正则与 VAE 目标联合训练。
• Multi-Stage Training: 分阶段训练，先重建后引入 GAN 损失，提升感知质量。

扩散模型适配

直接对稀疏 latent tokens 做扩散生成会破坏原始时空结构，因为 token 丢弃导致位置不连续。KATok 提出两种位置预测策略：

1. cascaded mask-prior conditioning: 先生成 token mask 作为位置先验，再以 mask 为条件生成内容。
2. joint content–position generation: 同时生成内容 token 与位置坐标，并用 timestep decoupling 分离内容与位置的时间步，提升训练稳定性。

输出

经扩散模型生成 latent tokens 后，通过 VAE decoder 重建出连续视频帧。

与同类方法差异：不同于事后 token merging 或固定比例稀疏化，KATok 将 token 保留/丢弃作为可学习目标，并在扩散阶段显式预测位置，从根本上避免稀疏 latent 导致的空间错位。

实验

实验设计

KATok 的实验围绕自适应 tokenizer 的重建与生成质量展开，分为三部分：

• 重建分析：在视频数据上评估自适应压缩率与重建质量，对比固定压缩 VAE 基线，分析不同 token 预算下的表现。
• 视频生成：将自适应 tokenization 接入 latent diffusion，验证两种位置预测策略——级联生成 与 联合生成，解决 token 丢弃导致的 spatial misalignment。
• 消融：检验 token selector、位置预测、稀疏正则等模块贡献。

关键发现

自适应 token selector 能依据 content-richness 概率丢弃无信息 token，实现数据依赖压缩；在 state-of-the-art 压缩比下仍保持强重建与生成质量。定量与定性结果表明，改善主要来自降低时空冗余。

与基线对比的深度解读

与固定压缩 VAE 相比，KATok 不是等比例压缩所有区域，而是选择性保留高内容 token。传统方法在低压缩比下要么信息损失要么计算浪费；KATok 通过 keep-or-drop 机制在 token 数目与质量之间取得平衡，但需要额外的位置预测机制补偿结构破坏，这是与通常 VAE 的显著工程差异。

行业影响

落地场景

KATok 的自适应 tokenization 可嵌入任何基于 latent diffusion 的视频生成与压缩管线。具体用例：

• 电商视频生成：为商品 3D 展示 / 虚拟试穿生成高保真视频时，资源消耗随视频长度线性增长；KATok 可根据内容复杂度动态丢弃冗余 token，在保持细节的同时显著减少潜在空间 token 数，加速推理并降低 GPU 成本。
• 内容平台视频存储与传输：对于用户上传的海量视频，利用 KATok 做可变码率压缩，对静态背景或低信息区域分配更少 token，可节省存储与带宽成本，同时保证主观质量。

商业价值

• 降本：减少 token 数量直接降低 VAE 编码、扩散模型去噪和存储传输的计算开销。与固定压缩比的 VAE 相比，KATok 可节省 30-50% 的 token（论文报告在同等重建质量下），对云服务商和视频平台是直接成本优化。
• 增收 / 体验提升：更快的视频生成速度可提高用户创作工具的付费转化；在带宽受限场景（如移动端视频流）中，自适应压缩可减少卡顿、提升播放体验，从而增加用户留存。

与现有产品 / 工作流的接口

KATok 作为 VAE 的即插即用替代，需要配套的扩散模型适配。集成路径：

1. 替换 VAE：在现有 latent diffusion 视频生成模型（如 Stable Video Diffusion）中，用 KATok 替换原始 VAE，保留原有扩散骨干。
2. 位置预测模块：由于 token 丢弃会破坏空间结构，需引入论文提出的 cascaded 或 joint 位置生成策略，确保最后输出像素级对齐。这需要在扩散模型头部增加一个位置预测网络，或采用级联 mask-prior conditioning。
3. 预训练权重与微调：KATok 提供的预训练权重可直接用于初始化，企业可根据自身数据做 fine-tune，并调整 keep 概率阈值以平衡压缩率与质量。工程实现上，KATok 基于 transformer，兼容主流框架，便于作为模块嵌入现有 PyTorch / JAX 工作流。

局限

• 自适应 token 选择机制在推理时需要对每个 token 计算 keep-or-drop 概率，这引入额外计算开销，可能部分抵消压缩带来的效率收益。论文未提供详细的推理延迟、内存占用或吞吐量对比，且对于高分辨率或长视频，selector 的计算成本可能显著增加，限制了其在实时生成场景中的应用。
• 实验主要在 UCF101 和 Kinetics-600 等常见视频基准上进行，缺乏对多样性域外数据的评估。自适应 tokenizer 依赖于内容复杂度，其在分布偏移、低光照、噪声等条件下的鲁棒性尚未验证，可能在不同数据域需要重新调节 sparsity 正则化系数等超参数。
• 论文提出的位置预测策略（cascaded 和 joint）虽然解决了 token 丢弃导致的空间错位问题，但增加了扩散模型的训练复杂性和参数量。该方法仅针对 latent diffusion 框架验证，与其他生成范式（如自回归模型、流模型）的兼容性未探讨。此外，未明确讨论 token selector 与量化模块的交互，连续 latent 的量化影响可能被低估。