dino-fsq-sweep

DINO-FSQ patch quantizer heads

Trained FSQ quantizer heads for a DINOv3 ViT-B/16 patch tokenizer, from the ar_wam_rl world-model codebase. Each head takes a 768-d patch feature and squeezes it through a finite scalar quantization lattice, producing one integer code per patch plus a reconstruction.

Six heads are released — the best config at each (mode, rate) point on the rate-distortion frontier. Deliberately-bad control arms from the sweep (the as-shipped default, a collapsed head) are not included.

The two modes

delta dominates feat at every rate below ~37 bits, because most of a patch feature is static scene content that the previous frame already supplies.

Results

1-step (h1) and 3-step closed-loop (h3) reconstruction, decoded to RGB through the frozen DINOv3 pixel decoder, on a held-out multiview manipulation val set. PSNR (dB) / SSIM.

Reference points on the same val set and decoder:

Read these numbers with the floor in mind. PSNR here is dominated by static scene content, which every head reconstructs well; the frame-to-frame motion is what they mostly miss. Against the 16.56 dB copy-previous-frame floor and the 25.56 dB ceiling, the best head closes roughly a quarter of the gap. feat_16bit_MLP sits below the copy floor — a 16-bit independent code per patch is a harder problem than it sounds.

Architecture

x (768-d)  ->  standardize (fixed affine, mu/sd in the state dict)
           ->  enc     Linear(768, d)          if hidden is None
                       or Linear(768,h) GELU Linear(h,d)
           ->  FSQ lattice, per-dim, straight-through   -> integer code in [0, prod(levels))
           ->  dec     Linear(d, 768)          if hidden is None
                       or Linear(d,h) GELU Linear(h,768)

All released heads use scheme="ifsq" (equal-width cells on [-1,1] via a sigmoid squash, which maps Gaussian activations to near-uniform codes) and input_norm="standardize". Nominal rate is log2(prod(levels)) bits/patch.

Two notes for anyone inspecting the weights:

• mu / sd are buffers carrying the training-set feature statistics, so a head restored from a checkpoint needs no access to the training data.
• q.proj_in.* / q.proj_out.* are untrained dead weight. The head uses only FSQ's lattice and substitutes its own enc/dec; FSQ's internal projections are bypassed and frozen. They are kept in the state dict so checkpoints load strictly. Do not read anything into their values.

Usage

QHead and the loader live in models/dinofsq_head.py in the ar_wam_rl repo (models/quantizers.py supplies the FSQ lattice). Clone it, then:

from models.dinofsq_head import CONFIGS, from_pretrained

print(sorted(CONFIGS))                                               # the six released names
head, cfg = from_pretrained("delta_37bit_MLP_d16L5", device="cuda")   # -> (module, config)

# feats: (B, T, N, 768) DINOv3 patch features (models/dinov3.py)
prev = feats[:, 0]
for t in range(1, feats.shape[1]):
    v_hat, codes, _ = head(feats[:, t] - prev)   # codes: (B, N) int64
    prev = prev + v_hat                          # closed loop: reference is the reconstruction

For feat mode, call head(feats[:, t]) and use the returned reconstruction directly.

Each .pt is a bare state_dict; configs.json holds the architecture config keyed by the same names, which is what models/dinofsq_head.py uses to rebuild the module before loading.

Training

Heads were fit on cached DINOv3 features (the only trainable part of the tokenizer), AdamW with a 200-step linear warmup, log-cosh reconstruction loss, no weight decay. Per-head lr and steps are in configs.json — 8k steps at 1e-3 for the linear heads, 16k at 1e-3 for the MLP heads. MLP heads are sensitive to learning rate: at 1e-2 the encoder collapses to a constant output and the effective vocabulary falls to a handful of codes.

Limitations

• 768-d DINOv3 ViT-B/16 features only; not transferable to another backbone width without retraining.
• Trained and evaluated on a multiview manipulation dataset. Behaviour on other domains is untested.
• Effective code usage is well below nominal: the 37-bit heads use ~61k distinct codes on the val set, so the usable rate is far under the nominal figure.
• These are reconstruction heads, not generative models. They define a discrete vocabulary for an autoregressive predictor to model; they do not predict future frames themselves.