Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization

https://arxiv.org/abs/2503.11056

Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization

摘要

本文提出FlowMo（Flow to the Mode），一种基于Transformer的扩散自编码器，在无需卷积、对抗损失、空间对齐二维潜在码或从其他分词器蒸馏的情况下，实现了多压缩率下的图像分词任务新SOTA性能。核心创新点包括：

1. 模式匹配预训练：通过扩散损失端到端训练系统
2. 模式搜索后训练：选择性丢弃与原始图像感知距离较远的重建模式
3. 一维潜在码Transformer架构：突破传统CNN架构限制

引言

生成模型发展现状

近年来，扩散模型（Diffusion Models）和离散自回归模型（Discrete Autoregressive Models）在图像/视频生成领域取得显著进展。现有视觉生成系统普遍采用两阶段范式：

1. 分词器（Tokenizer）：将像素数据压缩到低维潜在空间
2. 生成模型：在压缩后的潜在空间进行训练

图像分词技术演进

• 传统方案：以VQGAN为代表，采用CNN架构和二维空间对齐潜在码
• 改进尝试：TiTok引入Transformer架构但依赖CNN预训练，ViT-VQGAN使用Transformer编解码器
• 现存问题：对抗损失不稳定、硬件效率受限、依赖预训练分词器蒸馏

FlowMo核心思想

通过两个关键洞察实现突破：

1. 模式选择优先性：感知相似的模式比所有模式匹配更重要
2. 两阶段训练策略：
   • 预训练阶段：端到端扩散损失训练
   • 后训练阶段：选择性模式过滤

相关工作

图像分词技术

• 主流架构：VQGAN及其变体主导离散图像分词领域
• 近期探索：TiTok（Transformer+一维码）、ViT-VQGAN（Transformer编解码）

扩散自编码器

• 技术起源：首次提出于[46]，后续研究聚焦表示学习和感知压缩
• 关键技术：直校正流（Rectified Flow）、有限标量量化（FSQ）

并行工作

• DiTo：连续图像分词扩散自编码器
• FlexTok：基于传统VAE的扩散分词器

方法

系统架构

1. 编码器（Encoder）：基于MMDiT的Transformer架构，将图像映射到潜在码
2. 解码器（Decoder）：条件扩散模型，学习从噪声到数据的速度场v
3. 量化模块：采用无查找表量化（LFQ）生成二值潜在码

训练流程

阶段1A：模式匹配预训练

• 目标：最大化潜在码的信息量，匹配真实分布
• 损失函数：
  • 扩散匹配损失（Rectified Flow Loss）
  • 感知距离监督（Lperc）
  • 潜在码熵损失（Lent）和承诺损失（Lcommit）

阶段1B：模式搜索后训练

• 目标：过滤非感知相似重建模式
• 关键技术：
  • 通过采样ODE积分计算n步样本损失（Lsample）
  • 使用ResNet作为感知网络
  • 调整采样步长分布（Shift参数ρ）

采样过程

• 概率流ODE求解：通过固定步长积分生成图像
• 加速策略：采用非线性步长分布（ρ=4）

生成模型训练

• 使用MaskGiT在第二阶段训练生成模型
• 关键设置：隐藏层尺寸1024，层数28，总参数量397M

分析

模式分布可视化

局限性

• 推理速度：需多次前向计算（本文采用n=25步）
• 计算资源：大规模训练需8xH100节点支持

结论

FlowMo通过创新的扩散自编码器架构和两阶段训练策略，在ImageNet-1K图像分词任务上取得SOTA性能。其核心贡献在于：

1. 首次实现无CNN/对抗损失的分词器
2. 提出模式导向的训练范式
3. 验证Transformer在低维潜在空间建模的潜力

文章笔记

1 研究目的

• 解决图像分词难题
  突破传统分词器依赖CNN架构、对抗损失不稳定、需预训练蒸馏等限制
• 提升感知重建质量
  在ImageNet-1K等高难度基准测试中实现SOTA性能指标（rFID/PSNR/LPIPS）

2 研究方法

• 双阶段训练框架
  1. 模式匹配预训练：端到端扩散损失训练编码器-解码器
  2. 模式搜索后训练：选择性过滤非感知相关重建模式
• 技术创新点
  • 一维潜在码Transformer架构（突破空间对齐限制）