跳转至

Balancing Image Compression and Generation with Bootstrapped Tokenization

会议: NeurIPS 2026
arXiv: 2606.05552
OpenReview: https://openreview.net/forum?id=1fSMXPTPyJ
领域: 模型压缩
关键词: 图像标记化, 自启动学习, 最优传输, 全局局部解耦, 图像生成

一句话总结

提出自启动图像标记器 SelfBootTok,将图像标记显式解耦为紧凑的全局语义标记与多粒度局部标记,通过自监督自启动机制与最优传输对齐从全局标记直接预测局部细节,使得下游生成模型仅需建模 64 个全局标记即可达到 1.56 gFID,同时节约约 40% 的训练算力。

研究背景与动机

基于扩散模型(如 DiT、SiT)和自回归模型的大规模视觉生成在近年来取得了突飞猛进的发展。为了降低生成模型的计算复杂度并支持统一多模态表示,图像标记化(Image Tokenization)成为了关键的基础设施。早期工作大多沿用 VQ-VAE 或 SD-VAE 的二维网格表征(2D grid latents),但其严格的空间网格约束导致序列长度居高不下(通常为 256 甚至更高),难以在极低标记预算下高效压缩。为此,近期涌现出以 TiTok、SoftVQ 为代表的一维(1D)图像标记器,试图将图像压缩进极短的一维标记序列中,为生成任务带来巨大的序列效率提升。

然而,现存的一维图像标记方法面临着严重的表征与生成矛盾。现有一维标记器在压缩时通常不区分特征粒度,使得每个标记都混合了高层语义与底层纹理细节。这种混杂编码不仅导致标记之间存在大量冗余重叠,更把沉重的视觉细节重建负担强加给了下游生成模型。下游生成器(如 DiT 或自回归 Transformer)在接收文本或类别条件后,被迫在极其有限的标记序列上同时学习全局布局与细粒度外观,造成优化难度激增、收敛缓慢,并且在扩大标记器容量时必须完全重新训练庞大的生成模型。

针对这一瓶颈,本文切入的角度是改变视觉标记的职责分配:既然高层语义决定了场景骨架,细粒度细节应当可以从全局语义中自驱动涌现,而不应全由生成模型硬背。核心 idea:将图像标记显式解耦为全局语义标记与局部细节标记,利用自监督自启动学习直接从全局标记预测多粒度局部细节,并将细粒度视觉渲染的负担从生成器卸载至标记器内部,使生成器仅需专注于学习超紧凑的全局标记。

方法详解

整体框架

SelfBootTok 的核心设计思想是“全局负责语义、局部负责细节、自启动解耦生成”。输入图像首先由 Vision Transformer 编码器提炼为极简的一维全局标记序列 \(\mathbf{z}_G\)。随后,系统进入自监督自启动预测阶段(Self-supervised Bootstrapping),分别通过轻量 MLP 预测局部 1D 标记 \(\mathbf{z}_{L1}\)、通过因果 Transformer 预测局部 2D 标记 \(\mathbf{z}_{L2}\)。这些局部标记分别借助预训练视觉编码器(如 I-JEPA、DINOv2、SigLIP2)以及最优传输(Optimal Transport, OT)与真实局部结构紧密对齐。最后,全局与局部标记经过可微软矢量量化(Soft VQ)后进行融合并送入 ViT 解码器重建图像。在下游生成时,生成模型(如 SiT)仅需生成全局标记 \(\mathbf{z}_G\),局部标记完全由标记器自启动补全。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["输入图像 I"] --> B["ViT 编码器"]
    B --> C["全局标记提取:生成超紧凑全局语义序列 z_G"]
    C --> D["自监督自启动预测:双路局部标记预测"]
    D --> E["局部结构对齐与最优传输:OT 损失与表征对齐"]
    E --> F["多码本解耦量化与特征融合:软矢量量化与解码"]
    F --> G["图像重建 / 下游仅生成全局标记"]

关键设计

1. 全局标记提取与多码本解耦量化:分离不同粒度表征空间 传统单一码本或扁平表征强迫离散或连续向量承载所有尺度的视觉信息。SelfBootTok 首先使用基于 DINOv2 初始化的 ViT 骨干网络将图像编码为极度紧凑的全局标记 \(\mathbf{z}_G \in \mathbb{R}^{B \times N_z \times d}\)(例如仅 64 个标记)。为杜绝不同粒度信息在码本空间中互相干扰,模型引入了显式分离的多码本集合: $\(\mathcal{C}_{\mathrm{tot}} = \{\mathcal{C}_{\mathrm{global}}, \mathcal{C}_{\mathrm{local1D}}, \mathcal{C}_{\mathrm{local2D}}\}\)$ 为了实现端到端梯度稳定回传,量化过程采用全可微的软矢量量化(Soft VQ)。对于任意标记组 \(\hat{Z} \in \mathbf{z}_{\mathrm{tot}}\),其在对应码本上的后验概率分布与量化向量定义为: $\(q_{\phi}(Z \mid I) = \mathrm{Softmax}\left(-\frac{1}{\tau}\|\hat{Z} - \mathcal{C}\|_2\right), \quad Z = q_{\phi}(Z \mid I)\,\mathcal{C}\)$ 通过为全局与局部标记分配独立的码本,系统强制每个分支专精于自身尺度的信息表示,消除了标记组之间的冗余交互与语义污染。

2. 自监督自启动预测:从全局语义自发推导局部细节 为了将生成模型从琐碎的像素细节建模中解放出来,模型构建了自启动预测模块(Bootstrapped Tokenization)。给定全局标记 \(\mathbf{z}_G\),模型并行分支推导局部结构:针对局部 1D 标记 \(\mathbf{z}_{L1}\),利用轻量多层感知机直接将全局特征投影到扩展的局部标记空间;针对局部 2D 标记 \(\mathbf{z}_{L2}\),则使用一系列可扩展的因果 Transformer 模块(Causal Transformer Blocks)进行自回归序列建模: $\(p_{\theta}(\mathbf{z}_{L2} \mid \mathbf{z}_G) = \prod_{i=1}^{N_l} p_{\theta}\big(\mathbf{t}^{\mathrm{local}}_{i} \mid \mathbf{z}_G, \mathbf{t}^{\mathrm{local}}_{<i}\big)\)$ 这一设计使生成器在测试时只需专注于建模紧凑的 \(\mathbf{z}_G\) 序列,下游采样完成后,由冻结的标记器局部对齐器自动“脑补”高保真细节,从而把生成难度由原本的全局+局部联合建模简化为单纯的全局流形采样。

3. 局部结构对齐与最优传输:无监督几何先验注入 自启动预测出的局部标记必须与真实的局部几何结构相锚定,否则会导致幻觉或细节模糊。为此,作者提出了双重对齐策略:首先利用预训练基础模型(DINOv2、I-JEPA、SigLIP2)提供补丁级语义监督 \(\mathcal{L}_{\mathrm{repa}}\)。更关键的是,为解决二维空间潜变量(2D VAE latents \(\mathbf{X} \in \mathbb{R}^{B \times N_x \times d}\),其中 \(N_x = H \times W\))向一维离散序列 \(\mathbf{z}_{L2} \in \mathbb{R}^{B \times N_z \times d}\) 映射时的空间拓扑丢失问题,引入了熵正则化最优传输(Optimal Transport, OT)。定义欧氏距离传输代价矩阵 \(\mathbf{C}_{ij} = \|\mathbf{x}_i - \mathbf{z}_j\|_2^2\),求解最优传输方案: $\(\mathbf{P}^* = \arg\min_{\mathbf{P} \in \Pi(\mathbf{a}, \mathbf{b})} \langle\mathbf{P}, \mathbf{C}\rangle + \varepsilon \mathcal{H}(\mathbf{P})\)$ 借助 Sinkhorn 算法快速迭代求解唯一的最佳传输平面 \(\mathbf{P}^* = \mathrm{Sinkhorn}(\exp(-\mathbf{C}/\varepsilon), \mathbf{a}, \mathbf{b})\),最终以最优传输代价 \(\mathcal{L}_{\mathrm{OT}} = \langle\mathbf{P}^*, \mathbf{C}\rangle\) 作为损失项回传梯度,实现了 2D 连续流形到 1D 紧凑序列之间几何一致性与保真度的无监督精准转移。

损失函数 / 训练策略

整个标记器的端到端优化目标由五部分构成: $\(\mathcal{L} = \mathcal{L}_{\mathrm{recon}} + \lambda_1 \mathcal{L}_{\mathrm{percep}} + \lambda_2 \mathcal{L}_{\mathrm{adv}} + \lambda_3 \mathcal{L}_{\mathrm{align}} + \lambda_4 \mathcal{L}_{\mathrm{KL}}\)$ 其中判别器采用 StyleGAN 架构,结合 LeCAM 稳定正则化(权重 0.001)与一致性正则化;对齐损失涵盖全局语义对齐、局部表示对齐与最优传输损失: $\(\mathcal{L}_{\mathrm{align}} = \mu_1 \mathcal{L}_{\mathrm{repa}}^{\mathrm{global}} + \mu_2 \mathcal{L}_{\mathrm{repa}}^{\mathrm{local}} + \mu_3 \mathcal{L}_{\mathrm{OT}}\)$ 在训练策略上,针对大容量局部对齐器难优化的问题,提出两阶段解耦扩展策略:阶段一使用轻量局部 2D 对齐器联合训练全局骨干网络与生成器(SiT),训练 50 个 epoch 并完成生成器训练(800K 步);阶段二冻结 ViT 编码器与生成器,仅扩展局部 2D 对齐器的参数(从 10M 扩展至 600M)并微调 50-75 个 epoch。由于下游生成器只依赖全局标记 \(\mathbf{z}_G\),生成器完全无需重新训练,实现了标记器与生成器的解耦并行优化。

实验关键数据

主实验

在 ImageNet 256×256 基准上,评估一维标记器、二维标记器以及自回归标记器在重建(rFID)与下游生成质量(gFID,基于 SiT / DiT)上的综合表现:

标记器类型 方法 标记数量 (#Tokens) 下游生成器 重建 rFID ↓ gFID (无 CFG) ↓ gFID (含 CFG) ↓
2D Tokenizer VAE 4096 LDM-4 0.27 10.56 3.60
2D Tokenizer SD-VAE 256 DiT-XL/2 0.62 9.62 2.27
2D Tokenizer SD-VAE 256 SiT-XL/2 0.62 8.61 2.06
2D Tokenizer VA-VAE 256 LightningDiT 0.28 2.17 1.35
AR Tokenizer MaskGIT 256 MaskGIT 2.28 6.18 –
AR Tokenizer MAR-H 256 MAR-H 1.22 2.35 1.55
1D Tokenizer TiTok-B 64 MaskGIT-ViT 1.70 – 2.48
1D Tokenizer TiTok-S 128 MaskGIT-UViT-L 1.71 – 1.97
1D Tokenizer FlexTok d18-d28 1~256 LlamaGen 1.45 – 1.86
1D Tokenizer HieraTok 256 DiT 0.45 3.53 1.82
1D Tokenizer GigaTok XL-XXL 256 LlamaGen XXL 0.79 – 1.98
1D Tokenizer MAETok 128 LightningDiT 0.48 – 1.73
1D Tokenizer SoftVQ 64 SiT 0.88 5.98 1.78
1D Tokenizer SelfBootTok (本文) 64 SiT 0.66 2.22 1.56

消融实验

表 A 展示了核心设计组件(因果建模 C、独立码本 CB、局部 1D/2D 分支)在 64 标记配置下的影响;表 B 展示了局部 2D 对齐器从 10M 扩展至 600M 时的单调提升效果(生成器仅训练一次):

表 A:核心模块设计消融(64 标记,600M 标记器,SiT 生成)

配置方案 因果建模 (C) 独立局部码本 (CB) 生成 gFID ↓ 重建 rFID ↓ 结论说明
SoftVQ 基线 – – 1.78 0.88 缺乏全局局部解耦设计
G + L1 (仅局部 1D) – ✓ 1.70 0.76 MLP 局部投影带来基础增益
G + L1 + L2 (无因果) ✗ ✓ 1.65 0.72 缺少因果约束削弱结构建模
G + L1 + L2 (共享码本) ✓ ✗ 1.68 0.73 码本混杂导致特征互相污染
完整模型 (G + L1 + L2) ✓ ✓ 1.56 0.66 全解耦因果结构达到最优

表 B:局部对齐器规模扩展(生成器 SiT 仅训练一次,参数与算力大幅节约)

对齐器规模 重建 rFID ↓ 生成 gFID ↓ 重建 PSNR ↑ 重建 SSIM ↑ 算力与训练耗时节省
10M 0.89 1.75 22.21 0.7103 基础对齐能力
100M 0.83 1.74 22.34 0.7141 细节逐步增强
200M 0.81 1.68 22.35 0.7172 持续稳步改善
400M 0.71 1.65 22.56 0.7211 逼近最优平衡
600M 0.66 1.56 22.86 0.7288 总算力减少 40%,耗时节省 54%

关键发现

  • 生成与重建双破纪录:在仅使用 64 个一维标记的极限压缩条件下,SelfBootTok 取得了 1.56 的 gFID(含 CFG)与 2.22 的无引导 gFID,显著超越同等标记预算的 SoftVQ(1.78 / 5.98)和 TiTok-B(2.48),甚至匹敌依赖 256 个标记的顶级 2D/自回归模型(MAR-H 的 1.55、LightningDiT 的 1.35)。
  • 生成器免重训的解耦红利:传统方法在扩展标记器容量时必须重训生成网络,而在扩展局部 2D 对齐器(10M 至 600M)时,下游 SiT 始终基于冻结的全局标记运行,累计训练参数从 4.9B 骤降至 2.95B,训练算力从 60K GFLOPs 减少至 36.6K GFLOPs(减少约 40%),整体训练耗时由 28 天压缩至 13 天(缩减约 54%)。
  • 局部自启动对齐器的必要性:消融实验证明,单纯依赖全局标记无法有效还原细粒度纹理;引入最优传输约束的局部 2D 分支带来了高达 0.16 的 rFID 纯增益,而以 I-JEPA 作为 1D 局部监督比 SigLIP 表现更佳(rFID 0.66 vs 0.72),体现出结构化表征对局部纹理合成的关键支撑。

亮点与洞察

  • 将生成难题转移进标记器:传统的做法是让生成器负责全部像素细节的合成,而 SelfBootTok 通过自监督机制让标记器学习从全局语义到局部细节的“脑补”先验,将生成模型的上下文长度压缩至极简的 64,实现了系统级的高效分工。
  • 最优传输连接 2D 与 1D 流形:利用 Sinkhorn 求解带有熵正则化的最优传输代价,优雅地解决了二维网格先验与一维序列离散表征之间的维度失配问题,避免了机械硬压或破坏空间连续性。
  • 可拆卸的解耦扩展范式:由于生成器仅与全局标记接口耦合,工业落地中可以根据端侧或云端的算力预算,自由插拔不同规模(10M 至 600M)的局部对齐解码器,而不需要对昂贵的文本/类别条件生成模型做二次微调。

局限与展望

  • 标记融合方式对生成存在权衡制约:实验发现,在解码端使用拼接(Concatenation)融合虽能带来最佳的重建保真度,但因增加了输入序列长度和解码复杂度,反而对无条件/条件生成的泛化稳定性带来轻微扰动,如何在更大标记规模下设计更自适应的全局-局部注意力交互机制仍待深入研究。
  • 多任务泛化尚未验证:当前评测主要集中于 ImageNet 类别条件图像生成与自重建,未进一步扩展至开放域大模型文生图(Text-to-Image)以及视频时空标记化场景。未来可探索将自启动分解范式推广至自回归文本-图像跨模态底座。

相关工作与启发

  • vs SoftVQ / TiTok: TiTok 与 SoftVQ 开启了一维超短标记序列的先河,但它们将不同尺度的视觉特征混合在同一串向量中,在极短长度下极易丢失微观纹理。本文提出全局-局部显式解耦与自启动推导,以相同 64 标记大幅超越两者的重建与生成指标。
  • vs HieraTok / VAR: HieraTok 与 VAR 探索了多尺度或逐级尺度的递进预测,但生成器本身仍需分层建模所有尺度标记。本文通过自监督对齐器完全吸收了局部细节的生成负担,生成器无需感知多尺度层级,极大降低了下游采样架构的复杂度。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次在一维图像标记化中引入自启动全局-局部解耦与最优传输对齐范式,理念十分先进。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 64/32 标记消融、从 10M 到 600M 参数的完整扩展验证、训练计算量精细分析与 Grad-CAM 可视化。
  • 写作质量: ⭐⭐⭐⭐⭐ 结构清晰,理论推导严谨(附录提供完整 ELBO 与 Sinkhorn 存在唯一性证明),图表翔实。
  • 价值: ⭐⭐⭐⭐⭐ 为高压缩率 1D 图像标记器与高效多模态生成架构提供了极具实用价值的解耦新范式。