dingshizhe/molmo-ae-cc12m-latents

CC12M latents — VGT-AE with a frozen SigLIP2 (Molmo) encoder

Pre-encoded CC12M images as (32, 16, 16) float16 latents, paired with their captions. WebDataset format: 1097 tars, each member a <key>.npy + <key>.txt pair.

Why 512 px, and why this encoder

SigLIP2 is patch16 and the pixel_shuffle fold halves each side, so the encode stride is 32 — symmetric with the DC-AE f32c32 decoder. 512/32 = 16 gives the 256-token 16×16 latent grid. The earlier Qwen2.5-VL VGT-AE is patch14 → stride 28 against the same stride-32 decoder, and has to paper over the mismatch with an F.interpolate at the end of decode(); removing that resampling is the point of this variant.

Freezing the encoder is deliberate. Measured drift from the pretrained init in earlier runs: original VGT-AE 2.57%, ImageNet stage1@100K 7.68%, CC12M stage1@45K 7.67%, stage2/3 8.08%. A frozen encoder (drift ≡ 0) is the upper-bound test of the hypothesis that VGT-AE's "semantically right, pixel-wise wrong" behaviour comes from staying near the VLM feature space.

Reconstruction PSNR on held-out CC12M shards, each model at its native geometry:

Important: not interchangeable with the Qwen2.5-VL latents

These have the same `(32,16,16)` shape but live in a different space. Nothing downstream can tell them apart by shape, so do not mix them into one training stream and do not decode them with the Qwen2.5-VL VGT-AE.

Preprocessing

Resize(512, BICUBIC) → CenterCrop(512) → ToTensor() → Normalize((0.5,)*3, (0.5,)*3), i.e. images in [-1, 1]. No ImageNet re-normalisation — SigLIP2's resize_mode: siglip means the tower was trained on [-1, 1] directly.

Note on the per-rank meta files

_meta_rank*.json are per-run records that were rewritten by a post-reboot resume, so their sum (3,916,287) covers only the 629 shards processed after the resume. Use _meta.json, whose samples is an exact count of .npy members across all 1097 tars.

Loading

import io, numpy as np, webdataset as wds

ds = (wds.WebDataset("00000.tar")
        .to_tuple("npy", "txt")
        .map_tuple(lambda b: np.load(io.BytesIO(b)), lambda b: b.decode()))
z, caption = next(iter(ds))     # z.shape == (32, 16, 16), float16