Vision transformer on CIFAR-10, reproduced in mixlab

A 9.5M-parameter vision transformer trained from scratch on CIFAR-10 in mixlab — no pretraining, no distillation. It reaches 79.35% test accuracy, reproducing the reference recipe it was matched against (kentaroy47/vision-transformers-cifar10 @ 79fa30c). Our own three-seed run of that reference, with RandAugment off, averages 79.16% ± 0.25; upstream's checked-in log for the same architecture ends at 79.08%.

This is the baseline for an optimizer comparison: its companion, mrothroc/vit-cifar10-muon-mixlab, is the same architecture and training budget with the matrix parameter group routed to Muon, and reaches 86.01%. The point is how little changes between them: three lines in a JSON config.

For calibration: a from-scratch ViT at this size and budget lands near 79%, and a convolutional model of the same size does better. Transformers have no convolutional inductive bias to lean on, and 50,000 images is not much to learn it from.

What it is

• 9,523,722 parameters: 6 plain bidirectional self-attention blocks with GELU FFNs, model_dim 512, 8 heads. Input is a 32×32×3 image as 64 patches of 4×4, projected by mixlab's linear_patches adapter, with learned absolute positions and CLS pooling.
• Trained 19,600 steps (200 epochs) at batch 512 on the CIFAR-10 train split, AdamW at lr 1e-4 on every parameter group, cosine to zero, no warmup, weight decay 0. Augmentation is random crop (pad 4) plus horizontal flip. Initialization matches nn.Linear via weight_init: "pytorch_linear_all", with pos_embedding and cls_token at std 1.0.
• Reported accuracy is the final epoch on the 10,000-image test split — no checkpoint selection. The published checkpoint is seed 42. model.safetensors in this repo is sha256:143e7a139dd59b500042035d744c7ded35882bdc168d57ea9a94a90efc921c1a (38,102,552 bytes). It is an export of a native mixlab checkpoint whose own hash is 151eea10a3ab…; the two files differ in layout, so only the first hash verifies what you download here.

Load it

This is a custom-code export, so load with trust_remote_code=True. Inputs are patch sequences, not raw images: [batch, 64, 48], per-channel normalized with the reference constants.

import torch
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "mrothroc/vit-cifar10-reproduce-mixlab", trust_remote_code=True).eval()

# patches: [batch, 64, 48] float32, per-channel normalized with the reference constants
logits = model(input_values=torch.randn(1, 64, 48)).logits   # -> [1, 10]
label  = logits.argmax(-1)   # 0..9, CIFAR-10 class order

scripts/prep_cifar.py in the cookbook entry turns CIFAR-10 into exactly this layout.

Reproduce / learn more

Full recipe (reproduce this baseline, then improve it by changing the optimizer group, with both learning-rate sweeps and the failures): mixlab cookbook, vit-cifar10-muon.

Requires mixlab v0.117.0 or newer.

Citation

Please cite what this builds on:

• The reference implementation: https://github.com/kentaroy47/vision-transformers-cifar10 @ 79fa30c
• The ViT implementation it wraps: Phil Wang, vit-pytorch, https://github.com/lucidrains/vit-pytorch
• Vision Transformer: Dosovitskiy, A. et al. "An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR 2021. arXiv:2010.11929
• CIFAR-10: Krizhevsky, A. "Learning Multiple Layers of Features from Tiny Images." Technical report, University of Toronto, 2009.
• A paper that trains this repo's ViT on CIFAR-10: Zhu, H.; Chen, B.; Yang, C. "Understanding Why ViT Trains Badly on Small Datasets: An Intuitive Perspective." 2023. arXiv:2302.03751