ESC
输入关键词搜索文章
目录

WeTok

离散视觉 Tokenization · ICLR 2026
GQ 扩展 codebook，GD 补回高压缩细节
2026/05/20 20:24:06·2026/05/21 10:20:00
AI/视觉分词器·46 min read
visual-tokenizerdiscrete-tokenizationGQGD

参考来源

简介

WeTok 是一篇值得读的视觉 tokenizer 论文，原因不只是它把重建 rFID 做到了很低，而是它直接挑战了一个长期默认判断：离散 tokenizer 通常压缩率高，但重建保真度弱于 continuous VAE。论文用两个很小但抓住瓶颈的设计把这个 trade-off 往前推了一步：

  1. GQ（Group-wise Lookup-Free Quantization）：把 LFQ 的指数级 entropy loss 计算拆成分组计算，使大 codebook 训练不再被显存卡住，同时避免 BSQ 逐 bit 独立假设带来的过大近似误差。
  2. GD（Generative Decoder）：把 decoder 从确定性重建器改成带噪声先验的条件生成器，让高压缩 token 只负责结构和语义，细节由 learned prior 补齐。

这篇论文的价值在于，它给出了一个非常工程化的启发：视觉生成系统的上游瓶颈往往不是 transformer 主干，而是 tokenizer 的压缩-保真 Pareto frontier。WeTok 在 ImageNet 50K zero-shot reconstruction 上达到 rFID 0.12，超过 FLUX-VAE 0.18 和 SD-VAE 3.5 0.19；在 768× 高压缩下达到 rFID 3.49，也优于 Cosmos 在 384× 下的 4.57。这说明离散 tokenization 仍有很大空间，不应过早让位给 continuous latent。

WeTok 与 SOTA tokenizer 的 zero-shot reconstruction 对比
图 1：WeTok 在压缩率与重建保真度之间给出新的 Pareto 点。左侧是 rFID 与压缩率的量化对比，右侧展示了文字、纹理和人脸等细节的重建差异。

作者与机构

论文作者共 10 位，覆盖高校、腾讯 WeChat Vision、中国科学院深圳先进技术研究院与上海人工智能实验室等机构。作者列表中的 ♠ 表示工作完成于腾讯 WeChat Vision 实习期间，* 表示通讯作者。

作者机构标注说明
Shaobin Zhuang1, ♠一作；Shanghai Jiao Tong University；WeChat Vision 实习
Yiwei Guo3, ♠Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences；WeChat Vision 实习
Fangyikang Wang5, ♠Zhejiang University；WeChat Vision 实习
Canmiao Fu2WeChat Vision, Tencent Inc.
Zhipeng Huang2WeChat Vision, Tencent Inc.
Zeyue Tian4, ♠Hong Kong University of Science and Technology；WeChat Vision 实习
Xiaohui Li1Shanghai Jiao Tong University
Ying Zhang2WeChat Vision, Tencent Inc.
Chen Li2WeChat Vision, Tencent Inc.
Yali Wang3, 6, *通讯作者；SIAT CAS；Shanghai AI Laboratory
编号机构
1Shanghai Jiao Tong University
2WeChat Vision, Tencent Inc.
3Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences
4Hong Kong University of Science and Technology
5Zhejiang University
6Shanghai AI Laboratory

作者与机构

一作：Shaobin Zhuang（上海交通大学，WeChat Vision 实习生）

通讯作者：Yali Wang（中科院深圳先进院，上海 AI 实验室）

作者共 10 人，来自 6 个机构：Shanghai Jiao Tong University、WeChat Vision (Tencent)、SIAT CAS、HKUST、Zhejiang University、Shanghai AI Laboratory。其中多位学生作者在 WeChat Vision 实习期间完成本工作。

论文信息：32 pages, 15 figures, 39 tables；arXiv:2508.05599 v3；首次提交 2025-08-07，最近修订 2026-02-09。

相关工作

视觉 tokenizer 的技术谱系

视觉 tokenizer 把像素空间压缩为更短的 latent 表示，让生成模型在 latent space 中高效操作。这个方向分为两条路线：

  • Continuous tokenizer：用 VAE 把图像映射到连续浮点 latent，如 SD-VAE、FLUX-VAE。重建质量高，但 latent 是 float32/bfloat16，压缩率受限，且与语言模型式的离散 token prediction 不兼容。
  • Discrete tokenizer：用量化器把图像映射到有限词表中的离散 code，如 VQ-VAE、VQGAN、MAGVIT-v2、Open-MAGVIT2、BSQ、WeTok。压缩率高，token 与语言 token 同构，适合 AR/MLM 视觉生成，但重建质量通常不如 continuous tokenizer。

WeTok 的核心论点是：离散 tokenizer 的"压缩率高但保真度低"不是结构宿命，而是 codebook scaling 和 decoder 建模方式没有处理好。通过 GQ 和 GD，WeTok 在离散 tokenizer 上同时实现了高压缩和高保真。

主要前置工作对比

方法年份核心思路ImageNet rFID与 WeTok 关系
VQ-VAE2017显式 learnable codebook + nearest-neighbor lookup~8.30 (16×)祖先范式；WeTok 放弃显式 codebook
VQGAN2021VQ + perceptual loss + GAN loss4.99 (16K codebook)WeTok 继承 GAN/perceptual 训练框架
SD-VAE 3.52024Continuous VAE for latent diffusion0.19 (zero-shot)WeTok 的 continuous baseline；WeTok rFID 0.12 超过它
FLUX-VAE2025Leading continuous tokenizer0.18 (zero-shot)WeTok 的 continuous baseline
MAGVIT-v2/LFQ2024Lookup-free quantization，用隐式二值 codebook1.17 (16×)WeTok 的直接前身；GQ 解决 LFQ entropy loss 内存瓶颈
BSQ2024Binary Spherical Quantization，bit-wise independence 降复杂度~4.57 (384×)WeTok 的直接对比对象；GQ 理论误差严格小于 BSQ
Open-MAGVIT22025MAGVIT-v2 开源复现，CNN encoder/decoder1.17 (16×)WeTok 的架构基础
Cosmos2025NVIDIA image/video tokenizer suite4.57 (384×)WeTok 高压缩对比对象；WeTok 768× rFID 3.49 优于 Cosmos
WeTok2025GQ + GD0.12 (16×, zero-shot)本文
UniWeTok2026WeTok + 2^128 codebook + PPD/GAP + unified MLLM1.38 (generation)WeTok 后续扩展

核心技术脉络

2017
VQ-VAE / VQGAN（显式 codebook，lookup cost）
2024
MAGVIT-v2 / LFQ（lookup-free binary codebook，entropy loss 内存瓶颈）
2024
BSQ（bit-wise independence 降复杂度，approximation error 大）
2025
WeTok GQ（group-wise LFQ，在 LFQ 内存效率和 BSQ 近似误差之间取折中）
2025
WeTok GD（generative decoder，把 deterministic decoder 改为条件生成模型）

论文对已有方法的批评

  1. VQ-VAE/VQGAN：codebook utilization 不稳定，lookup 操作引入额外训练/推理成本
  2. LFQ：entropy loss 计算复杂度随 codebook size 线性增长，$d=24$ 起即 OOM，无法扩展到更大 codebook
  3. BSQ：bit-wise independence 假设引入 approximation error，导致性能下降
  4. 传统离散 tokenizer decoder：deterministic mapping，高压缩下容易学到"平均图像"，高频纹理丢失
前置知识
  • Vector Quantization (VQ)
    将连续向量映射到有限 codebook 中最近邻 code 的过程。VQ-VAE 用这个机制建立离散 bottleneck，使图像可以被表示为 token 序列。
  • Lookup-Free Quantization (LFQ)
    MAGVIT-v2 提出的量化方式，不使用可学习的 codebook，而是把 latent feature 的每个 channel 直接二值化为 $\{-1,1\}$。隐式 codebook 大小为 $2^d$，可以极大扩展词表，且不需要 nearest-neighbor lookup 操作。
  • Entropy Loss
    为防止 LFQ 的 codebook utilization collapse，引入 entropy regularization 鼓励 code 均匀分布。公式为 $H(X) = -\sum_{X_i \in X} X_i \log X_i$。但计算 codebook entropy 需要 $O(2^d)$ 复杂度，$d$ 增大时内存成为瓶颈。
  • Binary Spherical Quantization (BSQ)
    把 embedding 投影到低维超球面后做二值量化，假设各 bit 独立，从而把 entropy 计算降到 $O(d)$。但独立性假设牺牲了 bit 之间的统计耦合。
  • Generative Decoder
    与 deterministic decoder 不同，生成式 decoder 建模 $p(I \mid U_Q)$ 条件分布，通过额外噪声变量 $z \sim \mathcal{N}(0,I)$ 允许从条件分布中采样，而非只输出条件期望。
  • Group-wise Quantization
    将 latent feature 在 channel 维度分组，每组独立量化。这种方式在保持组内联合建模能力的同时，降低了整体 entropy 计算的内存需求。
  • rFID (recognition FID)
    使用在 ImageNet 预训练的特征提取器计算的 FID，反映重建图像与原图在特征空间中的分布距离，比普通 FID 更适合评估 tokenizer 的重建质量。
阅读前疑问
  1. GQ 的分组策略具体如何选择？论文说 $g=4, d'=8$ 通常最优，但这个 choice 是基于实验还是理论？不同压缩比是否需要不同 $g$？
  2. GD 的生成式 decoder 用 GAN loss 而非 diffusion decoder，single-step 采样确实更高效，但 GAN 训练的不稳定性是否会成为瓶颈？Stage 2 的 zero-initialization 是否真的能保证训练稳定？
  3. WeTok 的 rFID 0.12 超过 FLUX-VAE 0.18，但这种优势是否在所有类型图像上成立？对人脸、纹理密集图像、医学图像是否同样有效？
  4. 768× 压缩比下 rFID 3.49，与 Cosmos 384× 的 rFID 4.57 相比确实更好，但两者的视觉质量在人眼感知上差异有多大？
  5. WeTok 主要验证了图像重建，下游 AR 生成只在 ImageNet class-conditional 上测了 FID。对于 text-to-image 或 image editing 场景，WeTok 的离散 token 是否同样有效？
面临问题及 Insights

问题 1：Scalable Codebook

具体问题：要降低离散 tokenizer 的量化误差，需要扩大 codebook。但 LFQ 的 entropy loss 需要在 $\{-1,1\}^d$ 全空间上计算概率分布，内存随 $d$ 指数增长。LFQ 在 $d=24$ 起即 OOM，无法扩展到更大 codebook。

BSQ 用 bit-wise independence 近似降低内存，但独立性假设引入 approximation error，导致性能下降。

Insight：分组本身就是一种近似。如果把 entropy 计算从完整 $\{-1,1\}^d$ 空间分解为 $g$ 个 $\{-1,1\}^{d'}$ 空间的组合，内存从指数降到线性（$g \times 2^{d'}$ vs $2^d$）。关键是：分组造成的 approximation error 是否比 BSQ 的 independence approximation 更小？

论文给出理论答案（Proposition 3.1）：在非退化的 group choice 下，GQ 的 approximation error 小于 BSQ。这来自于 entropy 函数的凹性和 refinement ordering 的单调性——越细的分组近似误差越大，BSQ 是最细（逐 bit），LFQ 是最粗（整体），GQ 处于中间。当 $g=d, d'=1$ 退化为逐 bit 分组时，GQ 等价于 BSQ。

问题 2：Generative Modeling in High Compression

具体问题：高压缩下，单个离散 token $U_Q$ 可能对应多个 ground-truth 图像（例如不同的毛发纹理、叶片形态）。Deterministic decoder 被训练预测条件期望，输出的是"平均图像"，高频细节被平滑掉。

Insight：把 decoder 从 deterministic mapping $G(U_Q)$ 改为条件生成模型 $G(U_Q, z)$，其中 $z \sim \mathcal{N}(0,I)$。这样 decoder 不再输出唯一答案，而是从 $p(I \mid U_Q)$ 条件分布中采样，$U_Q$ 作为条件约束全局内容，$z$ 帮助 decoder 在条件分布中采样具体高频细节。

这个思路与 consistency decoder、$\epsilon$-VAE 一脉相承，但 WeTok 是首个将其引入离散 tokenizer 的工作，且用 GAN-style single-step 采样替代 diffusion multi-step，在效率和质量间取得更好平衡。

类比理解

可以把 tokenizer 想象成古代驿站传递军情的符节系统：

  • 连续 tokenizer 像把整张地图打包传递，细节完整但体积庞大
  • 传统离散 tokenizer 像用少量符节编码地图，压缩率高但重建时只能给出"平均位置"，细节丢失
  • WeTok 的 GQ 像把地图分成几个区域，每个区域用独立符节编码，兼顾压缩和局部精度
  • WeTok 的 GD 像在解码时引入"现场判断力"，符节给出大方向，解码者根据局部知识填充细节，但这个"现场判断力"是学习出来的 prior 而非随意发挥
模型和方法

整体架构

WeTok 的 pipeline 分为三个阶段：

  1. Encoder：$U = E(I)$，将图像 $I \in \mathbb{R}^{H \times W \times 3}$ 压缩为 $U \in \mathbb{R}^{h \times w \times d}$
  2. Quantizer：Group-wise Lookup-Free Quantization，将 $U$ 量化为离散 token $Q$
  3. Decoder：Generative Decoder，从 $U_Q$（或 $U_Q + z$）重建图像 $\hat{I}$

架构采用 Open-MAGVIT2 的 CNN architecture，encoder 和 decoder 参数量分别为 198M 和 261M（base channel 256, residual blocks 4）。

WeTok 的 Group-wise Lookup-Free Quantization 与 Generative Decoder 架构
图 2：WeTok 的整体结构。GQ 在 channel 维度分组做 lookup-free quantization，GD 在 decoder 端引入噪声变量，把高压缩重建从确定性映射改成条件生成。

GQ：Group-wise Lookup-Free Quantization

Channel 分组

将 $U \in \mathbb{R}^{h \times w \times d}$ reshape 为 $U_G \in \mathbb{R}^{h \times w \times g \times d'}$，满足 $d = g \cdot d'$。第 $k$ 组使用独立的隐式二值 codebook $C_{\mathrm{GQ},k} = \{-1,1\}^{d'}$，在组内执行 lookup-free quantization（直接取 sign）。

Token Entropy Loss 的分解

利用 entropy 的 additivity，token entropy loss 从 $\{-1,1\}^d$ 空间转为 $g$ 个 $\{-1,1\}^{d'}$ 空间的线性组合：

$$L_{\text{Token Entropy}} = \frac{1}{hw}\sum_{i,j} H(q(c|U[i,j])) = \frac{1}{hw}\sum_{i,j}\sum_{k=1}^{g} H(q_G(c_k|U_G[i,j,k]))$$

内存复杂度从 $O(2^d)$ 降到 $O(g \cdot 2^{d'})$。当 $d'=8$ 时，即使 $d=40$（$g=5$），内存仍约 10.6GB，与 $d=8$ 时相当，而 LFQ 在 $d=24$ 即 OOM。

Codebook Entropy Loss 的近似

由于 codebook entropy loss 中的 $H(\sum \cdot)$ 操作无法直接因 entropy additivity 分解，论文对 grouped conditional probabilities 提出乘积分解近似：

$$\frac{1}{hw}\sum_{i,j} q(\mathbf{c}|U[i,j]) \approx \prod_{k=1}^{g} \frac{1}{hw}\sum_{i,j} q_G(c_k|U_G[i,j,k])$$

从而 codebook entropy loss 可转为各 group entropy 的加和形式（具体公式见原文 Eq. 7/8）。该近似误差由 Proposition 3.1 与 Appendix A 分析，结论为在非退化的 group choice 下 GQ 的 approximation error 小于 BSQ。

GQ 的可调折中性

  • $g=1, d'=d$：等价于 LFQ，最小 approximation error，最大内存开销
  • $g=d, d'=1$：等价于 BSQ，最小内存，但 approximation error 最大
  • $g \in (1, d)$：中间态，可调 trade-off

主文消融显示增大 $g$ 能显著提升重建性能且不触发 LFQ 的内存瓶颈；大规模 SOTA 配置中 $d'=8$，$g$ 随压缩率取 2/4/8。

Loss 结构：从 VQ-VAE 到 GQ/GD

WeTok 的 loss 设计可以理解为：先继承 VQ-VAE/VQGAN 的重建、感知与对抗训练框架，再用 LFQ/GQ 的 entropy regularization 替代 learnable codebook 的部分约束，最后在 GD 阶段把 decoder 的对抗训练改成条件生成形式。

VQ-VAE 五部分 loss（Eq. 2）

标准目标由 reconstruction loss、codebook loss、commitment loss、perceptual loss 与 GAN loss 五部分组成：

$$L_{\text{VQVAE}} = \|I - \hat I\|^2 + \|Q - \text{sg}[U]\|^2 + \alpha\|U - \text{sg}[Q]\|^2 + \beta L_{\text{LPIPS}}(I, \hat I) + \gamma L_{\text{GAN}}(U_Q)$$

其中第二项是 codebook loss，用来学习 codebook 向量；第三项是 commitment loss，让 encoder 输出靠近所选 code。LFQ/GQ 使用固定二值隐式 codebook，因此不再依赖传统 learnable codebook loss。

LFQ/GQ entropy loss（Eq. 4）

为防止 codebook utilization collapse，LFQ/GQ 引入 entropy loss：

$$L_{\text{Entropy}}(U) = \frac{1}{hw}\sum_{i,j} H(q(\mathbf{c} \mid U[i,j])) - \zeta H(\frac{1}{hw}\sum_{i,j} q(\mathbf{c} \mid U[i,j]))$$

第一项是 token entropy loss，鼓励单个 token assignment 更确定；第二项是 codebook entropy loss，鼓励整体 code 使用更均匀。$\zeta$ 控制两者权重。

GQ 的 grouped form

GQ 将 $U \in \mathbb{R}^{h \times w \times d}$ reshape 为 $U_G \in \mathbb{R}^{h \times w \times g \times d'}$，其中 $d = g \cdot d'$。每组使用独立二值 codebook $C_{\mathrm{GQ},k} = \{-1,1\}^{d'}$，把完整 $2^d$ 空间上的 entropy 计算拆成 $g$ 个 $2^{d'}$ 空间。

token entropy 的 grouped form 为：

$$L_{\text{Token Entropy}} = \frac{1}{hw}\sum_{i,j}\sum_{k=1}^{g} H(q_G(c_k \mid U_G[i,j,k]))$$

codebook entropy 则使用 grouped conditional probabilities 的乘积分解近似：

$$\frac{1}{hw}\sum_{i,j} q(\mathbf{c} \mid U[i,j]) \approx \prod_{k=1}^{g} \frac{1}{hw}\sum_{i,j} q_G(c_k \mid U_G[i,j,k])$$

这样做的复杂度从 $O(2^d)$ 降到 $O(g \cdot 2^{d'})$。当 $1 \lt g \lt d$ 时，GQ 介于 LFQ 与 BSQ 之间：比 LFQ 省显存，比 BSQ 保留更多组内 bit 相关性。

GD 阶段 GAN loss 改变

Stage 2 打开 Generative Decoder 后，decoder 不再只输入 $U_Q$ 做 deterministic reconstruction，而是输入 $(z, U_Q)$，其中 $z \sim \mathcal{N}(0,I)$。对应的 GAN loss 改为条件生成形式：

$$L_{\text{GAN}} = \mathbb{E}_{z \sim \mathcal{N}(0,I)}[\log(1 - D(G(z, U_Q)))]$$

直观地说，$U_Q$ 负责约束结构和语义，$z$ 负责让 decoder 在 $p(I \mid U_Q)$ 的条件分布中补出具体高频细节。

Loss 结构：从 VQ-VAE 到 GQ/GD

VQ-VAE 的标准训练目标包含五部分（原文 Eq. 2）：

$$L_{\text{VQVAE}} = \|I-\hat I\|^2 + \|Q-\text{sg}[U]\|^2 + \alpha\|U-\text{sg}[Q]\|^2 + \beta L_{\text{LPIPS}}(I,\hat I) + \gamma L_{\text{GAN}}(U_Q)$$

其中第二项是 codebook loss（学习 codebook 向量），第三项是 commitment loss（让 encoder 输出靠近所选 code）。LFQ/GQ 使用固定的二值隐式 codebook，因此不再需要 learnable codebook loss。它们引入 entropy loss 来防止 codebook utilization collapse（原文 Eq. 4）：

$$L_{\text{Entropy}}(U) = \frac{1}{hw}\sum_{i,j} H(q(\mathbf{c}|U[i,j])) - \zeta H\Bigl(\frac{1}{hw}\sum_{i,j} q(\mathbf{c}|U[i,j])\Bigr)$$

第一项 token entropy loss 鼓励单个 token assignment 更确定，第二项 codebook entropy loss 鼓励整体 code 使用更均匀。GQ 的贡献是在 grouped space 中近似计算这两项（Eq. 6、8），从而避免 LFQ 在大 $d$ 下的显存爆炸。

GD 阶段进一步把 GAN loss 改为条件生成形式（原文 Eq. 10），使 decoder 学习 $p(I \mid U_Q)$ 而不是只学习条件均值。

GD：Generative Decoder

训练目标转变

传统 discrete tokenizer 的 decoder 输入只有 $U_Q$，输出 deterministic reconstruction，训练目标是最小化与 ground truth 的距离。

GD 的 decoder 输入为 $(z, U_Q)$，其中 $z \sim \mathcal{N}(0,I)$ 沿 channel 维度与 $U_Q$ concatenation，输出从条件分布 $p(I \mid U_Q)$ 中采样的图像。

Two-stage Training

  • Stage 1：用 reconstruction loss + perceptual loss + GQ entropy loss（token entropy + grouped codebook entropy，对应原文 Eq. 2、6、8）训练 tokenizer，不含 $z$
  • Stage 2：将 decoder 的 conv-in layer 扩展 channel dimension 以接收 $z$，新增 channel 使用 zero-initialization，保证 Stage 2 开始时 decoder 行为与预训练状态完全一致

GAN Loss

Stage 2 的 GAN loss 从 deterministic 改为 conditional：

$$L_{\text{GAN}} = \mathbb{E}_{z \sim \mathcal{N}(0,I)}\left[\log(1-D(G(z, U_Q)))\right]$$

$U_Q$ 作为 condition，$z$ 提供额外随机性，使 decoder 可以建模条件分布而非条件期望。

符号表

符号定义
$I$输入图像，$I \in \mathbb{R}^{H \times W \times 3}$
$E$encoder
$G$decoder / generator
$D$discriminator
$U$encoder latent，$U = E(I)$，$U \in \mathbb{R}^{h \times w \times d}$
$U_G$reshape 后的 grouped latent，$U_G \in \mathbb{R}^{h \times w \times g \times d'}$
$d$latent channel 维度
$g$group number
$d'$每组 channel 数，$d = g \cdot d'$
$C_{\mathrm{GQ},k}$GQ 第 $k$ 组 codebook，$C_{\mathrm{GQ},k} = \{-1,1\}^{d'}$
$Q$量化后的离散 codes
$U_Q$$U + \text{sg}[Q - U]$，backward gradient 只关于 $U$
$\text{sg}[\cdot]$stop-gradient operation
$z$GD 噪声变量，$z \sim \mathcal{N}(0,I)$
$H(X)$entropy，$H(X) = -\sum_{X_i \in X} X_i \log X_i$
训练

数据集

数据集规模用途
ImageNet training set1.2Mablation 和 ImageNet SOTA 模型训练
General-domain dataset400M大规模预训练，覆盖多个压缩率

评估数据集：ImageNet 50k validation set、MS-COCO 2017 validation set、TokBench。

超参数（ablation 默认设置）

参数值
训练步数250K
optimizerAdam，$\beta_1,\beta_2=0.5,0.9$
weight decay0
learning rate$1e^{-4}$
LR scheduleconstant（warm-up + cosine decay 被发现 suboptimal）
warmup steps0
cos decay end ratio1
per GPU batchsize16
global batchsize128
训练图像尺寸$256 \times 256$，random crop
EMATrue
channel mult$[1,1,2,2,4]$
base channel128
residual blocks4
generative decoder (ablation)False

GQ/LFQ/BSQ 对照设置

参数GQLFQBSQ
$g$2116
$d'$8161
downsample16×1616×1616×16

Large-scale ImageNet SOTA 设置

参数16×8×
training dataIN-1K training setIN-1K training set
downsample16×168×8
$g$44
$d'$88
channel num256256
residual blocks44
total steps400K350K
per GPU batchsize88
global batchsize10241024
GPU number128 H20128 H20
generative decoderFalseFalse

General-domain Stage-1 / Stage-2 训练设置（768× compression ratio SOTA）

general-domain SOTA 模型使用 400M general-domain dataset，并采用两阶段训练：Stage-1 训练普通 tokenizer，Stage-2 打开 generative decoder。

参数Stage-1Stage-2
training datageneral-domain datasetgeneral-domain dataset
image size$256 \times 256$$256 \times 256$
downsample32×3232×32
$g$44
$d'$88
optimizerAdam $\beta_1,\beta_2=0.5,0.9$Adam $\beta_1,\beta_2=0.5,0.9$
learning rate$1e^{-4}$, constant$1e^{-4}$, constant
total steps550K550K
channel mult$[1,1,2,2,4,8]$$[1,1,2,2,4,8]$
base channel256256
residual blocks44
generative decoderFalseTrue
per GPU batchsize66
global batchsize10561056
GPU number176 H20176 H20

Stage-2 与 Stage-1 的核心差异是打开 generative decoder，并通过扩展 decoder conv-in channel 接收 $z$。新增 channel zero-initialization，使 Stage-2 初始行为与 Stage-1 预训练 decoder 对齐。

General-domain Stage-1 / Stage-2 训练设置（768× compression）

参数Stage-1Stage-2
training datageneral-domain datasetgeneral-domain dataset
image size256×256256×256
downsample32×3232×32
$g$44
$d'$88
optimizerAdam $\beta_1,\beta_2=0.5,0.9$Adam $\beta_1,\beta_2=0.5,0.9$
learning rate$1e^{-4}$, constant$1e^{-4}$, constant
total steps550K550K
channel mult[1,1,2,2,4,8][1,1,2,2,4,8]
base channel256256
residual blocks44
generative decoderFalseTrue
per GPU batchsize66
global batchsize10561056
GPU number176 H20176 H20

Stage-2 的核心差异是打开 generative decoder，新增 channel 使用 zero-initialization 保证初始行为与 Stage-1 对齐。

计算成本

  • 最佳架构 encoder 参数量：198M
  • 最佳架构 decoder 参数量：261M
  • WeTok-AR-XL 参数量：1.5B
  • Large-scale 训练使用 128-176 H20 GPU
  • 训练总时长：原文未明确给出
  • 推理 latency/FLOPs：原文未明确给出
实验

主要评估指标

rFID (recognition FID)、PSNR、SSIM、LPIPS、下游生成任务：FID、IS、Precision、Recall。

ImageNet 50k zero-shot reconstruction

方法Token 类型压缩比rFID ↓
WeTokDiscrete400%0.12
FLUX-VAEContinuous-0.18
SD-VAE 3.5Continuous-0.19
Open-MAGVIT2-I-PTDiscrete4391.39

（注：Open-MAGVIT2 rFID 1.17 来自 Table 3 ImageNet in-distribution 训练设置；此处为零-shot 设置的数值。）

ImageNet 50K reconstruction（ImageNet 训练）

这一组表格对应 ImageNet in-distribution 训练设置，说明 WeTok 不只是 general-domain zero-shot 强，在 ImageNet 训练集上也能刷新 reconstruction SOTA。

方法Token resolutionDownsampleCodebook sizerFID ↓PSNR ↑Codebook usage ↑
VQGAN16×161610248.3019.51-
VQGAN16×1616163844.9920.00-
LlamaGen16×1616327682.2620.5985%
Open-MAGVIT216×1616$2^{18}$1.1722.64100%
MGVQ16×1616$2^{52}$0.6423.71100%
BSQ1D 1024-$2^{36}$0.4528.14100%
WeTok16×1616$2^{32}$0.6124.50100%
WeTok32×328$2^{32}$0.1929.69100%

Zero-shot reconstruction 多压缩率对比

这一组结果来自 general-domain 训练后的 zero-shot reconstruction，对比 WeTok 在 768×、384×、192×、48×、24× 等不同压缩率下与 Cosmos、SD-VAE 3.5 的表现。

方法类型训练数据压缩比MS-COCO rFID ↓ImageNet rFID ↓ImageNet PSNR ↑ImageNet SSIM ↑
WeTokDiscrete400M7688.943.4920.770.55
CosmosDiscrete-3847.232.5220.490.52
WeTokDiscrete400M3845.300.8121.990.58
WeTokDiscrete400M1923.800.4023.750.67
WeTokDiscrete400M482.090.1827.540.82
SD-VAE 3.5Continuous-61.640.2428.390.86
WeTokDiscrete400M241.460.1229.510.88

关键观察是：WeTok 在 24× compression 时可以接近甚至超过 leading continuous VAE 的 rFID；在 768× compression 时仍能优于 Cosmos 的高压缩 discrete tokenizer。

High-compression comparison

方法压缩比ImageNet rFID ↓
WeTok768×3.49
Cosmos384×4.57

（注：原文 Introduction 第 30 行写 3.59，Abstract 与 Table 4 为 3.49；此处以 Table 4 为准。）

Class-conditional image generation (256×256 ImageNet)

模型参数量FID ↓IS ↑Precision ↑Recall ↑
WeTok-AR-XL1.5B2.31276.550.840.55
Open-MAGVIT2-AR-XL1.5B2.33271.770.840.54
LlamaGen-XL775M3.39227.080.810.54
DiT-XL/2675M2.27278.20.830.57
效果

核心效果

  1. 在 ImageNet 50k zero-shot 上，WeTok 以 400% 压缩比实现 rFID 0.12，优于 FLUX-VAE (0.18) 和 SD-VAE 3.5 (0.19)
  2. 在 768× 压缩比下 rFID 3.49，优于 Cosmos 在一半压缩率（384×）下的 4.57
  3. 以 WeTok 作为 tokenizer 的 LlamaGen 在 ImageNet 256×256 class-conditional 生成上 FID 2.31，优于使用 Open-MAGVIT2 tokenizer 的 LlamaGen (FID 2.33)

Memory ablation（$d'=8$）

方法$d=8$$d=16$$d=24$$d=32$$d=40$
LFQ10.5GB10.6GBOOMOOMOOM
BSQ10.5GB10.5GB10.6GB10.6GB10.6GB
GQ10.5GB10.6GB10.6GB10.6GB10.6GB

LFQ 在 $d=24$ 起即 OOM，无法扩展；GQ 与 BSQ 内存相当，但性能更优。

Stage 2 (GD) ablation

配置rFID ↓LPIPS ↓SSIM ↑PSNR ↑
Stage1 only5.370.170.5420.53
Stage1 + Stage23.900.160.5520.72

去掉 GD，rFID 差 1.47，是所有模块中掉点最多的，说明 GD 对高压缩重建贡献最大。

Generative Decoder 在 MS-COCO 上的定性消融对比
图 7：GD 的定性消融。Stage1-only 的重建更容易变成平均化纹理；加入 GD 后，模型能在 $U_Q$ 条件约束下采样更自然、更锐利的高频细节。

组件贡献总结

组件 / 选择对照主要影响已知数值
GQLFQ / BSQ在接近 BSQ 显存的情况下保留组内 joint modelingLFQ 在 $d=24$ OOM；GQ 在 $d=40$ 仍约 10.6GB
GDStage1 only高压缩下恢复纹理和细节，最大增益来源rFID 5.37 → 3.90；LPIPS 0.17 → 0.16；PSNR 20.53 → 20.72
constant LRwarm-up + cosine decaydiscrete tokenizer 训练中 constant LR 更稳定原文 Figure 6 显示 constant schedule 优于 warm-up + cosine decay
general-domain dataImageNet only提升泛化与 PSNR/SSIM原文：distribution gap 可能导致 in-distribution rFID/LPIPS 更差
model sizesmaller encoder/decoderbase channel 256 + residual blocks 4 最优encoder 198M，decoder 261M
group number $g$$g=1,2,4,16,32$$g$ 太小接近 LFQ 显存高；$g$ 太大接近 BSQ 误差大大规模配置中 $g=2/4/8$ 随压缩率变化

失败案例与负面发现

  1. LFQ 直接扩大 codebook 不可行：当 latent dimension 增大到 $d=24$，LFQ 的 entropy loss 显存开销已经 OOM，说明 lookup-free 本身不足以解决 scalable codebook，还必须处理 entropy 计算。
  2. BSQ 的 bit-wise independence 太强：BSQ 可以把显存降到近似常数，但逐 bit 独立假设丢掉了 bit 间相关性，带来比 GQ 更大的 approximation error。
  3. warm-up + cosine decay 并非默认最优：论文发现对 discrete tokenizer，constant learning rate 优于 warm-up + cosine decay。这说明 continuous VAE / diffusion 中常用的 schedule 不能直接迁移。
  4. general-domain data 不是所有指标都更好：400M general-domain 训练提升泛化和 PSNR/SSIM，但在 ImageNet in-domain 的 rFID/LPIPS 上可能因 distribution gap 变差。
  5. 小规模 AR 下游模型收益有限：WeTok 作为 tokenizer 的优势需要足够大的 AR 模型释放，小参数 AR 模型上不一定优于 Open-MAGVIT2-based AR。
  6. GD 可能牺牲严格忠实性：GD 通过采样补细节，高压缩视觉质量更好，但不保证像素级 faithful reconstruction，医学、遥感、司法图像压缩等场景需要谨慎。
  7. 高压缩下 trade-off 仍存在：768× rFID 3.49 仍远高于 16× 的 0.12，压缩-保真 trade-off 没有消失，只是 Pareto frontier 往前推了。
结论

WeTok 通过两项简洁设计解决了离散 tokenizer 的压缩-保真矛盾：

  1. GQ：把 lookup-free quantization 分组，在保留组内联合建模能力的同时，将 entropy loss 的内存复杂度从 $O(2^d)$ 降到 $O(g \cdot 2^{d'})$，且在非退化情形下 approximation error 小于 BSQ
  2. GD：引入带噪声 prior 的生成式 decoder，把 decoder 从 deterministic reconstruction 改为条件生成模型，从条件分布中采样而非输出条件期望

实验验证了这两项设计在高保真和高压缩两个场景下的有效性，并已开源代码和预训练模型。

未来方向（论文明确提及）：

  • 视频 tokenizer：WeTok 目前主要验证图像，视频场景需要时序建模扩展
  • 与下游生成模型更深度结合：当前 AR 生成只是初步验证，更大的 vocabulary 建模策略（如 asymmetric token factorization）值得探索
  • 语义增强：UniWeTok 已开始解决 binary codebook tokenizer 语义抽取不足的问题
收获
  1. 分组近似的工程价值：当遇到 $O(2^d)$ 复杂度的瓶颈时，检查是否可以按某种结构（channel/space/time）分组，把指数复杂度转为线性组合。GQ 的关键是：分组造成的 approximation error 比直接做 independence assumption 更小——这需要理论支撑，不能只靠直觉。
  2. Decoder 从 deterministic 到 generative 的转变值得重视：在高压缩场景中，如果 latent 不够承载所有细节，与其让 decoder 输出"平均"，不如让它学习在给定 latent 条件下的分布，并从中采样。这个思路在 consistency model、diffusion decoder 中都有体现，WeTok 将其迁移到离散 tokenizer 并用 GAN 实现 single-step 采样，效率上有优势。
  3. Two-stage training + zero-initialization 是个好技巧：当需要在已训练的模型上引入新模块时（如 GD 的 $z$ 输入），直接添加新 channel 并从零初始化会让模型在训练初期行为与预训练状态一致，避免新模块破坏已学到的能力。这个技巧在 VQGAN 的 codebook initialization、LLM 的 adapter 等场景中也常用。
  4. Constant LR 对离散 tokenizer 可能更合适：这与直觉相反（通常 warm-up + cosine 是 standard），但可能是因为 discrete token 的 optimization landscape 与 continuous VAE 不同。这个发现值得在训练其他离散 tokenizer 时尝试 constant LR。
  5. Tokenizer 是视觉生成的上游瓶颈：MAGVIT-v2、Open-MAGVIT2、WeTok、UniWeTok 都强调这个判断。对于实际系统设计，与其不断改进生成模型主干，不如先确保 tokenizer 的压缩-保真 Pareto frontier 是否已经足够好——上游瓶颈没解决，下游改进事倍功半。
Categories:AI / 视觉分词器
Tags:visual-tokenizer, discrete-tokenization, GQ, GD