KATok：按内容复杂度自适应取舍令牌的紧凑视频表示分词器 Keep-or-Drop? Adaptive Tokenizer for Compact Video Representation
可微分令牌取舍机制让视频 VAE 按内容复杂度自适应压缩，生成更快更省
前置知识
潜在扩散模型（LDM）
在 VAE 压缩出的低维潜在空间而非像素空间执行扩散去噪的生成范式。像素空间扩散要处理数百万维数据，而潜空间维度通常小一到两个数量级，使高分辨率图像与视频的迭代去噪在计算上可行，Stable Diffusion 等系统均属此类。
本文的 VAE 分词器正是为 LDM 提供潜在空间，其压缩质量与 token 数直接决定下游生成器的训练成本和最终保真度。
变分自编码器（VAE）/ 视觉分词器
把图像或视频压缩为低维 latent 的编码-解码网络。编码器输出高斯分布参数并经重参数化采样 $\hat z_i=\mu_i+\sigma_i\odot\epsilon_i$，解码器从 latent 重建输入；训练用重建、KL 散度、感知与对抗损失联合优化。
KATok 本体就是一个 Transformer VAE，理解重参数化与 KL 项才能看懂它的训练目标与消融实验。
Gumbel-Softmax 松弛
把离散类别采样（如保留/丢弃二选一）松弛为可微操作的技巧：给 logits 加 Gumbel 噪声后做带温度 $\tau$ 的 softmax，训练时梯度可传，$\tau\to 0$ 时逼近硬性 one-hot 采样。
keep-or-drop 本质是离散决策，KATok 靠它端到端学习 token 掩码；消融显示去掉它训练会直接坍缩到只剩 2 个 token。
流匹配（Flow Matching）
一类连续归一化流的生成训练框架：学习速度场 $v_\theta(z_t,t)$，把高斯噪声 $z_1$ 沿线性插值路径 $z_t=(1-t)z_0+t z_1$ 输运到数据 $z_0$，损失为预测速度与真实速度 $(z_1-z_0)$ 的均方误差。
KATok 训练下游视频生成器用的就是流匹配加 SiT-XL；内容-位置联合预测的位置损失也是在此框架内扩展的。
rFVD / gFVD 指标
FVD（Fréchet Video Distance）在预训练特征空间中度量两个视频分布的 Fréchet 距离，越低越好；rFVD 用于重建任务（真实视频 vs 重建视频），gFVD 用于生成任务（真实分布 vs 生成分布），两者都依赖约数千样本估计。
论文所有质量结论都锚定在这两个指标上，例如重建 5.12 rFVD、UCF-101 生成 61.53 gFVD，是横向对比的核心标尺。
3D RoPE 与 register tokens
RoPE 把相对位置以旋转矩阵形式注入注意力打分，3D 版本分别编码时间与两个空间维度的坐标；register tokens 是不携带具体位置的内容无关全局 token，充当信息汇聚的寄存器，能缓解注意力伪影并稳定全局结构。
消融显示去掉 register tokens 后 token 用量从 365.57 涨到 410.20、PSNR 降至 29.17，它们是激进稀疏化下保持空间连贯的锚点。
研究动机
潜在扩散模型（LDM）已成为高保真图像与视频合成的主流框架，其计算效率高度依赖前端 VAE 分词器的压缩能力。但主流视觉分词器（LDM 的卷积 VAE、DC-AE、LTX-Video、OmniTokenizer 等）都采用固定压缩比：无论视频是静态天空还是剧烈运动场景，编码出的 token 数完全由时空分辨率决定。以 OmniTokenizer-VAE 为例，一段 256²×17 的视频要编码成 5120 个 token，512²×33 更膨胀到 36864 个。然而视频时空冗余极强，大量 token 落在均匀、静止、无信息的区域，白白占用下游扩散模型的训练与推理容量。已有可变长方案（FlexTok、One-D-Piece、SEED、ALIT、ElasticTok）虽允许用户指定 token 数，但本质是“用户可控”而非“自适应”：每个样本的最优 token 数仍需额外模型或推理期搜索（ElasticTok 采用 KL 模型加逐样本二分搜索），给下游任务引入不可忽视的开销。
本文的目标是本文的目标是构建一个端到端可微的自适应视频分词器：模型自己学会判断每个样本需要多少 token（有效 token 数 $N_{eff}(X)=\sum_i m_i(X)$），自动丢弃冗余、无信息的 token，从而在不预设 token 预算、不做推理期搜索的前提下实现数据依赖的压缩。具体而言，作者希望新分词器在重建指标（PSNR、SSIM、LPIPS、rFVD）上超越固定长度的 Transformer VAE 基线 OmniTokenizer 与灵活长度基线 ElasticTok，同时把 token 数降低一个数量级；并且能作为扩散模型的潜在空间，训练出收敛更快、质量更高的生成器，并解决稀疏 token 带来的内容-位置错位问题，最终实现视频生成训练提速约 6.9 倍、推理提速约 3.2 倍。
与已有工作不同的是，作者的独特切入是严格区分“灵活”与“自适应”两种范式：灵活分词器把决定权留给用户、仍需搜索；KATok 则把 token 数变成训练的产物。技术上，它在 Transformer 编码器内部对每个连续 latent token 做可微的 keep-or-drop 门控（Gumbel-Softmax 松弛），并把同一份软掩码以 $b_j=\log(\tilde m_j+\epsilon)$ 的形式加到解码器注意力 logits 上，保证“丢 token”真正等价于“丢信息”。针对稀疏化导致的内容-位置错位，它不靠简单增容，而是提出内容-位置联合预测（带时间步解耦的双噪声调度）与级联 mask 先验两条互补路线。这与推理期剪枝（DynamicViT、EViT、ToMe）和尾部截断式的弹性分词有本质区别：KATok 是训练期、端到端、数据依赖的压缩。
核心方法
直觉上，视频里绝大多数 token 是冗余的：静态背景、均匀区域几乎不携带信息，真正需要表征容量的是运动与纹理复杂的位置。KATok 的思路因此是“先编码、再打分、后丢弃”：视频切成 16²×8 时空 patch，经带 3D RoPE 和 2 个 register tokens 的 Transformer 编码器得到中间表征；每个 patch 经重参数化 $\hat z_i=\mu_i+\sigma_i\odot\epsilon_i$ 产生连续 latent；轻量网络 $g_\theta$ 对每个 token 输出 keep/drop logits，经 Gumbel-Softmax 得到软掩码 $\tilde m_i$，以 $z_i=\tilde m_i\cdot\hat z_i$ 软门控。软掩码同时调制解码器注意力，解码端用 FLUX 式双流块与更细的 8²×4 查询网格重建画面。总损失 $\mathcal{L}=\mathcal{L}_{recon}+\lambda_{KL}\mathcal{L}_{KL}+\lambda_{sparse}\mathcal{L}_{sparse}+\lambda_{adv}\mathcal{L}_{adv}+\lambda_{align}\mathcal{L}_{align}$ 中，ℓ1 稀疏正则配合先弱后强的退火调度，让有效 token 数 $N_{eff}$ 随训练自然下降。在 Panda-70M 上三阶段训练，下游配 SiT-XL 流匹配生成。
核心创新有三层。第一，把 token 数从人为超参数变成可微优化的对象：keep-or-drop 概率与 latent 联合学习，稀疏正则 $\mathcal{L}_{sparse}=\mathbb{E}_X\sum_i\tilde m_i(X)\approx\mathbb{E}_X|N_{eff}(X)|$ 实现逐样本压缩，无需 ElasticTok 式二分搜索或手工 dropout 调度。第二，软注意力掩码机制：同一 $\tilde m$ 既门控 latent 又以 $b_j=\log(\tilde m_j+\epsilon)$ 移位解码器注意力 logits，这是训练不坍缩的关键——消融显示去掉它只剩 2 个 register tokens 活跃（PSNR 19.00、SSIM 0.54）。第三，针对稀疏 latent 的内容-位置错位，提出级联 mask 先验：8.3M 参数的 SiT-Tiny 先生成 token 掩码，阈值化出位置后以位置嵌入条件化主流匹配模型，把“选位置”与“生内容”解耦，比联合生成（73.16 gFVD）更稳更好（61.53）。非对称粗编码（16²×8）细解码（8²×4）则是效率与细节的巧妙折中。
方法步骤详情
流程分八步。(1) 切片编码：视频 $X\in\mathbb{R}^{T\times H\times W\times C}$ 切成 16²×8 patch，线性投影后经带 3D RoPE、2 个 register tokens 的编码器得到 $e_i$。(2) 重参数化：$f_\theta$ 把 $e_i$ 映射为 $(\mu_i,\sigma_i)$，采样 $\hat z_i=\mu_i+\sigma_i\odot\epsilon_i$。(3) 重要性打分：$g_\theta$ 输出 logits $\alpha_i\in\mathbb{R}^2$，训练时经 Gumbel-Softmax（温度 $\tau$）得 $\tilde m_i$，推理时取硬掩码 $m_i=\mathbb{I}(\alpha_{i0}\ge\alpha_{i1})$。(4) 门控与正则：$z_i=\tilde m_i\hat z_i$，latent 注入 [0,0.2] 噪声并加冻结 vJEPA-2 特征的对齐损失 $\mathcal{L}_{align}$。(5) 解码：可学习 query（8²×4 细网格、带 3D RoPE）与 latent 在 FLUX 式双流块交互，注意力 logits 加 $b_j=\log(\tilde m_j+\epsilon)$ 偏置，latent 不带位置编码。(6) 优化：Video-LPIPS（S3D 特征）、KL、对抗与退火稀疏损失联合训练。(7) 三阶段：256²×16 单分辨率→多分辨率→GAN 微调。(8) 下游生成：SiT-XL（686.29M）流匹配，级联方案用 8.3M mask 先验生成掩码并提供位置嵌入条件，训练时以活跃 token 数作附加条件。
技术新颖性
与已有工作的本质区别有四点。其一，推理期剪枝方法（DynamicViT、EViT、ToMe）只在推理时省计算，latent 仍稠密；KATok 在训练期就学到稀疏 latent，训练与推理双双受益。其二，灵活长度分词器（FlexTok、One-D-Piece、SEED、ALIT、ElasticTok）依赖预设预算、dropout 调度或推理期搜索；KATok 的 Gumbel-Softmax 门控让 token 数成为连续可微优化的自然输出，无预算、无搜索。其三，“位置先验”是全新问题：作者首次系统指出自适应稀疏化会让扩散模型无法恢复 token 的原始时空位置，并用双噪声调度联合生成与级联 mask 先验两种方案解决，gFVD 从 naive 的 95.69 降到 61.53。其四，Video-LPIPS、latent 噪声与 vJEPA-2 对齐的组合体现了“轻微牺牲重建、大幅换取生成”的设计哲学——去掉 latent 正则虽让重建 PSNR 从 30.85 升到 31.26，但下游 gFVD 从 101.23 恶化到 161.23，经验上这种权衡非常划算。
实验结果
重建（Table 1，Panda-70M 验证集约 5000 段视频）：256²×16 下 KATok 仅用 366.24 个 token（64 通道），取得 31.24 PSNR / 0.04 LPIPS / 0.94 SSIM / 5.12 rFVD，压缩比 134.21；而 Omni-VAE 用 5120 token 只有 28.10 PSNR / 7.84 rFVD，ElasticTok-KL 用 3845.56 token 只有 30.52 PSNR / 12.37 rFVD。512²×32 下优势更大：33.23 PSNR / 6.40 rFVD（1554 token，压缩比 253.00）对 Omni 的 24.07 / 16.85（36864 token）。定性上（Fig. 6）纯白视频仅需 28 个 token，海洋片段 1104→836，室内场景 1712→330。Fig. 7 显示 token 数与时间熵相关性最强（r=0.865），高于空间熵（0.618）。生成（Table 3，SiT-XL，gFVD）：UCF-101 上 Cascaded 61.53 对 Omni 100.00、Elastic-KL 712.56；Kinetics-600 上 160.84 对 206.58；SkyTimelapse 上 23.19（Joint 最佳 21.36）对 23.28，且生成 token 少约 11 倍。Table 4 显示 naive 95.69 → Joint 73.16 → Cascaded 61.53 递进；去 latent 正则 gFVD 从 101.23 恶化到 161.23。效率（Fig. 2）：80k 步达 73.81 gFVD 超过 Omni 200k 步的 82.31（加速 6.9×），比 ElasticTok 快 46.7×；推理吞吐 15.71 视频/秒，为 Joint 的 3.1×、Omni 的 3.2×、ElasticTok 的 3.7×。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 视频重建（Panda-70M 验证集，256²×16） | PSNR↑ / rFVD↓ | 31.24 / 5.12（366.24 token，压缩比 134.21） | OmniTokenizer-VAE：28.10 / 7.84（5120 token）；ElasticTok-KL：30.52 / 12.37（3845.56 token） | 较 Omni +3.14 PSNR，rFVD 降低 34.7%，token 减少 92.9% |
| 视频重建（Panda-70M，512²×32） | PSNR↑ / rFVD↓ / 压缩比 | 33.23 / 6.40 / 253.00（1554.24 token） | OmniTokenizer-VAE：24.07 / 16.85 / 96（36864 token） | +9.16 PSNR，rFVD 降低 62%，压缩比提升约 2.6 倍 |
| 类条件视频生成（UCF-101） | gFVD↓ | 61.53（Cascaded） | Omni-VAE 100.00；ElasticTok-KL 712.56 | 较 Omni 降低 38.5%，生成 token 减少约 11 倍 |
| 类条件视频生成（Kinetics-600） | gFVD↓ | 160.84（Cascaded） | Omni-VAE 206.58 | 降低 22.1% |
| 无条件视频生成（SkyTimelapse） | gFVD↓ | 23.19（Cascaded）/ 21.36（Joint） | Omni-VAE 23.28；ElasticTok-KL 95.53 | Joint 最优，Cascaded 持平略优；较 ElasticTok 提升 4.1 倍 |
| 生成训练效率（UCF-101，8×H200） | 达到超基线质量所需 wall-clock | 80k 步 73.81 gFVD（较 Omni 6.9× 加速）；30k 步 203.51（较 ElasticTok 46.7× 加速） | OmniTokenizer 200k 步 82.31；ElasticTok 200k 步 571.41 | 训练加速 6.9× / 46.7× |
| 生成推理吞吐（DOPRI5 求解器） | videos/sec | 15.71 | Ours-Joint 5.01；OmniTokenizer 4.91；ElasticTok 4.24 | 3.1× / 3.2× / 3.7× |
局限与改进
作者承认的局限：其一，latent 正则（噪声注入 + vJEPA-2 对齐）会轻微损害重建指标（Table 2 中 PSNR 从 31.26 降到 30.85），是刻意用重建换生成的权衡；其二，联合生成方案对内容与位置两条噪声调度的超参非常敏感，需仔细调参，因此默认采用级联方案；其三，受资源限制，ElasticTok 只在 SkyTimelapse 与 UCF-101 上对比，且需离线分块与二分搜索，对比存在打折；部分指标只在 16/32 帧上评估。我自己的观察：级联方案引入 mask 先验这一额外环节，先验在分布外场景（罕见运动模式、极端分辨率）失准时错误会传播到内容生成；自适应 token 数意味着 batch 内序列长度不等，工程上需要分桶或 padding，实际吞吐收益可能低于报告值；所有生成实验限于无条件与类条件的 256²×16 短视频，未验证文本到视频等大规模场景；rFVD/gFVD 依赖参考分布，跨数据集可比性有限；与 64 通道 latent 的对比虽按压缩比定义占优，但与 8 通道基线的逐比特对比仍需谨慎解读。
独立分析的弱点
第一，动态 token 数带来系统层面复杂性：batch 内 $N_{eff}$ 差异极大（28 到数千），训练与部署需要动态 padding 或分桶调度，实际吞吐提升可能低于论文数字；改进方向是引入软预算上下界约束或按复杂度分桶。第二，硬掩码丢弃不可逆：推理时丢掉的细节无法恢复，对精细纹理任务不利；可保留极低维残差 token 或引入分层细化机制按需补回。第三，级联 mask 先验是单点故障：其预测错误的时空位置会系统性误导主模型，可改为多先验集成、与内容模型联合微调，或用不确定度加权的位置嵌入。第四，评估规模偏小：SiT-XL 只在三个短视频数据集（SkyTimelapse、UCF-101、Kinetics-600）上训练 10 万步，尚不清楚优势能否迁移到文本到视频与长时序生成，应在更大规模设置下复验。第五，$N_{eff}$ 与熵的相关（r≈0.87）只是统计关联，缺乏信息论层面的理论刻画，未来可建立自适应压缩率与内容复杂度之间的下界分析。
未来方向
作者方向：论文展示了 token 数可作为推理期的运动/细节旋钮——Fig. 5 中 200 token 样本联合熵 2.61（低运动）对 400 token 的 5.09（强运动），无需重训练即可调制动态性，这一涌现可控性有望用于交互式生成与运动强度编辑；级联 mask 先验也可推广为更通用的结构先验（语义、深度引导的稀疏化）。基于成果可延伸：把 KATok 嵌入 LTX-Video 级别的文本到视频大模型，验证在文本条件、高分辨率、长视频下的加速收益；把掩码先验与生成模型结合为“先规划时空构图、再生成内容”的两阶段范式；利用 $N_{eff}$ 作为免费的视频复杂度度量，服务数据清洗、课程学习与训练算力分配；与神经视频编解码结合，探索压缩与生成统一的编码器；以及在多模态统一 token 空间中让文本、图像、视频 token 共享自适应机制。
复现评估
复现难度中等偏上。有利因素：所用数据全部公开（Panda-70M、SkyTimelapse、UCF-101、Kinetics-600）；项目页 kakao.github.io/KATok/ 存在，论文称完整实现细节与超参在补充材料中；架构组件均为公开技术（ViT-B 规模 Transformer、3D RoPE、register tokens、FLUX 双流块、SiT-XL、流匹配），无黑盒模块；mask 先验仅 8.3M 参数，可作为最小验证单元先行复现。不利因素：截至论文未见明确的代码与权重开源承诺；算力门槛高——tokenizer 三阶段训练，生成器 SiT-XL（686.29M）在 8 张 H200、全局 batch 256 下训练 20 万步；Gumbel-Softmax 温度、$\lambda_{sparse}$ 退火节奏、双噪声调度等超参敏感，消融显示去掉软注意力掩码或 Gumbel-Softmax 会训练坍缩（只剩 2 个 token、PSNR 约 19），意味着复现容错很低。建议先复现单分辨率 stage-1（100K 步）加 UCF-101 小规模生成，验证 61.53 gFVD 量级后再扩展。
论文图表
