ResTok:学习自回归图像生成中一维视觉 tokenizers 的层次残差 ResTok: Learning Hierarchical Residuals in 1D Visual Tokenizers for Autoregressive Image Generation
在 1D 视觉 tokenizer 中引入层次残差设计,显著提升 AR 图像生成效率和质量
前置知识
自回归模型(Autoregressive Model, AR)
自回归模型通过预测序列中的下一个元素来生成数据,对图像而言就是按顺序预测下一个视觉 token。它依赖条件概率分布 P(x_t | x_{<t}),每次预测都基于之前生成的所有元素。这种方式天然支持可控生成和可扩展性,但采样步数与序列长度成正比,长序列导致生成速度慢。代表性工作如 GPT 系列在文本领域、LlamaGen 在图像领域。
本文的自回归图像生成范式直接继承自语言模型,理解 AR 建模机制是理解为什么 tokenization 设计对生成质量至关重要的基础。
向量量化(Vector Quantization, VQ)
VQ 将连续特征映射到离散的 codebook 索引,实现数据压缩和离散化表示。给定输入特征 z,通过计算与 codebook 中所有嵌入的欧氏距离或余弦相似度,选择最近的码向量作为量化结果。VQ-VAE 将 VQ 与变分自编码器结合,学习紧凑的潜在表示,使高维图像可以用少量离散 token 表示。
本文的 tokenizer 基于 VQ-VAE 框架,codebook 设计直接影响 token 分布的熵值和 AR 模型的学习难度。
Vision Transformer (ViT)
ViT 将图像分割为固定大小的 patch,展平为 token 序列后输入标准 Transformer 处理。每个 patch 通过线性投影获得 token 向量,加上可学习的位置编码后,经多层自注意力层提取特征。ViT 采用各向同性(isotropic)设计,各层结构相同,与 CNN 的金字塔层次结构不同。这使得 ViT-based tokenizer 在多尺度特征融合方面需要额外设计。
本文的 tokenizer 基于 ViT 架构,核心贡献之一是在 ViT 的各向同性设计中恢复层次化表示,理解 ViT 结构限制对读懂论文关键创新至关重要。
Fréchet Inception Distance (FID)
FID 衡量生成图像与真实图像在特征空间分布的相似性。首先用预训练的 Inception 网络提取两组图像的特征,计算均值和协方差矩阵,然后计算这两个高斯分布之间的 Fréchet 距离。FID 越低表示生成质量越高。gFID 指生成 FID(evaluation on generated samples),rFID 指重建 FID(evaluation on reconstructed images)。
本文使用 FID 作为主要评估指标,理解 FID 的物理意义有助于判断性能提升的幅度和实际价值。
研究动机
现有的 1D 视觉 tokenizer 存在两个核心问题。第一是缺乏跨层级特征融合。如图 1(a) 所示,主流方法如 SEED、TiTok、SpectralAR、DetailFlow 等仅沿网络深度提取特征,无法在特定层级融合多尺度信息。这与特征金字塔网络等经典视觉表示学习的发现相矛盾——跨层级融合对于强视觉表示至关重要。第二是 codebook 熵值过高。由于层级之间的信息重叠很少被处理,当前方法在 codebook 中产生相似的嵌入,导致概率分布相对均匀。例如,基线模型的 codebook 熵达到 11.89 bits,这意味着 latent 分布非常分散,不利于 AR 建模。在 ImageNet-256×256 任务上,这些方法要么需要长序列(如 LlamaGen-L 需要 576 tokens),要么依赖手工约束(如 SpectralAR 将 token 长度绑定到频带,DetailFlow 绑定到空间分辨率),限制了灵活性和语义对齐。
本文的目标是本文旨在恢复视觉 tokenization 中固有的层次和残差先验,构建更适合 AR 图像生成的 1D tokenizer。具体目标包括:在 ViT-based tokenizer 中引入渐进式合并的层次化表示,使 latent tokens 能够查询多层级特征;学习语义结构化的残差以减少层级间的信息重叠,得到更集中的 latent 分布和更低熵的 codebook;提出层次化 AR 生成器加速采样,从逐 token 预测改为逐层级预测。最终在 ImageNet-256×256 上用更少的 tokens(128)和采样步数(9)实现竞争的 gFID(2.34),证明层次残差设计的有效性。
与已有工作不同的是,本文的独特切入点是认识到视觉和语言的本质差异。现有 1D tokenizers 直接沿用语言模型的各向同性 transformer 设计,将视觉数据视为扁平的序列 token 流,丢弃了视觉领域长期证明有效的层次和残差网络设计。如图 1 所示,ResTok 不仅沿深度查询,还沿层次查询。与手工约束(SpectralAR 的频带、DetailFlow 的分辨率)或简单堆叠不同,ResTok 通过渐进式图像 token 合并和语义残差学习,隐式地实现跨层级绑定——粗糙的 latent tokens 自然对齐高层图像 tokens,精细的 latent 捕获低层残差细节。这种设计与 ViT 的 in-context learning 范式兼容,避免了 Q-Former 或扩散解码器的优化复杂度,是首个在纯 AR 框架下系统恢复视觉层次残差先验的工作。
核心方法
ResTok 的整体思路是将 CNN/ViT 混合架构的图像编码器改造为层次化 tokenizer。直观上,就像 CNN 的特征金字塔从低分辨率高语义到高分辨率低细节,ResTok 让 1D latent token 也遵循这个语义层次。技术路线上,ResTok 在 ViT encoder 的特定层插入残差合并块,将图像 token 从细粒度渐进合并到粗粒度,形成多尺度图像 token 表示。同时,latent token 初始化时也采用残差策略——先对图像 token 池得到 coarsest latent,然后上采样回原尺寸减去原始特征得到残差,再池化得到下一层 latent。这确保每个 latent 层级捕获的是相对于更粗层的补偿信息。在训练时,编码器注意力 mask 阻止粗层访问细层,保持层级因果性。解码器简化处理,直接用量化后的 latent tokens 查询 masked 图像 tokens。
核心创新点有两个。第一是渐进式合并的层次化图像 token。与现有方法在所有层级共享单一图像 token 不同,ResTok 每 N/S 层(N 是 transformer 深度,S 是阶段数)插入合并块,用平均池化将细尺度 token 合并到粗尺度,然后通过上采样减法得到残差。这允许每层的 latent tokens 融合跨层级特征,特别是 coarsest scale 充当全局语义源,被 latent tokens 频繁查询。第二是语义残差的 latent 初始化。与 ImageFolder 的硬加性残差或 VAR 在量化瓶颈添加残差不同,ResTok 在输入阶段就学习语义结构化的残差。通过上采样-减法-池化的迭代过程,每个 latent 层级捕获的是增量语义而非重复信息。这直接降低了 codebook 熵——从基线的 11.89 bits 降到 8.79 bits(完整模型),使 latent 分布更集中,AR 模型更容易建模。
方法步骤详情
编码过程分四个步骤。首先,输入图像经 CNN encoder 下采样得到初始图像 token。其次,初始化 1D latent tokens:对初始图像 token 用 1×1 池化得到第一层 latent,然后上采样回原尺寸减去原始特征得到残差,再交替加倍宽高进行池化得到后续层 latent,最终拼接成完整的 latent token 序列。第三,将展平的图像 token 和 latent token 输入 ViT encoder。在每 N/S 层,用残差合并块替换标准 transformer block:先执行自注意力,然后将细尺度图像 token 平均池化合并到粗尺度,计算残差后通过 MLP 输出下一层特征。编码器注意力 mask 限制粗层只能访问自己和更粗层,确保层级因果性。第四,经 N 层后得到最终图像 token 和 latent token,后者经向量量化得到量化码。解码过程相对简单:masked 图像 token 和量化 latent 输入 ViT decoder,输出恢复的图像 token,经 CNN decoder 解码为重建图像。训练目标包括 MSE 损失、感知损失、GAN 损失和视觉基础模型对齐损失,总损失为各项的加权和。
技术新颖性
技术新颖性体现在三个方面。首先是在 ViT 的各向同性设计中成功引入层次化表示。与使用 Q-Former 的 GigaTok 或直接查询预训练模型的 VFMTok 不同,ResTok 在 tokenization 过程中让图像 token 进化,通过渐进式合并实现多尺度特征提取,保持了 in-context learning 的优势。其次是语义残差的双层级设计。与在量化瓶颈简单相加的残差不同,ResTok 对图像和 latent tokens 都学习语义结构化的残差——图像 token 的残差是层级的(合并-减法),latent token 的残差是初始化的(上采样-减法-池化)。这种设计确保每个层级捕获增量语义而非重复信息。最后是隐式对齐机制。ResTok 不手工将 latent token 绑定到特定图像层次(如 SpectralAR 的频带或 DetailFlow 的分辨率),而是通过嵌套 dropout 和残差学习让模型自动学习跨层级绑定。实验表明粗 latent tokens 自然对齐高层图像 tokens 捕获物体/位置/颜色,细 latent tokens 聚焦细粒度细节处理纹理/边界,无需显式约束。
实验结果
ResTok 在 ImageNet-256×256 上的实验验证了层次残差设计的有效性。首先,从重建和生成的权衡看,训练到 750k 步时达到最佳:rFID 为 1.28,gFID 为 2.34(仅用 9 步采样),IS 为 257.8,Precision 为 0.79,Recall 为 0.60。这意味着仅用 128 个 latent tokens(远少于 LlamaGen-L 的 576)就实现了竞争的生成质量。其次,消融研究证实了三个关键发现:(1) 层次显著提升表示能力。添加 4 层图像 token 层次将 rFID 从 1.87 降到 1.67,证明跨层级特征融合有效。(2) 残差有效降低 codebook 熵。同时应用图像和 latent token 残差将 HC 从 11.89 降到 8.79 bits,gFID 从 6.01 降到 4.56(无 CFG),显示集中分布对 AR 建模的重要性。(3) 层次 AR 生成器大幅加速采样。从 vanilla AR 的 128 步降到 HAR 的 9 步,gFID 仅从 4.56 增加到 5.53(无 CFG),加速比超过 14×。可视化分析揭示了语义堆叠行为:前 16 个 latent tokens 主要编码抽象语义(物体/位置/颜色),对应粗尺度 S1 和 S2;后续 token 逐步细化细节,依赖细尺度 S3 和 S4。编码器注意力图显示粗尺度 S1 充当全局语义源,被 latent tokens 频繁查询,验证了隐式对齐机制。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet-256×256 图像生成 | gFID(采样步数) | 2.34(9步) | TiTok-L-32: 2.77(8步),ImageFolder: 2.60(10步),VFMTok: 2.75(256步) | 相比 TiTok-L-32 提升 15.5%,相比 ImageFolder 提升 10.0%,同时采样步数减少或相当 |
| ImageNet-256×256 图像重建 | rFID | 1.28(128 tokens) | TiTok-L-32: 2.21(32 tokens),GigaTok-B-L: 0.81(256 tokens),DetailFlow-16: 1.22(128 tokens) | 相比 TiTok-L-32 改善 42.1%,相比 DetailFlow 改善 5.0%,使用的 token 数适中 |
| Codebook 熵值(分布集中度) | HC(bits) | 8.79(完整模型) | 基线(无层次残差):11.89,仅层次图像:11.47,仅残差 latent:10.58 | 相比基线降低 26.1%,证明语义残差有效集中 latent 分布 |
局限与改进
作者承认的局限性包括:首先,层次设计增加了 tokenizer 的复杂度,残差合并块和注意力 mask 的实现比标准 ViT 更复杂。其次,HAR 生成器虽然大幅加速,但在小采样步数下(8-9 步)gFID 仍略高于 vanilla AR(2.34 vs 2.18),说明加速与质量存在权衡。第三,当前实验集中在 ImageNet-256×256,尚未扩展到更高分辨率(如 512×512)或视频生成,层次设计的可扩展性有待验证。第四,ResTok 不手工绑定 latent token 到特定语义层次,虽然实现了隐式对齐,但在某些需要精确控制的场景下可能不如显式约束(如 SpectralAR 的频带绑定)。此外,从 Table 1 可以看到,ResTok 的 rFID(1.28)仍高于某些使用更多 tokens 的方法(如 GigaTok 的 0.81),这意味着在极致重建质量上仍有提升空间。训练稳定性方面,tokenizer 在 750k 步后 gFID 开始上升,提示存在优化困境——过度训练会提升重建但损害生成,需要仔细选择检查点。
独立分析的弱点
首先,层次合并使用简单的平均池化,可能丢失重要的高频细节。可以探索学习型合并操作(如带注意力机制的动态池化)或保留更多信息的内容感知合并方法。其次,残差初始化的池化策略(交替加倍宽高)是手工设计的,可能不是最优的。可以学习池化模式或自适应地确定每个层级的 token 数量。第三,HAR 生成器在极小步数(如 4-5 步)下的性能可能进一步下降,可以探索更精细的层次预测策略(如并行预测多个中间尺度)或结合扩散解码器来补充细节。第四,当前的对齐机制依赖 DINOv3 的预训练特征,这些特征可能不完全匹配生成任务的目标。可以探索端到端学习的对齐目标或使用生成导向的基础模型(如专为生成训练的 VF 模型)。第五,层次设计的可扩展性在高分辨率下可能受限——256×256 的 4 层合并在 512×512 可能需要更多层级,导致计算开销增加。可以探索稀疏层次或动态层级选择来平衡效率和质量。
未来方向
作者提出的未来方向包括进一步增强保真度和扩展到统一的理解-生成模型。基于本文成果,可以延伸的方向有:(1) 高分辨率生成:将 ResTok 扩展到 512×512 或更高分辨率,探索层次数量与分辨率的缩放规律,可能需要自适应层级选择或细粒度控制合并比例。(2) 视频生成:将层次残差设计推广到视频 tokenization,视频的时空层次结构更丰富,可以探索时间维度的渐进合并和残差学习。(3) 统一多模态模型:ResTok 的 1D latent tokens 天然适合与语言 token 对齐,可以构建统一的图像-文本理解-生成模型,利用层次结构实现更细粒度的跨模态对齐。(4) 条件生成:探索在 ResTok 中注入条件信息(如类别标签、文本描述、空间布局)的机制,层次化表示可能支持更精细的条件控制。(5) 动态层级选择:根据内容复杂度自适应地决定使用哪些 latent 层级,简单图像用少 tokens,复杂图像用多 tokens,实现计算资源的动态分配。(6) 可解释性分析:进一步研究隐式对齐机制,分析每个 latent 层级捕获的具体语义属性,可能揭示更可解释的视觉表示结构。
复现评估
论文承诺代码将在 https://github.com/Kwai-Kolors/ResTok 开源,这为复现提供了良好基础。实现基于 TiTok-L 的架构,每个 ViT encoder/decoder 包含 24 层、1024 维度、16 个注意力头,CNN encoder/decoder 使用 MaskGIT 的配置。训练在 ImageNet-1K(1,281,167 张图像,1000 类)上进行 200 个 epoch,batch size 为 256,使用 AdamW 优化器,初始学习率为 1e-4 经余弦衰减到 1e-5。Codebook 包含 8,192 个条目,维度为 8。嵌套 dropout 的最小保留 token 数为 4,保持概率从 80%(128 tokens)指数衰减到 0.65%(4 tokens)。算力需求方面,论文未明确报告 GPU 类型或训练时长,但基于模型规模(662M 参数)和训练配置,估计需要 8-16 张 A100/H100 GPU 训练数周。消融研究使用较短训练(30-50 epoch)验证各组件有效性,主结果使用 750k tokenizer checkpoint + 300 epoch HAR generator。复现难度中等,主要挑战在于正确实现残差合并块、层级注意力 mask 和 M-RoPE 的修改版本。依赖项包括 TiTok(Apache-2.0)、LlamaGen(MIT)、DINOv3(自定义许可证)等开源项目,许可证允许研究使用但有限制。
论文图表