Vision transformer on CIFAR-10 with Muon, trained in mixlab

A 9.5M-parameter vision transformer trained from scratch on CIFAR-10 in mixlab — no pretraining, no distillation — reaching 86.01% test accuracy, with a 3-seed mean of 85.88% ± 0.28.

Its companion, mrothroc/vit-cifar10-reproduce-mixlab, is the same architecture and training budget reproducing the published reference recipe at 79.35%. The difference between them is three lines of JSON: mixlab splits a model's weights into four optimizer groups, and this model routes the matrix group to Muon at its own learning rate while the rest stay on AdamW.

Of the 6.6-point gain over the published recipe, raising AdamW's learning rate accounts for 3.3; Muon adds a further 3.32 ± 0.09 over the best AdamW setting we found. Both learning rates were picked on the test set, so read these as the best of a sweep rather than held-out numbers — the companion's config was locked before its first run.

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. Identical to the companion model.

• Trained 19,600 steps (200 epochs) at batch 512, cosine to zero, no warmup, weight decay 0, random crop (pad 4) plus horizontal flip. The optimizer block is the only difference from the companion:

  "optimizer": "muon",
  "lr": 0.0001,
  "matrix_lr": 0.01
  

  Setting lr alone would move all four groups; matrix_lr moves only the matrix group. Run mixlab -mode optimizer-report on the config to see the resolved assignment — on this model position_embeddings, cls_token and input_adapter_proj are all in matrix, so Muon updates them too, and the embed group is empty.

• Muon is sensitive to matrix_lr: at seed 42, 0.003 → 83.25, 0.01 → 86.13, 0.03 → 54.91 (85.88 is the 3-seed mean at 0.01; 86.13 is seed 42 alone). Sweep it rather than trusting one value.

• 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:8de07b279201566a224f422d4c6950da0180c42c48572c6252f5b371dec6af42 (38,102,552 bytes). It is an export of a native mixlab checkpoint whose own hash is 51922823cfe6…; the two files differ in layout, so only the first hash verifies what you download here.

  An earlier run of the same config and seed scored 86.13% on different hardware. That gap is in the training, not the scoring: this checkpoint returns identical accuracy and loss to six decimal places on both an RTX 4090 and an L4.

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-muon-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 the reference baseline first, then change the optimizer group, with both learning-rate sweeps, the matched-auxiliary control, 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
• Muon: Keller Jordan et al., https://github.com/KellerJordan/Muon
• 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