--- license: other tags: - image-tokenizer - video-tokenizer - autoencoder - dinov3 --- # MAVT v2 — Grain-Prior tokenizer (checkpoints + notes, 2026-10-30) Research checkpoints of the Grain-Prior C-D tokenizer from [Computer-Vision-Lab-LinhDang/Antokenizer](https://github.com/Computer-Vision-Lab-LinhDang/Antokenizer) (`src/mavt/grain_prior`, branch `feat/grain-splat-decoder`, code at commit `a5ab394`; the flags are documented in `src/mavt/grain_prior/README.md`). Work in progress, not a release. ## What the model is A 256×256 image is encoded on top of a **frozen DINOv3 ViT-L/16** into a continuous state of **11,300 numbers** (the comparison ceiling is 12,288 = 16 × 16 × 48, AToken's image latent): | Part | Shape | Numbers | Role | |---|---|---|---| | Grain geometry | 16 × 16 × 4 (u 2, ρ 1, σ 1) | 1,024 | position in the cell, log inverse depth, radius | | Grain content `f` | 16 × 16 × **32** | 8,192 | local appearance (`d_f = 32`) | | Coarse memory `c_0` | 4 × 4 × 128 | 2,048 | layout / colour / semantics of 64 px regions | | Header | 36 | 36 | global scalars | Decoder (`rgb_decoder=splat`): a geometry-biased context transformer over the grains and coarse slots, per-grain frequency atoms, a depth-aware Gaussian splat onto a 64 × 64 query grid, then a pointwise RGB head with a coarse start image. Video adds per-frame records (camera, motion, innovation slots) on top of the anchor state. **DINOv3 weights are not included** (gated licence). The loader skips the missing `stems.dino.model.*` keys; pass your own copy with `dino_repo=`. ## Files | File | Step | What | Use | |---|---|---|---| | `image/bs_step38113_ema.pt` | 38,113 | **Best image model** (image-only from scratch, stopped mid-decay), EMA | eval / inference | | `image/bs_step38113.pt` | 38,113 | same, raw weights | `--init` for further stages | | `image/bs_step35000.pt` | 35,000 | same run, before the LR decay, raw | `--init` for a fresh schedule | | `image_video/s2_step5500_ema.pt` | 5,500 | stage 2 (image + video) from `bs_step38113`, EMA | eval | | `image_video/s2_step5500_resume.pt` | 5,500 | stage 2 full state (optimizer, EMA, scheduler) | `--resume` | | `ablations/df48_ema.pt` | +3,000 | `d_f = 48` (15,396 numbers), widened from `bs_step38113` | rate study | | `ablations/df64_ema.pt` | +3,000 | `d_f = 64` (19,492 numbers) | rate study | All checkpoints: `cfg` keys `d_f=32` (48 / 64 for the ablations), `rgb_decoder=splat`, `query_grid=64`, `rgb_head=pointwise`, `w_lpips=0.1`. ```python from mavt.grain_prior import load_grain_prior model, cfg, step = load_grain_prior("image/bs_step38113_ema.pt", "cuda", dino_repo="") out = model(x, "image", return_outputs=True) # x: [B, 3, 256, 256] in [-1, 1] rgb = out["outputs"]["rgb"] ``` ## Results (ImageNet val, first 1,000 images, 256², centre crop) | Model | Numbers | PSNR ↑ | SSIM ↑ | LPIPS ↓ | rFID ↓ | |---|---|---|---|---|---| | `bs_step38113_ema` | 11,300 | **27.76** | 0.790 | 0.130 | 17.6 | | `bs_step38113` (raw) | 11,300 | 27.69 | 0.788 | 0.131 | 17.1 | | `s2` step 3,500 (EMA, image after adding video) | 11,300 | 27.58 | 0.784 | 0.136 | 18.0 | | `df48_ema` | 15,396 | 28.16 | 0.804 | 0.121 | 16.3 | | `df64_ema` | 19,492 | 28.43 | 0.813 | 0.114 | 15.5 | | *reference: AToken-So/c (paper)* | 12,288 | 29.72 | 0.848 | 0.085 | 0.21 | rFID on 1,000 images is high-variance and not comparable to 50k-image rFID; use it only between rows of this table. ## What was learned this round (all 3k-step probes from `bs_step38113`, lr 1e-4, same protocol, EMA PSNR) | Probe | PSNR | Δ vs control | |---|---|---| | control | 27.74 | — | | more frequency atoms per grain (6 → 12 / 24) | 27.67 / 27.65 | 0 | | atoms evaluated at the pixels (no 64-grid aliasing) | 27.74 | 0 | | + pixel ResBlocks at 128 / 256 | 27.75 | +0.01 | | + ConvNeXt context at 64 | 27.77 | +0.03 | | + Open Images 50/50 for 5k steps | 27.72 | 0 | | **`d_f` 35 / 48 / 64** | **27.83 / 28.16 / 28.43** | **+0.09 / +0.42 / +0.69** | * The decoder is not the bottleneck: three ways of adding decoder capacity at the pixel level change nothing, although the new modules do train. * The information in the state is: PSNR grows by about +1 dB per doubling of `d_f` (raw: +0.13 / +0.61 / +0.95), and SSIM / LPIPS / rFID improve with it. * Channel ablations at step 20k: replacing `c_0` by its mean → 15.0 dB, shuffling `f` across cells → 13.3 dB, `f` uses all 32 dimensions (first PC 16 % of variance), 94 % of the remaining error energy is above 1/8 cycle/px. * Adding video (stage 2) costs about 0.2 dB on images early on; part of it is the LR going back up from a decayed point. * Remaining error is high-frequency texture (fur, foliage, small text). Perceptual losses (edge / LPIPS / Gram) are expected to help LPIPS and rFID rather than PSNR; flags `w_edge` and `feat_sg_f` exist at `a5ab394`, not yet trained. ## Training recipe of `bs_step38113` Image only, ImageNet-1k train, 4 GPUs, batch 256, Muon / AdamW, lr 2e-4, warmup 1k, WSD with the decay over the last 4k of 40k steps (stopped at 38,113, mid-decay), loss = L1 + 0.1 LPIPS + DINOv3 feature loss (+ init / code terms), about 11 h. Stage 2 adds OpenVid + SSv2 clips (batch 128 images + 32 clips).