WeTok: Powerful Discrete Tokenization for High-Fidelity Visual Reconstruction

TL;DR

WeTok is a discrete visual tokenizer from SJTU + WeChat Vision that claims to close the long-standing reconstruction-quality gap to continuous VAE tokenizers (SD-VAE 3.5, FLUX-VAE) at comparable spatial compression. Two pieces: (1) Group-wise lookup-free Quantization (GQ) partitions the channel dimension of the latent into G groups and runs an LFQ-style implicit codebook on each group, which lets the effective codebook scale to enormous sizes (the paper uses very large per-group bit budgets at the 8× compression setting) without the memory blowup of standard LFQ’s entropy loss. (2) Generative Decoder (GD) appends a Gaussian noise channel to the quantized latent and adds a GAN-style conditional generation objective on top of the standard reconstruction loss, so the decoder learns a distribution over images conditioned on the discrete code rather than a deterministic point estimate — which matters at high compression where one code corresponds to many plausible images. On ImageNet-50K at 8× compression, WeTok hits rFID 0.12 (vs SD-VAE 3.5 0.19, FLUX-VAE 0.18, all continuous); the AR generator built on WeTok (WeTok-AR-XL @ 1.5B) hits gFID 2.31 on ImageNet class-conditional, beating Open-MAGVIT2-AR-XL at matched scale.

Key claims

The “discrete tokenizers must trade compression for reconstruction” trade-off is not fundamental; WeTok matches or beats continuous SD-VAE 3.5 / FLUX-VAE on rFID at the same 8× spatial compression on ImageNet-50K [Tab. 3].
LFQ’s entropy loss [Eq. 4] has GPU memory cost linear in codebook size; this is what bounds prior LFQ tokenizers at codebooks of ~2^18 entries in practice [Tab. 2].
GQ rewrites both the token-entropy and codebook-entropy losses as sums over G independent groups [Eq. 6, Eq. 8], reducing the per-loss variable space from the full codebook size to the much smaller per-group bit budget — eliminating entropy-loss memory as the bottleneck [§3.2, Tab. 2].
GQ interpolates between LFQ (G=1, exact, memory-heavy) and BSQ (G=number-of-bits, independence assumption per bit, lossy); choosing an intermediate G recovers most of LFQ’s accuracy while taking BSQ-level memory [§3.2].
The Generative Decoder is a two-stage recipe: stage 1 is standard reconstruction training; stage 2 expands the decoder’s conv_in to accept a sampled Gaussian noise channel (newly added weights zero-initialized for stability) and adds a conditional GAN objective [Eq. 10] so the decoder models p(image | code) [§3.3, Tab. 2 (Stage1 vs Stage1+Stage2: rFID 5.37 → 3.90)].
WeTok zero-shot rFID @ 8× compression on ImageNet-50K: 0.12, vs FLUX-VAE 0.18, SD-VAE 3.5 0.19, UniTok (16× ratio) 0.41, BSQ 3.81 at higher compression [Tab. 3].
WeTok at 32× compression (768× compression ratio) achieves rFID 3.49 on ImageNet-50K, beating Cosmos (16×, ratio 384) at 4.57 — i.e. WeTok at 2× higher compression still beats Cosmos [Tab. 3].
WeTok-AR-XL @ 1.5B on ImageNet class-conditional 256×256 reaches FID 2.31 / IS 276.55, beating Open-MAGVIT2-AR-XL (1.5B) at FID 2.33 with the same tokenizer-AR recipe [Tab. 5].
For tokenizer training, a constant learning rate outperforms the standard warmup + cosine-decay schedule for discrete tokenizers — an unusually clean negative result against the default LM schedule [§4.2, Fig. 6].
Training data ablation finds an in-distribution vs out-of-distribution tradeoff: a 400M general-domain dataset improves SSIM/PSNR/validation-loss vs 1.2M ImageNet but hurts rFID/LPIPS on the in-distribution ImageNet validation set — distribution gap, not capacity, dominates rFID at scale [§4.2, Fig. 6].
Token-count vs decoder-parameter scaling: increasing group count G monotonically improves reconstruction without hitting LFQ’s memory wall [§4.2, Fig. 4]; the optimal encoder/decoder has 256 base channels and 4 residual blocks (198M / 261M params) [§4.2, Fig. 5].

Method

WeTok inherits Open-MAGVIT2’s CNN encoder/decoder/discriminator architecture and trains in two stages.

Stage 1 (reconstruction). Given an image $x$ , the encoder produces a per-spatial-position latent $z \in \mathbb{R}^d$ , which is split along the channel dimension into G equal groups $\{z_1, \ldots, z_G\}$ , each of dim $d/G$ . Each group is quantized independently via sign-binarization (the LFQ rule from MAGVIT-v2), giving a binary code in $\{-1, +1\}^{d/G}$ per group. The effective discrete vocabulary per spatial position is $2^d$ — but crucially the entropy losses are computed per group. Token entropy becomes a sum of G entropies over $2^{d/G}$ -sized spaces (Eq. 6); codebook entropy becomes the same sum under the independence approximation $\mathbb{E}[\sum \log p_k] \approx \sum \log \mathbb{E}[p_k]$ (Eq. 7). G is the dial: G=1 = LFQ (exact, expensive), G=d = BSQ (every bit independent, lossy), intermediate G = sweet spot. Total loss = reconstruction L1 + LPIPS + GAN + grouped token entropy + grouped codebook entropy + commitment.

Stage 2 (generative decoder). Freeze the encoder + quantizer. Expand the decoder’s first conv to accept extra channels carrying a sampled $z' \sim \mathcal{N}(0, I)$ concatenated to the discrete code (newly added input weights zero-initialized so the stage-2 decoder starts behaviorally identical to stage 1). Re-train the decoder + discriminator with the same reconstruction stack plus a conditional GAN loss where $z'$ is the generator noise and the discrete code is the condition (Eq. 10). The decoder now learns p(image | code) rather than the conditional expectation, which gives sharper, more realistic high-frequency detail at high compression.

AR generation. For class-conditional ImageNet, the WeTok codes feed LlamaGen’s transformer at standard scales (L 343M, XL 775M, XXL 1.5B). No architectural change to the AR side; the win comes from a better tokenizer.

Results

Setting	Tokenizer	rFID (ImageNet-50K, resize)	Notes
16× compression, ratio 192	WeTok	0.60	beats SD-VAE 1.x discrete 1.13, Open-MAGVIT2-I-PT 1.67
8× compression, ratio 48	WeTok	0.20	beats SD-VAE 2.x 0.70, SDXL-VAE 0.67, UniTok (ratio 64) 0.41
8× compression, ratio 24	WeTok	0.12	beats SD-VAE 3.5 0.19, FLUX-VAE 0.18 — first discrete tokenizer to beat both at parity
32× compression, ratio 768	WeTok	3.49	vs Cosmos (16×, ratio 384) 4.57 — WeTok at 2× the compression still wins

ImageNet class-conditional 256×256 AR generation:

WeTok-AR-XL (1.5B): FID 2.31 / IS 276.55 [Tab. 5].
Open-MAGVIT2-AR-XL (1.5B): FID 2.33 / IS 271.77.
VAR-d30 (2.0B): FID 1.92 / IS 323.1.
DiT-XL/2 (675M): FID 2.27 / IS 278.2.

Ablations (Tab. 2): GD stage gives Δ rFID 5.37 → 3.90 at fixed encoder/quantizer; GQ memory holds at ~10.6 GB across all group settings while LFQ goes OOM beyond a moderate codebook size.

Why it’s interesting

For Luma, the practical question is whether a discrete tokenizer can match continuous-tokenizer reconstruction quality at the same spatial compression, so that AR-style generators can use it without paying a quality tax. WeTok is the cleanest positive datapoint to date — it beats FLUX-VAE / SD-VAE 3.5 head-to-head on ImageNet at matched 8× compression, with the AR-side benefit demonstrated end-to-end (WeTok-AR-XL > Open-MAGVIT2-AR-XL).

Direct contrast with the wiki’s two closest tokenizer papers. UniTok: A Unified Tokenizer for Visual Generation and Understanding (UniTok) and WeTok solve adjacent problems: UniTok’s multi-codebook quantization (MCQ) splits the latent vector into N chunks each with its own explicit learned codebook of size V (effective vocab $V^N$ ); WeTok’s GQ splits the channels into G groups each with an implicit lookup-free codebook of size $2^{d/G}$ (effective vocab $2^d$ ). MCQ scales by adding sub-codebooks; GQ scales by enlarging each per-group bit budget within the lookup-free regime. Both arrive at the same “very large effective vocabulary, but factorized” recipe by different routes — MCQ is the product-quantization analog (related to TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate‘s PQ tradition), GQ is the LFQ-with-grouping descendant. The user-flagged comparison to QLIP is worth foregrounding: WeTok’s reported per-spatial-position bit budget is large (Jiaming notes “2^32 codebook size”), which inherits QLIP’s lineage of “if your effective vocab is enormous, of course rFID drops”. The question to keep watching is whether AR side performance scales with this enormous discrete vocab — Tab. 5 says yes at 1.5B, but only barely beats Open-MAGVIT2-AR with a much smaller effective vocab, suggesting the rFID gains over-promise the gFID gains.

Also worth pairing with TC-AE: Unlocking Token Capacity for Deep Compression Autoencoders (TC-AE) which independently shows that for high-compression continuous tokenizers, structuring the bottleneck (staged token compression + iBOT) matters more than raw compression ratio for downstream generation. WeTok and TC-AE point at the same conclusion from opposite sides: discrete tokenizers can match continuous ones if you (a) factor the codebook to escape the memory wall (WeTok) or (b) factor the spatial compression to preserve semantics (TC-AE).