QuadTok Tokenizers

Content-adaptive QuadTok VQ image tokenizers trained on ImageNet-1K. This repo hosts multiple tokenizer variants and their checkpoints under a shared directory convention so we can keep adding new checkpoints over time.

Repository layout

Tokenizer/
  <variant>/                 # one tokenizer model family
    checkpoint-<step>/       # a checkpoint at a given global training step
      unwrapped_model/pytorch_model.bin   # raw (online) model weights  -> inference / warm-start
      ema_model/pytorch_model.bin         # EMA weights (decay 0.999)   -> use these for eval metrics
      model.safetensors                   # accelerate state: unwrapped model
      model_1.safetensors                 # accelerate state: discriminator (if GAN stage reached)
      optimizer.bin / optimizer_1.bin     # optimizer state (model / discriminator) -> resume
      scheduler.bin / scheduler_1.bin     # LR scheduler state -> resume
      random_states_*.pkl                 # per-rank RNG state -> exact resume
      config.yaml                         # full OmegaConf training/model config
      metadata.json                       # {"global_step": <step>}

Naming convention

• <variant> — the tokenizer family. Current variants:
  • 3level — 256×256, 3-level (8×8 → 16×16 → 32×32, finest patch size 8), token budget ~64 → ~1024.
  • 512 — 512×512, high-res 2-level (16×16 → 32×32, finest patch size 16), token budget ~256 → ~1024.
• checkpoint-<step> — <step> is the global training step (e.g. checkpoint-350000). A checkpoint-final alias may be added for the last step of a finished run.

What to download

• Inference only → unwrapped_model/pytorch_model.bin (or ema_model/… for the eval-quality weights) + config.yaml.
• Resume training → the whole checkpoint-<step>/ folder (accelerate load_state expects the model*.safetensors, optimizer*, scheduler*, random_states_* alongside).

---

Variants

Tokenizer/3level — 256×256, 3-level (finished, step 350,000)

3-level QuadTok VQ tokenizer, trained from scratch on ImageNet-1K at 256×256. Extends the 2-level QuadTok (8×8 → 16×16) with a third level of detail (32×32, patch size 8).

Final evaluation (EMA weights, step 350,000)

Metric	Value
rFID	0.881
PSNR	23.02 dB
Inception Score	209.3
Codebook usage	100%

• LODs: num_patch_side_list: [1,2,4,8,16,32], patch_size_list: [16,16,16,16,16,8]; tokens at lod ≥ 3.
• Encoder ViT-small, decoder ViT-large; token_size: 8, codebook 16384, L2-normalized codes.
• Losses: L2 + LPIPS(ConvNeXt-S) + VQ + LeCam from step 0; PatchGAN after step 200,000.
• 350k steps, global batch 512 (8×H200), LR 1e-4 cosine w/ 10k warmup, EMA 0.999, bf16.
• Training tree: random probabilistic quadtree per step, guaranteed_depth: 3, expansion_probs: [0.95, 0.7].

Tokenizer/512 — 512×512, high-res 2-level (training in progress)

High-res 2-level QuadTok VQ tokenizer, trained from scratch on ImageNet-1K at 512×512. The 256 2-level recipe (lod3 8×8 + lod4 16×16) shifted one level deeper: coarse = lod4 16×16 (256 tokens, mandatory) + fine = lod5 32×32 (up to 1024); tokens live at lod ≥ 4, finest patch size 16.

Intermediate evaluation (EMA weights) — run targets step 350,000; latest checkpoint here is checkpoint-280000.

Step	rFID	PSNR
50,000	2.209	22.50 dB
100,000	1.347	22.96 dB
150,000	1.200	23.16 dB
200,000	1.127	23.27 dB
250,000	0.878	22.22 dB

• LODs: num_patch_side_list: [1,2,4,8,16,32], patch_size_list: [16,16,16,16,16,16]; tokens at lod ≥ 4.
• Encoder ViT-small, decoder ViT-large; token_size: 8, codebook 16384, L2-normalized codes.
• Losses: L2 + LPIPS(ConvNeXt-S) + VQ + LeCam from step 0; PatchGAN after step 200,000.
• 350k steps target, global batch 512 (8×H200, per-GPU 32 × grad-accum 2), LR 1e-4 cosine w/ 10k warmup, EMA 0.999, bf16.
• Training tree: random probabilistic quadtree per step, guaranteed_depth: 4, expansion_probs: [0.73].

PSNR dips slightly after the GAN turns on at step 200k while rFID keeps improving — expected: the discriminator trades a little pixel-wise PSNR for better perceptual / distributional fidelity.

---

Loading (inference)

import torch, json
from omegaconf import OmegaConf
from modeling.quadtok import QuadTok   # from the quadtok repo

variant, step = "512", 280000
cfg = OmegaConf.load(f"Tokenizer/{variant}/checkpoint-{step}/config.yaml")
model = QuadTok(cfg)
sd = torch.load(f"Tokenizer/{variant}/checkpoint-{step}/ema_model/pytorch_model.bin", map_location="cpu")
model.load_state_dict(sd, strict=False)
model.eval()

Use unwrapped_model/pytorch_model.bin for the raw (non-EMA) weights.