QKVAE-1m
QKVAE-1m is a 1.06 million parameter quantizing convolutional autoencoder that turns RGB images into discrete visual tokens from a 15,360-code finite-scalar-quantization (FSQ) codebook. The tokens are ordinary integer ids, so you can graft them onto a small language model's vocabulary and let a 1M language model see images without an external encoder, adapter, or vision tower.
This release is the adaptive, native-resolution version. Two things changed from the first checkpoint:
- Any resolution. The model is fully convolutional. It takes any H×W image and reconstructs it at the exact same H×W. No resize, no aspect-ratio squash. 96×96 is just the size it was trained at.
- Adaptive token count. Instead of always spending the full grid, it merges flat regions on a quadtree over the latent and emits a variable-length token stream. Flat images cost a fraction of the tokens; detailed images keep theirs.
auto_tau lets the model pick its own detail level per image. On a detailed image it holds roughly 96% structural match at about 80% of the tokens, and 97%+ if you force the full grid.
Notice
QKVAE-1m is a research prototype. Trained on STL-10 (96×96 natural images of animals, vehicles, and similar subjects). Out-of-distribution inputs like text, faces, diagrams, very large structures, and very small structures will reconstruct with the characteristic STL-10 "vibe." Use accordingly. Pair with a real language model before drawing strong conclusions.
At a Glance
| Property | Value |
|---|
| Parameters | 1,063,304 trainable (1.06M) |
| Architecture | CNN encoder + FSQ + CNN decoder |
| Image Size | Any H × W (native; trained at 96 × 96) |
| Tokens per Image | Adaptive; up to 2,304 on a 96 × 96 grid |
| Codebook Size | 15,360 (8 × 8 × 8 × 6 × 5 levels) |
| Latent Dim | 5 |
| Compression | 6.9× (27,648 → ~4,001 bytes) |
| Training Dataset | STL-10 train (5,000 images, 96×96) |
| Test Dataset | STL-10 test (8,000 images, 96×96) |
| PSNR (test) | 32.61 dB |
| SSIM (test) | 0.9743 |
| Training Steps | 6,000 |
| Training Time | ~1.3 minutes on a single RTX 5090 |
| Precision | bfloat16 autocast, fp32 weights |
| Output Range | tanh → clamped to [−1, 1] |
Background
The QKVAE was built to answer one question: how small can a quantizing autoencoder get before reconstruction quality collapses, and is the result still useful as a discrete image tokenizer for a tiny language model? For 96×96 RGB on STL-10 the answer is 1.06 million parameters at 32.61 dB PSNR, and the 2,304 code tokens slot into a small language model's vocabulary as ordinary integer ids, with no vision tower required.
The training loop is deliberately ordinary. AdamW with OneCycleLR. bfloat16 autocast. L1 + 0.25×MSE + 0.5×Sobel-edge loss. 6,000 steps at batch 32 on the STL-10 train split, evaluated on the held-out test split. No EMA, no GAN, no perceptual loss, no discriminator, no augmentations beyond the natural STL-10 diversity. The only thing that makes this work is the FSQ codebook being large enough (15,360) to cover the STL-10 manifold and small enough that the gradient is informative end-to-end.
The 1m here is the first checkpoint. We are actively working to make this way way way more efficient. Smaller architecture, fewer codebook entries, less compute per forward, and lower token counts per image. The 1m is published because it works today, and because every smaller version we try will be measured against it.
Model Specification
| Parameter | Value |
|---|
| Architecture | Quantizing Convolutional Autoencoder |
| Stem | 3 → 104 channels, 3×3 conv, padding 1 |
| Encoder Downsample | 1× 4×4 stride-2 conv (96 → 48) |
| Encoder ResBlocks | 2 × (GroupNorm → SiLU → 3×3 conv → GroupNorm → SiLU → 3×3 conv + residual), 104 channels |
| Encoder Head | GroupNorm → SiLU → 1×1 conv to 5 latent channels |
| Quantizer | Finite Scalar Quantization, levels (8, 8, 8, 6, 5) |
| Codebook Size | 15,360 |
| Decoder Input | 3×3 conv 5 → 104 |
| Decoder ResBlocks | 2 × ResBlock(104) |
| Decoder Upsample | 1× nearest-neighbour 2× → 3×3 conv 104 → 104 → SiLU |
| Decoder Head | GroupNorm → SiLU → 3×3 conv 104 → 3, tanh |
| Normalization | GroupNorm (8 groups, 104 channels) |
| Activation | SiLU |
| Output Range | tanh, no clamping outside the model |
Training Details
| Parameter | Value |
|---|
| Dataset | STL-10 train, 96×96, 5,000 images |
| Test Set | STL-10 test, 96×96, 8,000 images |
| Augmentation | None (dataset is naturally diverse at 96×96) |
| Steps | 6,000 |
| Batch Size | 32 |
| Optimizer | AdamW, β=(0.9, 0.95), weight decay 1e-4 |
| Learning Rate | 3e-4 peak, OneCycleLR, 5% warmup, cosine decay |
| Loss | L1 + 0.25 × MSE + 0.5 × L1(∇Sobel(recon), ∇Sobel(target)) |
| Precision | bfloat16 autocast, fp32 reductions |
| Gradient Clip | 1.0 |
| Hardware | NVIDIA RTX 5090, 32 GB |
| Wall Time | 1.3 minutes (78 seconds) |
| Throughput | ~77 steps/second |
Usage
Python API
1import torch
2from modeling import load_qkvae, reconstruction_psnr
3from inference import load_image
4from png_io import save_png
5
6model = load_qkvae("model.safetensors") # 1.06M params, ~4 MB
7
8image = load_image("examples/cityscape.png") # (3, H, W) in [-1, 1], native size
9
10with torch.no_grad():
11 tau = model.auto_tau(image) # model picks its own detail level
12 stream = model.encode_adaptive(image, tau) # variable-length list[int] token stream
13 recon = model.decode_adaptive(stream) # (3, H, W), same size as the input
14
15save_png(recon, "out.png")
16print(f"tokens: {len(stream) - 2} PSNR: {reconstruction_psnr(image, recon):.2f} dB")
Fixed full grid (no merging, max quality) is the plain encode/decode path:
1recon = model.reconstruct(image) # (3, H, W), full grid, same size in/out
2indices = model.encode(image.unsqueeze(0)) # (1, N) int64 ids in 0..15359 (N = grid)
Command-Line Interface
1# tau defaults to "auto"; pass 0 for the full grid, or any float to merge harder
2python inference.py path/to/image.png model.safetensors out.png auto
3# image=image.png size=1570x198 tokens=66253/77715 (85%) tau=0.812 ... psnr=33.32dB -> out.png
Using the Tokens with a Language Model
The 2,304 code tokens are just integers in 0..15,359. To make a language model see images, append the 15,360 code ids plus 3 special ids (<img_start>, <img_end>, <img_newrow>) to its vocabulary, project each code's dequantized FSQ vector into a learnable embedding, and add a small amount of training data mixing image and text tokens. The result is a single model that handles both modalities. No vision tower, no cross-attention.
Adaptive Tokenization
The full grid is one token per latent cell (2,304 at 96×96). Most images do not need all of them: sky, walls, and out-of-focus background are nearly flat in latent space. encode_adaptive walks a quadtree over each 16×16 tile of the latent grid and, whenever a sub-block's latent spread is within tau, replaces it with a single shared code marked by a per-depth merge id (codebook_size + depth, kept distinct from any real FSQ id). The result is a variable-length stream [H, W, *symbols] that never exceeds the full grid and shrinks on flat content.
tau is the quality/size knob: 0 forbids merges and keeps every token, higher values merge harder. auto_tau bisects on latent distortion to pick the most aggressive tau that still stays within a small fixed error of the full-grid latent, so flat images merge hard and detailed images keep their tokens at a steady quality, with no decoder passes in the search.
1stream = model.encode_adaptive(image, tau=model.auto_tau(image))
2recon = model.decode_adaptive(stream) # exact, same H×W as the input
3grid, (h, w), (H, W) = model.adaptive_to_grid(stream) # replay back to a flat index grid
Repository Contents
| File | Description |
|---|
README.md | This card |
config.json | Architecture spec and reported metrics |
model.safetensors | 1.06M trainable parameters, fp32, ~4.1 MB |
modeling.py | Self-contained QKVAE, FSQ, adaptive encode/decode, auto_tau, load_qkvae |
inference.py | CLI + helpers for native-resolution reconstruction |
png_io.py | stdlib PNG reader/writer (no Pillow or torchvision) |
requirements.txt | torch, safetensors. Nothing else. |
Limitations
- Trained at one scale. It reconstructs any resolution at native size, but it learned its texture prior from 96×96 STL-10. Very large images are tokenized in that same per-cell vocabulary, so fine detail far from the training scale still looks STL-10.
- Fixed vocabulary. 15,360 FSQ codes cover STL-10 at this scale. They are not enough for ImageNet, faces, or text legibly.
- Distribution. Trained on STL-10 train (5,000 images). Faces, text, architecture, and diagrams will come out looking like STL-10 subjects. This is the expected behaviour of a 1M autoencoder.
- No perceptual loss. L1 + MSE + edge are honest, blunt metrics. Larger models trained with LPIPS or DISTS produce more pleasing results. We skipped those to keep a clean ablation of scale.
- No FID / IS. Evaluation is PSNR + SSIM + visual inspection. The grid above is the only qualitative evidence.
- Not for production. Research artifact.
Citation
1@misc{qkvae1m2026,
2 author = {Glint Research},
3 title = {QKVAE-1m: A 1.06M-Parameter Quantizing Autoencoder for 96x96 Image Tokens},
4 year = {2026},
5 publisher = {Hugging Face},
6 url = {https://huggingface.co/Glint-Research/QKVAE-1m}
7}
Built by Glint Research. Small models trying their best since 2026. 
