| --- |
| 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=<local dir or facebook/dinov3-vitl16-pretrain-lvd1689m>`. |
|
|
| ## 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="<dinov3-vitl16 dir>") |
| 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). |
| |
