Aditya Sridhar

Model · Oct 2026 · runs in your browser

ViT-Tiny

A transformer that reads an image the way Mini-GPT reads text: cut into 64 patches of 4×4 pixels, each a token, with attention between all of them from the first layer on. No convolution after the patch embedding, so nothing tells it that neighbouring pixels belong together; it has to learn that from 50,000 images.

Test accuracy
82.68%
Parameters
1,806,538
Model file
7.0 MB, ONNX
Training time
12.1 min, GPU

Try it

CIFAR-10 test images

Pick an image
32×32 pixels: the whole image, exactly as the model sees it.

Loading the model…

  1. airplane
  2. automobile
  3. bird
  4. cat
  5. deer
  6. dog
  7. frog
  8. horse
  9. ship
  10. truck

Runs in your browser. Nothing is uploaded.

The heatmap is attention rollout (Abnar & Zuidema, 2020): how much each patch reached the class token through all 6 layers. It is the same for every class, unlike the activation maps on the ResNet and CNN pages, which use the same test images.

Against the convolutional networks

MetricValueSource
ViT-Tiny, 1807k params, 50 epochs82.68%training/vit_cifar10/log.csv, epoch 50
ResNet-56, 853k params, 50 epochs91.70%training/resnet_cifar10/resnet-56.csv
ResNet-20, 270k params, 50 epochs90.12%training/resnet_cifar10/resnet-20.csv
CNN, 289k params, 40 epochs89.69%training/cnn_cifar10/log.csv
ViT-Tiny, exported ONNX file82.71%public/models/vit-cifar10/meta.json

With 6.7× the parameters of ResNet-20, the transformer ends 7.4 points behind it. A convolution builds in two assumptions, that nearby pixels matter together and that a pattern means the same thing anywhere in the image. The transformer has to learn both from data, and 50,000 images of 32×32 are not much data for that. With large pretraining sets the picture reverses: that is the regime in which Dosovitskiy et al. (2021) showed ViTs overtaking ResNets.

Train accuracy at the end was 98.59% against 82.68% on the test set, with the same crop and flip augmentation as the ResNets.

Training run

Accuracytrain (augmented)test
40%60%80%100%191725334150
Losstraintest
0.0000.5001.0001.500191725334150

How it works

Patches as tokens. A 4×4 convolution with stride 4 turns the image into 64 vectors of width 192. A learned class token is prepended and a learned position embedding added, since attention by itself does not know where a patch came from.

Encoder. 6 pre-norm transformer blocks, each multi-head self-attention with 3 heads and an MLP of width 384. The class token's final state, after layer norm, goes through one linear layer to the ten classes.

Training. AdamW, weight decay 0.05, a one-cycle learning rate peaking at 0.001, batch 128, 50 epochs, the same random crops and flips as the ResNet runs, bf16 autocast. No pretraining, no mixup or RandAugment.

Reproduce

Hardware
NVIDIA GeForce RTX 3050 Laptop GPU, PyTorch 2.6.0+cu124
Commit
267a304
Data
CIFAR-10, 50,000 train / 10,000 test

Environment

pip install torch numpy onnx onnxruntime

Train, export and score the ONNX file

python training/vit_cifar10/train.py