STAT: Soft Tail-dropping for Adaptive Visual Tokenization¶
会议: ECCV 2026
论文: ECCV 2026
项目主页: https://zeyuan-chen.com/STAT/
领域: 视觉语言模型 / 模型压缩 / 视觉生成
关键词: 自适应视觉分词, 自回归图像生成, 软末端丢弃, 离散表征, 概率分词
一句话总结¶
STAT 提出了一种带有软末端丢弃(Soft Tail-dropping)机制的自适应 1D 离散视觉分词器,通过预测逐 token 保留概率并施加内容自适应与单调递减先验,使分配的 token 数量与图像感知复杂度对齐,在更低平均 token 数下实现了 SOTA 重建并解锁了极具扩展潜力的自回归变长生成。
研究背景与动机¶
自然图像在视觉复杂度上具有天然的多样性:既有由平滑均匀区域构成的极简几何构图,也有包含高频纹理、精细细节和繁杂物体的丰富场景。在传统数字压缩领域(如用于图像的 JPEG、用于视频的 H.264 与 HEVC)以及自然语言处理的大语言模型中,基于内容复杂度自适应分配码率或 token 数量早已是公认的基本范式。然而,在现代主流视觉生成模型(如 VQ-GAN、DiT、LlamaGen)中,图像分词器几乎无一例外地依赖于刚性 2D 规则网格或固定长度的 1D token 序列,强制对所有输入施加完全相同的固定压缩比。
这种固定长度的表达范式造成了严重的表征失衡:简单图像往往被过度分配了冗余的视觉 token,浪费了计算开销并增加了生成建模负担;而纹理高度复杂的图像则受限于固定的 token 容量而无法获得充分细致的重构与表征。尽管近年来已有工作(如 FlexTok、One-D-Piece)尝试利用嵌套丢弃(nested dropout)在训练期学习有序的 1D 序列以支持可变长度重建,但它们在推理时仍严重依赖于外部人工预设的全局固定 token 长度,无法由模型自主判断输入图片的内在复杂度;而少数尝试单遍自适应的分词方案又往往过度依赖外部大模型或复杂的启发式损失,难以与经典的纯因果自回归模型(vanilla causal AR)自然契合。
本文的切入点在于将视觉分词的长度自适应性内生化,让分词器在输出离散编码的同时自主输出反映感知重要度的概率分布。核心 idea:提出软末端丢弃自适应分词器(STAT),利用轻量概率预测头生成逐 token 保留概率,引入内容感知先验(匹配感知重建误差)与递减重要性先验(约束概率自头向尾单调递减),在单遍前向中自适应决定视觉 token 长度,并无缝赋能标准因果自回归模型实现自适应长度采样与卓越的缩放特性。
方法详解¶
整体框架¶
STAT 基于 1D Transformer 架构构建分词器,彻底解耦 2D 空间网格约束。输入图像 \(x \in \mathbb{R}^{B \times C \times H \times W}\) 经由步长为 \(f\) 的 patch embedding 层映射为展平的图像 patch 特征 \(P \in \mathbb{R}^{B \times N \times D}\)。模型将这组 patch 嵌入与 \(L\) 个可学习的 1D 隐式 token \(L \in \mathbb{R}^{B \times L \times D}\) 进行拼接,送入 ViT 编码器中进行跨注意力交互与信息提炼。编码器输出后仅保留 1D 隐式 token 序列 \(z_l\),随后通过向量量化器(Vector Quantizer, VQ)映射为离散 token \(z_q\)。
为了让 token 序列具备内容自适应性,STAT 的编码器后接一个位置感知概率预测头,为每个 token 输出一个保留概率 \(p_{j,i}\)。在训练阶段,模型通过伯努利采样决定哪些 token 被保留,经过直通估计器(STE)进行可微采样与掩蔽,最后由 ViT 解码器结合一组可学习输出 token \(O\) 完成图像的高保真重建。整体生成推理阶段则通过采样保留阈值 \(\tau\) 动态插入终止符(EoS),将离散 token 序列直接对接因果自回归 Transformer。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入图像与可学习 1D 隐式 Token"] --> B["ViT 编码器与特征交互"]
B --> C["向量量化与位置感知保留概率预测"]
C --> D["两阶段概率先验约束<br/>内容感知先验 + 递减重要性先验"]
D --> E["软末端丢弃与直通估计采样掩蔽"]
E --> F["ViT 解码器变长重建"]
E -.->|生成阶段| G["自适应 EoS 截断与因果自回归生成"]
关键设计¶
1. 两阶段渐进式训练与软末端丢弃:打破固定长度与刚性截断的局限 传统固定长度分词器一旦直接在未充分对齐的表征上联合学习自适应丢弃,极易陷入平凡解或训练崩塌。STAT 采用优雅的两阶段训练策略:在第一阶段,采用硬末端丢弃(Hard Tail-dropping),每轮迭代从均匀分布 \(K \sim \mathcal{U}(L_{\min}, L_{\max})\) 中随机采样前缀长度,强制分词器仅使用前 \(K\) 个 token 解码全图,赋予其强大的变长前缀重建能力;在第二阶段,分词器引入位置感知 MLP \(g_\theta\) 预测每个 token 的保留概率 \(p_{j,i} = \sigma(g_\theta(z_l[j,i]))\),并通过独立的伯努利分布采样 \(m_{j,i} \sim \text{Bernoulli}(p_{j,i})\) 构造二值掩码向量。针对伯努利采样不可导的痛点,STAT 采用直通估计器(Straight-Through Estimator, STE)传递梯度,赋予网络在反向传播中优化概率预测的能力。这种设计既保留了前缀解码的有序性,又在训练中引入了随机扰动,显著增强了后续自回归生成对不完美 token 序列的鲁棒性。
2. 双重结构化先验与预算正则化:内生对齐图像感知复杂度与因果因果序 为了引导保留概率 profile 形成符合视觉感知与生成特性的合理分布,STAT 设计了两个互补的物理先验。首先是内容自适应先验(Content-adaptive Prior),该设计针对不同复杂度的图像应当耗费不同码率的直觉,计算批次中每张图的感知重建误差(如 LPIPS)作为其复杂度代理 \(L_{\text{perc}, j}\),并将其与该图的期望 token 总数 \(T_j = \sum_{i=0}^{L-1} p_{j,i}\) 在整个 batch 内计算皮尔逊相关系数: $\(\mathcal{L}_{\text{content}} = 1 - \mathrm{corr}(L_{\text{perc}}, T)\)$ 该损失通过批次内不同难易度样本的相对对比,强力促使简单图保留较少 token,而纹理密集的复杂图获得更多 token,无需复杂的绝对数值归一化。其次是递减重要性先验(Decreasing Importance Prior),针对因果自回归模型从左到右预测的天然因果结构,STAT 惩罚概率曲线沿序列方向的非单调回弹: $\(\mathcal{L}_{\text{decrease}} = \sum_{i=1}^{L-1} \max\big(0, p_{j,i} - p_{j,i-1}\big)\)$ 这一惩罚迫使 token 重要度从头至尾单调递减,促成前部 token 承载骨架全局信息、后部 token 承载局部细节的层次化结构。最后,为了防止概率整体膨胀或坍缩,引入Token 预算正则化,利用伯努利 KL 散度将整图平均保留概率 \(\bar{p}_j = \frac{1}{L} \sum_{i=0}^{L-1} p_{j,i}\) 软约束在目标稀疏度 \(p^* = 0.5\) 附近: $\(\mathcal{L}_{\text{sparse}} = \mathrm{KL}(\mathrm{Bern}(p^*) \parallel \mathrm{Bern}(\bar{p}_j))\)$
3. 自适应阈值 EoS 截断机制:无缝桥接标准因果自回归模型 以往自适应分词方案难以与标准 GPT 风格的因果自回归模型兼容,往往需要改变注意力掩码或引入分层尺度预测。STAT 则利用递减重要性先验保证的近似单调性,巧妙引入终止符(End-of-Sequence, EoS)。在因果 AR 模型的训练中,STAT 随机采样一个概率阈值 \(\tau\),并将首个保留概率低于 \(\tau\) 的位置定义为 EoS 位置:\(t_{\text{eos}} = \min \{i \mid p_i < \tau\}\)。序列在 \(t_{\text{eos}}\) 处插入 EoS 并丢弃后续所有 token。自回归模型仅需按照最标准的最大似然估计进行 Next-Token 预测: $\(p_\phi(q) = \prod_{t=0}^{t_{\text{eos}}} p_\phi(q_t \mid q_{<t})\)$ 在训练期间随机采样 \(\tau\) 相当于提供了一种自适应的数据增强,使 AR 模型能够适应从粗粒度到精细度的不同截断程度;而在推理采样时,AR 模型一旦预测出 EoS 标记即刻停止解码,实现了真正的内容自适应变长生成,无需任何架构上的特殊魔改。
损失函数 / 训练策略¶
分词器第二阶段的总体优化目标综合了多项监督信号: $\(\mathcal{L} = \mathcal{L}_{\text{recon}} + \mathcal{L}_{\text{GAN}} + \mathcal{L}_{\text{VQ}} + \lambda_{\text{content}} \mathcal{L}_{\text{content}} + \lambda_{\text{decrease}} \mathcal{L}_{\text{decrease}} + \lambda_{\text{sparse}} \mathcal{L}_{\text{sparse}}\)$ 其中 \(\mathcal{L}_{\text{recon}}\) 包含像素级 \(\ell_1\) 损失与 LPIPS 感知损失,\(\mathcal{L}_{\text{GAN}}\) 为判别器对抗损失,\(\mathcal{L}_{\text{VQ}}\) 为 codebook 向量量化损失。在 ImageNet 256×256 分辨率下,下采样倍率设为 \(f=16\),隐式 token 序列总长 \(L=256\),码本大小为 4096,码字维度为 12。推理重建时,默认直接取阈值 \(p > 0.5\) 的 token 送入解码器。
实验关键数据¶
主实验¶
在 ImageNet 256×256 图像重建与类别条件生成基准上,STAT 与包括经典 2D 网格分词器、先进 1D 分词器以及各类扩散与自回归模型进行了系统评测。
表 1:ImageNet 256×256 图像重建对比(重建 FID 与 PSNR)
| 分词器类型 | 模型方法 | Token 数量 | 码本大小 (CB) | rFID ↓ | PSNR ↑ |
|---|---|---|---|---|---|
| 2D 网格 | Taming VQ-GAN | 256 | 16,384 | 4.98 | 19.40 |
| 2D 网格 | MaskGIT VQ-GAN | 256 | 1,024 | 2.28 | - |
| 2D 网格 | Open-MAGVIT2 | 256 | 262,144 | 1.17 | 22.64 |
| 2D 网格 | LlamaGen | 256 | 16,384 | 2.19 | 20.79 |
| 1D 序列 | TiTok-S | 128 | 4,096 | 1.71 | 17.80 |
| 1D 序列 | FlexTok | 256 | 64,000 | 1.08 | 17.70 |
| 1D 序列 | One-D-Piece-L | 256 | 4,096 | 1.08 | 19.04 |
| 1D 序列 | STAT (默认 Thr 0.5) | 220 (自适应平均) | 4,096 | 1.15 | 20.22 |
| 1D 序列 | STAT (Thr 0.01) | 230 (自适应平均) | 4,096 | 0.99 | 20.25 |
| 1D 序列 | STAT (Thr 0.01 +10) | 240 (自适应平均) | 4,096 | 0.88 | 20.02 |
表 2:ImageNet 256×256 类别条件自回归生成对比(LlamaGen-XL 775M 参数统一评测)
| 分词器 | 生成模型架构 | 平均 Token 数 | 码本大小 | gFID ↓ | IS ↑ | Precision ↑ | Recall ↑ |
|---|---|---|---|---|---|---|---|
| LlamaGen | 纯因果 AR (LlamaGen-XL) | 256 | 16,384 | 3.39 | 227.1 | 0.81 | 0.54 |
| LlamaGen | 纯因果 AR (LlamaGen-XL) | 576 | 16,384 | 2.62 | 244.1 | 0.80 | 0.57 |
| TiTok-S | 纯因果 AR (LlamaGen-XL) | 128 | 4,096 | 4.90 | 191.7 | 0.77 | 0.56 |
| One-D-Piece-L | 纯因果 AR (LlamaGen-XL) | 256 | 4,096 | 2.99 | 235.1 | 0.81 | 0.59 |
| STAT | 纯因果 AR (LlamaGen-STAT-XL) | 223 | 4,096 | 2.36 | 244.0 | 0.78 | 0.62 |
消融实验¶
为了验证软末端丢弃机制、自适应分配及训练先验的必要性,作者在统一的 775M 参数模型上进行了系统消融(见表 3)。
表 3:训练阶段丢弃策略与 EoS 阈值设计的消融对比
| 训练丢弃配置 | 重建 Token 数 | rFID ↓ | 重建 PSNR ↑ | 生成 Token 数 | gFID ↓ | IS ↑ | Precision ↑ | Recall ↑ |
|---|---|---|---|---|---|---|---|---|
| 固定 Token 分配 (Fixed Token) | 220 | 1.15 | 20.35 | 220 | 2.73 | 234.2 | 0.77 | 0.62 |
| 硬末端丢弃 (Hard Tail-dropping) | 222 | 1.15 | 20.30 | 223 | 2.67 | 231.2 | 0.76 | 0.63 |
| 固定 EoS 阈值 (\(\tau = 0.5\)) | 220 | 1.15 | 20.22 | 221 | 2.49 | 240.7 | 0.79 | 0.61 |
| 软末端丢弃完整方案 (STAT) | 220 | 1.15 | 20.22 | 223 | 2.36 | 244.0 | 0.78 | 0.62 |
关键发现¶
- 自适应性显著促进生成性能:尽管固定 token 方案在重建指标上因任务简单而表现相当(rFID 1.15 对 1.15),但在生成任务中,自适应 token 分配将 gFID 从 2.73 显著压低至 2.36。这证实了允许生成模型根据所生成画面的复杂度自主决定分配的 token 长度,能大幅提升 token 建模效率并释放生成潜力。
- 软丢弃相比硬截断更具泛化鲁棒性:硬末端截断方案在生成上的 gFID 仅为 2.67,落后于软末端丢弃的 2.36。软丢弃在训练期间引入的独立伯努利概率扰动,有效扮演了序列正则化的角色,使得解码器和后续 AR 模型对生成过程中的微小误差更具容忍度。
- 卓越的模型容量扩展性(Scaling Law):随着因果自回归模型从 775M(XL)扩展至 1.4B(XXL)再到 3.1B,LlamaGen-STAT 的 gFID 从 2.36 持续降至 1.91 和 1.75(Recall 提升至 0.64);相比之下,原版使用 256 tokens 的 LlamaGen 在 1.4B(gFID 3.09)扩展至 3.1B(gFID 3.06)时已陷入严重饱和。STAT 使简单的纯因果自回归模型重新获得了可比拟甚至超越顶级扩散模型的缩放红利。
- 与文件尺寸高度共振的复杂度对齐:可视化分析证实,STAT 预测的保留 token 数量与图像经标准 JPEG 压缩后的文件尺寸(KB)呈现极高的正相关性(从简单平滑图的 ~175 tokens / 63 KB 平滑增长至复杂纹理图的 256 tokens / 164 KB),表明模型真正学到了图像的感知信息熵。
亮点与洞察¶
- 将表征压缩比内生化为概率分布:传统的自适应压缩多依赖外部规则、迭代重推理或大模型打分,而 STAT 仅通过位置感知概率头与感知损失的相关性惩罚(\(\mathcal{L}_{\text{content}}\)),便在单遍前向中实现了与视觉复杂度高度一致的 token 预算分配。
- 软末端丢弃与自回归因果特性的深度共鸣:通过引入单调递减先验(\(\mathcal{L}_{\text{decrease}}\)),STAT 将 1D token 序列排列成“粗粒度语义 \(\to\) 细粒度高频细节”的信息级联,天然贴合自回归模型由前向后的上下文预测机制。
- 极简的 EoS 机制解放纯因果生成模型:无需像 VAR 引入层次化尺度重构,也不用像 MAR 引入连续扩散头,STAT 仅凭一个动态采样的 EoS 符号,就让最朴素的标准 GPT 式因果 Transformer 达成了 SOTA 级生成效果与变长采样能力。
局限与展望¶
- 极端长尾复杂场景下的截断保守倾向:虽然预算正则项设定了全局平均稀疏度目标 \(p^* = 0.5\),但在全景复杂高密图像(如密集人群、密集小文字标注)中,模型趋向于将几乎所有 token 的概率都推至接近 1.0,导致序列长度逼近最大截断上限(\(L=256\)),未能在极端复杂图像上展现超出上限的动态扩展能力。
- 训练收敛对超参数平衡的敏感性:损失函数同时融合了重构、对抗、VQ 量化、内容相关性、递减先验和稀疏度正则,多任务目标较多,需要精细微调 \(\lambda_{\text{content}}\) 和 \(\lambda_{\text{decrease}}\) 等权重系数,否则可能出现概率退化或概率抖动。
- 未来方向:跨模态全自适应序列大一统:未来可将这种逐 token 自适应保留机制推广至视频多帧联合表征以及视觉-语言端到端统合大模型中,使多模态输入能依据信息密度自适应动态折叠与展开。
相关工作与启发¶
- vs LlamaGen (Sun et al., 2024):LlamaGen 确立了使用离散 2D VQ 分词器进行因果自回归的可行性,但依赖 256 或 576 的固定长度网格,不仅在简单图像上浪费计算,且在扩展到 3B 参数时生成质量过早饱和;STAT 采用 1D 自适应分词,用更少 token(平均 223)将 3B 模型 gFID 进一步提升至 1.75,彻底打破扩展瓶颈。
- vs TiTok & FlexTok (Yu et al., 2024; Bachmann et al., 2025):TiTok 将图像压缩为固定紧凑的 1D token;FlexTok 利用 nested dropout 实现了变长前缀重建,但测试时必须由人工显式指定全局 token 数量;STAT 则通过概率预测实现了图像级的内容自适应动态决策。
- vs ALIT & KARL (Duggal et al., 2024/2025):ALIT 依赖循环多遍前向推理寻找最小 token 集合,推理开销成倍放大;KARL 依赖外部损失启发式停止规则;STAT 实现了真正的端到端单遍前向预测,兼顾极高的推理效率与自回归生成的天然适配性。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ 提出软末端丢弃与内容/递减双重先验,设计巧妙直击痛点
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 ImageNet 重建、类别条件生成、GenEval 文生图、512 分辨率与 UCF-101 视频扩展
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密,图文对应清晰,方法推导与动机剖析极其深入
- 价值: ⭐⭐⭐⭐⭐ 为纯因果自回归视觉生成重新赋予了强大的扩展生命力,范式具有广泛参考价值