Improving Flexible Image Tokenizers for Autoregressive Image Generation¶

会议: NeurIPS 2026
arXiv: 2601.01535
OpenReview: https://openreview.net/forum?id=neJJASNqpr
代码: https://github.com/zfu006/ReTok
领域: 图像生成
关键词: 图像生成, 自回归模型, 视觉分词器, 嵌套丢弃, 语义正则化

一句话总结¶

针对传统嵌套丢弃导致信息向头部过度聚集并在下游自回归模型中产生生成瓶颈的问题，本文提出 ReTok 框架，通过冗余 Token 填充与分层语义正则化激活尾部 Token 并实现自粗至细的高保真生成。

研究背景与动机¶

自回归（AR）模型凭借强大的可扩展性与统一多模态建模的潜力，已成为离散图像生成领域的核心范式。这类生成器通常依赖预训练的视觉分词器（Visual Tokenizer）将像素空间压缩至紧凑的离散潜在空间。然而，绝大多数现存分词器将图像编码为固定网格大小（如 1D 或 2D 固定长度）的表征，难以动态匹配具有异质复杂度的自然图像，并对生成速度与计算开销构成刚性约束。

为了打破固定网格的限制，近期工作尝试构建可变长度的灵活分词器（Flexible Tokenizers），主要通过在训练阶段引入嵌套丢弃（Nested Dropout）策略——随机截断尾部 Token，并仅依赖剩余的前缀子序列重建输入图像。然而，这种朴素的截断训练策略存在严重的结构性矛盾：分词器为了最小化任意截断长度下的重建误差，被迫将绝大部分核心视觉信息压缩进序列开头的若干 Token 中，导致尾部 Token 的利用率极低甚至退化为冗余。当下游自回归模型基于此类分词器建模时，生成性能遭遇明显的瓶颈——增加序列长度无法持续改善图像质量，甚至因后部 Token 承载信息微弱而在长序列下出现生成 FID（gFID）倒挂恶化，显著落后于相同 Token 预算下的固定长度分词器。

本文的切入视角在于重新审视嵌套丢弃的激活机制与不同位置 Token 的语义分配需求。既然朴素截断让原有尾部 Token 处于“被丢弃边缘”，就应当从训练机制上将整个目标序列整体向“序列前半段”偏移，同时从粗粒度语义到细粒度纹理显式引导各阶段 Token 的职责分工。核心 idea：通过引入冗余 Token 填充拉长训练序列以高频激活原有尾部 Token，并辅以强度自头至尾线性衰减的分层语义正则化与解码器微调，构建全局语义对齐且细粒度纹理充分利用的可变长度 1D 分词器 ReTok。

方法详解¶

整体框架¶

ReTok 的整体架构继承了混合结构 1D 离散分词器（基于 GigaTok）的设计规范，核心流程包含特征抽取、可变长度离散潜在建模与多尺度特征映射恢复。

输入分辨率为 \(256 \times 256 \times 3\) 的原始图像 \(\mathbf{X}\) 首先经过步长下采样为 16 的轻量 CNN 编码器，提取出紧凑的 2D 局部图像块嵌入 \(\mathbf{P} \in \mathbb{R}^{\frac{H}{16} \times \frac{W}{16} \times D}\)（长度 \(L=256\)）。随后，ViT 架构的 Q-Former 编码器交替执行自注意力与以图像嵌入 \(\mathbf{P}\) 为键/值的交叉注意力，将图像特征抽象聚合至一组可学习的 1D 查询向量（Query Tokens）中。量化模块对查询输出进行向量量化（VQ），生成紧凑离散索引序列。在解码重建阶段，截断的前缀离散 Token 经由 MLP 投影视觉维度后与可学习的 2D 掩码 Token 拼接，依次送入 ViT 解码器与轻量级 CNN 解码器还原高保真图像。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 X (256x256x3)"] --> B["CNN 编码器 (16x 下采样)<br/>生成图像块嵌入 P"]
    B --> C["冗余 Token 填充与联合编码<br/>拼接基础查询 Q 与冗余查询 Q_pad"]
    C --> D["矢量量化与嵌套丢弃<br/>截断保留 k 个离散 Token"]
    D --> E["分层语义正则化约束<br/>对齐 DINOv2 高层特征 (权重衰减)"]
    E --> F["ViT 与 CNN 组合解码器<br/>结合掩码 Token 恢复原图"]
    F --> G["解码器微调<br/>固定编码器与量化表，放宽语义约束"]
    G --> H["高保真重构与自回归生成输出"]

关键设计¶

1. 冗余 Token 填充：消除信息向头部过度聚集并充分激活尾部 Token

朴素嵌套丢弃在长度为 \(N\) 的序列上随机截断保留 \(k \le N\) 个 Token，处于末端位置的 Token 在多数采样步中均被遮蔽抛弃，分词器自然倾向于将所有图像关键内容塞进前几个 Token。为了迫使模型在整个目标长度 \(N\) 内都保持高激活度，ReTok 在基础查询序列 \(\mathbf{Q} \in \mathbb{R}^{N \times D}\) 尾部拼接 \(M\) 个额外的冗余查询 Token \(\mathbf{Q}_{pad} \in \mathbb{R}^{M \times D}\)，联合编码后得到扩展离散表征：

\[[\mathbf{Z}; \mathbf{Z}_{pad}] = \mathcal{Q}\left(\mathcal{E}\left([\mathbf{P}; \mathbf{Q}; \mathbf{Q}_{pad}]\right)\right) = [\mathbf{z}_1, \dots, \mathbf{z}_N, \mathbf{z}_{N+1}, \dots, \mathbf{z}_{N+M}]\]

在训练期间，嵌套丢弃在总长度为 \(N+M\) 的全序列上执行截断，采样保留长度满足 \(k \le N+M\)。这一转换让原本处于末端的 Token（如第 256 位）在扩展视界下退居为“前半段 Token”，被抽中保留并参与图像重建的概率大幅提升。由于模型必须准备好仅依赖前缀子集进行重建，原来的全部 \(N\) 个 Token 都被赋予了不可或缺的信息承载压力。在下游自回归模型训练与实际推理部署时，模型直接舍弃附加的 \(\mathbf{Z}_{pad}\)，仅保留激活充分、信息分布均衡的前 \(N\) 个有效 Token。

2. 分层语义正则化：兼顾前缀全局概念对齐与末端细粒度纹理重建

高压缩比的 1D 分词器天然擅长捕获高层语义表征，但若对全序列施加单一强度的视觉大模型对齐约束，会抑制尾部 Token 学习低级像素细节的能力。为此，ReTok 引入依赖于当前截断长度 \(k\) 的动态语义正则化损失：

\[\mathcal{L}_{total} = \mathcal{L}_{rec} + \lambda(k) \mathcal{L}_{reg}, \quad \text{其中} \quad \mathcal{L}_{reg} = -\cos\left(\text{MLP}(f_l^{dec}), f^{DINO}\right)\]

其中 \(f_l^{dec}\) 为 ViT 解码器第 \(l\) 层的中间特征，\(f^{DINO}\) 为预训练冻结的 DINOv2-B 提取的全局语义特征。关键机制在于使平衡系数 \(\lambda(k)\) 成为保留长度 \(k\) 的单调递减函数：在极短序列（如 \(k=32\)）时赋予最高的对齐权重（ReTok-S-S 为 2.0，ReTok-S-B 为 2.5），迫使前缀 Token 专精于描摹主体轮廓与核心类别语义；随着截断长度 \(k\) 向 256 延伸，\(\lambda(k)\) 线性递减至基准下限 0.5（并在超过 256 的填充区域保持 0.5）。这种渐进式松绑机制使后续 Token 摆脱了严苛的语义瓶颈，得以全力投入到局部高频结构与细腻材质纹理的重建补偿中。

3. 解码器微调：解除前缀过约束并强化极短序列下的图像重建表现

尽管分层语义正则化极大增强了前缀 Token 的判别性语义表征，但在训练初期过大的语义损失约束在一定程度上损害了仅依赖极少数 Token（如 32 或 64 个）时的浅层像素还原质量。为补齐这一短板，ReTok 在整体分词器完成预定轮次训练后引入独立的解码器微调阶段。具体而言，该阶段完全冻结已具备优质信息分布的 CNN 编码器、Q-Former 编码器以及 VQ 量化码本，仅放开解码器参数更新，并将全序列长度范围内的正则化系数 \(\lambda(k)\) 统一固定为较小的常数 0.5。微调专注于让解码器在离散码本表征固定的前提下，重新适应从高度紧凑的前缀语义中平滑解码像素，显著拔高了短序列阶段的图像重建逼真度。

损失函数 / 训练策略¶

分词器第一阶段联合训练的总目标由重构损失 \(\mathcal{L}_{rec}\) 与分层语义正则化项 \(\lambda(k)\mathcal{L}_{reg}\) 组成。其中 \(\mathcal{L}_{rec}\) 采用 VQGAN 标准方案，包括像素级 \(\ell_1\) 损失、基于 LPIPS 的感知损失、基于 PatchGAN 的判别器对抗损失以及码本承诺损失。网络使用 AdamW 优化器（\(\beta_1=0.9, \beta_2=0.95\)），学习率设为 \(1\times 10^{-4}\) 并按余弦退火衰减至 \(1\times 10^{-5}\)。ReTok-S-S（136M）与 ReTok-S-B（232M）分别训练 200 与 250 个 epoch，随后进行 50 个 epoch 的解码器微调。下游自回归生成器选用 LlamaGen 框架（覆盖 B、L、XL 规格），采用步长为 32 的渐进离散 Token 建模，推理时采用前 18% Token 关闭 CFG（CFG=1.0）、后续 Token 开启全局搜索最佳 CFG 的阶梯引导策略。

实验关键数据¶

主实验¶

在 ImageNet \(256 \times 256\) 类别条件图像生成基准上，ReTok 与主流连续/离散、固定/灵活长度分词器进行了全面对比评测（如下表所示）：

建模范式	分词器类型	分词器 (参数量)	重构 rFID \(\downarrow\)	生成器 (参数量)	Token 数量	生成 gFID \(\downarrow\)	生成 gIS \(\uparrow\)	Precision \(\uparrow\)	Recall \(\uparrow\)
离散固定	1D 固定网格	TiTok-S (72M)	1.71	MaskGIT-UViT-L (287M)	128	1.87	281.8	–	–
离散固定	1D 固定网格	GigaTok-S-S (136M)	1.01	LlamaGen-L (343M)	256	2.86	261.2	0.81	0.57
离散固定	1D 固定网格	GigaTok-S-B (232M)	0.89	LlamaGen-L (343M)	256	2.71	246.3	0.81	0.58
灵活长度	1D 灵活 (嵌套)	FlexTok (~2.5B)	1.08	LlamaGen (1.33B)	256	~2.50	–	–	–
灵活长度	1D 灵活 (嵌套)	One-D-Piece (641M)	1.08	LlamaGen-L (318M)	256	3.86	231.7	0.81	0.51
灵活长度	1D 灵活 (分层)	DetailFlow-32 (270M)	0.80	LlamaGen-L (326M)	256	2.75	250.8	0.81	0.58
灵活长度	1D 灵活 (频域)	SpectralAR (–)	4.03	SpectralAR-d20 (600M)	64	2.49	305.4	–	–
灵活长度	1D 灵活 (本文)	ReTok-S-S (136M)	1.09	LlamaGen-L (343M)	256	2.92	243.5	0.81	0.57
灵活长度	1D 灵活 (本文)	ReTok-S-B (232M)	1.01	LlamaGen-L (343M)	256	2.66	231.7	0.82	0.57
灵活长度	1D 灵活 (本文)	ReTok-S-B (232M)	1.01	LlamaGen-XL (775M)	256	2.27	245.9	0.82	0.60

消融实验¶

基于 ReTok-S-S 分词器与自回归生成器，针对冗余填充数量、分层正则化机制及嵌套步长进行了系统消融（如下表所示）：

实验切片 / 配置模块	考察维度 / 变量设置	评估 Token 长度	重构 rFID \(\downarrow\)	感知 LPIPS \(\downarrow\)	生成 gFID \(\downarrow\)	核心结论与机制说明
冗余填充 Token 数 (\(M\))	填充 64 个	256	1.21	0.238	4.74	填充过少无法彻底消除尾部信息衰减
冗余填充 Token 数 (\(M\))	填充 224 个 (默认)	256	1.18	0.234	4.68	期望采样长度居中，全序列利用率达到平衡
冗余填充 Token 数 (\(M\))	填充 384 个	256	1.34	0.245	5.03	填充过多稀释了前 256 个 Token 的关键注意力
语义正则化策略 \(\lambda(k)\)	固定低强度 (\(\lambda=0.5\))	32 / 256	9.92 / 1.15	0.411 / 0.232	16.40 / 4.92	短序列缺乏语义约束导致生成质量显著坍塌
语义正则化策略 \(\lambda(k)\)	固定高强度 (\(\lambda=2.0\))	256	1.24	0.241	4.72	尾部过度语义化剥夺了细粒度纹理重建能力
语义正则化策略 \(\lambda(k)\)	分层衰减 (\(\lambda: 2.0 \to 0.5\))	32 / 256	9.02 / 1.18	0.415 / 0.234	14.79 / 4.68	实现从粗轮廓到细纹理的平滑层级自适应分配
嵌套丢弃步长选择	步长 16	32 / 256	9.56 / 1.22	0.422 / 0.234	15.79 / 4.71	过于密集的截断采样干扰了短序列语义收敛
嵌套丢弃步长选择	步长 32 (默认)	32 / 256	9.02 / 1.18	0.415 / 0.234	14.79 / 4.68	在不同阶段表征学习与计算效率间取得最佳平衡

关键发现¶

尾部 Token 的激活是解锁长序列潜力的决定性因素：Token 贡献度摄动分析表明，朴素嵌套丢弃的分词器仅在头部前 32 个 Token 呈现极高重建贡献，后续 Token 贡献趋近于零；而引入 224 个冗余 Token 填充后，所有 256 个 Token 均呈现接近固定网格分词器的均匀贡献分布，从根源上治愈了长序列生成 FID 倒挂的顽疾。
渐进生成的边际效益与图像异质复杂度相匹配：在渐进采样过程中，简单背景主体图像仅需 32-64 个 Token 即可生成无伪影的清晰图像，而富含高频细节的复杂自然场景（如织物纹理、繁复背景）随序列延展至 256 个 Token 时表现出明显的质感跃升，验证了灵活分词在计算效率与保真度之间的动态权衡价值。

亮点与洞察¶

将目标序列推入“前半段”的冗余填充技巧：不改动最终生成模型接口与 Token 预算，仅通过训练期在尾端拼接虚设查询 Token，利用概率平移巧妙激活原有尾部表征，思路极为简洁优美且具备跨架构复用性。
与表征粒度对齐的分层正则衰减：敏锐洞察视觉大模型（DINOv2）的高层概念偏向，在短序列时强制对齐概念，长序列时松绑让位给像素级重构，化解了离散量化中高层语义与底层细节争抢表征容量的固有矛盾。
自回归生成中的阶梯引导（Step-CFG）设计：前 18% Token 关闭 CFG 确保全局语义与构图的多样性，后续 Token 开启适度 CFG 强化局部边缘与细节的准确性，这一推理技巧显著优化了长序列 AR 采样的稳定性。

局限与展望¶

图像分辨率受限于 \(256 \times 256\)：当前验证主要聚焦在固定正方形单分辨率基准上，面向多宽高比以及 \(512 \times 512\) 或更高分辨率的扩展适应性尚待进一步探索。
时序冗余维度的延伸空间：视频生成领域天然存在巨大的帧间时间冗余，探索 ReTok 的可变长度思想在 3D 时空视频分词器（Video Tokenizer）中的泛化应用将是极具价值的后续方向。

评分¶

新颖性: ⭐⭐⭐⭐ [针对嵌套丢弃早期聚集痛点提出的冗余填充与分层语义衰减构思巧妙，切中要害]
实验充分度: ⭐⭐⭐⭐⭐ [完备的对比基准、逐步累加的路线图消融、Token 敏感度分析与渐进可视化]
写作质量: ⭐⭐⭐⭐⭐ [逻辑推演严密，痛点剖析直观透彻，技术细节交代详尽完整]
价值: ⭐⭐⭐⭐⭐ [填补了灵活自回归图像分词器与固定网格分词器之间的性能鸿沟，实用性极佳]