anbinh39/MAVT_V2

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 (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.

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).