Technical Summary: MergeTok

Problem Statement

Current visual tokenizers for image generation are bifurcated into two distinct families, each suffering from complementary structural limitations:

1. Continuous VAE-based tokenizers: While offering high-fidelity reconstruction, they rely on latent spaces lacking intrinsic organization. Under strong compression, semantic factors become entangled within dense latent representations, resulting in weakly disentangled codes that hinder semantic controllability.
2. Discrete VQ-based tokenizers: While enabling autoregressive generation, they face optimization bottlenecks. The non-differentiable nature of quantization leads to sparse and uneven gradient updates, causing many codebook entries to receive insufficient supervision. This increases the risk of codebook collapse and reduces representation efficiency.

These limitations create a persistent tension between reconstruction fidelity, semantic structure, efficiency, and training stability.

Methodology: MergeTok Framework

MergeTok proposes a unified dual-branch tokenizer that jointly optimizes continuous (VAE) and discrete (VQ) tokenization within a single encoder-decoder architecture. The core innovation is the use of Token Merging (ToMe) as a semantic bridge between the two branches.

Architecture Overview

Given an input image $X$, a shared CNN encoder $E_c$ produces a feature map $Z$. This map is processed by two parallel branches:

1. VAE Branch (Continuous):

   • Token Merging: An attention-based encoder equipped with ToMe modules compresses the token sequence from $L$ tokens to $K$ tokens ($K < L$) by merging similar tokens. This process generates a condensed representation $Z^{(vae)}_K$ and a source map $S$ that records the assignment of original tokens to merged groups.
   • Semantic Alignment: The merged tokens are aligned with a frozen DINO-style teacher model (also equipped with ToMe) using a cosine distance loss ($L_{align}$). This regularizes the continuous latent space toward semantic-aware, disentangled representations.
   • Reconstruction: A hybrid decoder utilizes both the condensed tokens and the source map $S$ to "unmerge" the groups and reconstruct the pixel-space image.

2. VQ Branch (Discrete):

   • Group-Aware Quantization: This branch bypasses the ToMe operation during the forward pass to preserve the full sequence length ($L$ tokens). However, it utilizes the source map $S$ generated by the VAE branch as a structural prior.
   • Regularization: The source map partitions the $L$ tokens into $K$ semantic groups. Two group-aware losses are applied to the codebook assignments:
     • Intra-group Diversity ($L_{div}$): Maximizes entropy within each group to prevent collapse to a single code.
     • Inter-group Exclusivity ($L_{cons}$): Minimizes overlap between groups to ensure different semantic groups utilize distinct code subsets.
   • Reconstruction: A standard VQ decoder reconstructs the image from the quantized tokens.

Training Strategy

• Unified Objective: The model is trained end-to-end with a total loss combining reconstruction, adversarial, and merge-aware regularizers:
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• Granularity-Aware Sampling: During training, the merge ratio (and thus the number of kept tokens $K$) is dynamically sampled from a discrete distribution centered around a hyperparameter $K^*$. This exposes the encoder to varying bottleneck widths, encouraging robustness across semantic granularities.
• Inference: At inference, the tokenizer always outputs a fixed sequence of 256 tokens. Downstream generators (both autoregressive and diffusion) operate on this fixed sequence and are agnostic to the internal merging structure.

Key Contributions

1. Unified Dual-Branch Tokenizer: MergeTok is the first architecture to jointly optimize continuous and discrete latents under a single objective, using token merging as a semantic interface to reinforce the stability of VAE training and the structural advantages of VQ.
2. Merge-Aware Training Constraints: The authors introduce two novel objectives derived from the merging process:
   • Merged-token alignment: Improves semantic structure in the continuous latent space.
   • Group-aware quantization: Stabilizes VQ training and improves codebook utilization via intra-group diversity and inter-group exclusivity.
3. Granularity-Aware Merge Ratio Sampling: A training-time strategy that samples different token granularities, enabling the model to learn unified tokenization that is robust across varying compression rates.

Experimental Results

Experiments were conducted on ImageNet-256 (256×256 resolution) with a fixed token budget of 256 tokens.

• Reconstruction Fidelity:
  • Discrete (VQ): MergeTok-BL achieved an rFID of 0.78, outperforming strong baselines like GigaTok-BL (0.81) and MergeVQ-GR (1.12). With a frozen DINO discriminator, MergeTok-BL* reached 0.48 rFID.
  • Continuous (VAE): MergeTok-BL achieved an rFID of 0.47, marginally outperforming RAE (0.49) and significantly beating SoftVQ-L (0.61) and DC-AE-f32 (0.69).
• Representation Quality:
  • MergeTok-BL attained 78.3% linear probing accuracy, the highest among listed discrete tokenizers, surpassing UniTok (70.8%) and MergeVQ-GR (77.9%). This indicates superior semantic organization.
• Downstream Generation:
  • Autoregressive (VQ Branch): Paired with LlamaGen-XXL, MergeTok-BL achieved a gFID of 1.93 (without CFG) and 2.13 (with CFG), outperforming GigaTok-BL (2.03).
  • Diffusion/Flow (VAE Branch): Paired with SiT-XL/1, MergeTok-BL achieved a gFID of 1.79 (without CFG) and 1.29 (with CFG), matching the best entries in the table.

Significance and Claims

The paper claims that MergeTok demonstrates that a single architecture can endow visual tokenizers with both robust semantic organization and generator-friendly discreteness.

• Unification: It resolves the trade-off between the high-fidelity reconstruction of VAEs and the discrete structure required for autoregressive generation.
• Semantic Bridge: By leveraging token merging as a structural prior, the model successfully transfers semantic grouping information from the continuous branch to stabilize the discrete branch, and vice versa.
• Compatibility: The tokenizer produces 256-token sequences that are compatible with both autoregressive (e.g., LlamaGen) and diffusion/flow (e.g., DiT, SiT) generators without requiring modifications to the generators themselves.

The authors position MergeTok not merely as a reconstruction tool, but as a unified interface that supports discrete autoregressive generation, continuous diffusion generation, and semantically structured representation learning simultaneously.