CustomViT-Nano: 4.24M Parameter Compact Vision Transformer
CustomViT-Nano is a compact, modernized Vision Transformer architecture designed for efficient ImageNet-1K image classification under a small parameter budget. With only 4.24M trainable parameters, the model combines a lightweight convolutional stem with modern Transformer components including 2D Rotary Positional Embeddings, Pre-RMSNorm, SwiGLU feed-forward layers, and PyTorch SDPA attention. The model is designed to deliver strong classification performance while remaining significantly smaller than standard large Vision Transformer baselines.
Key Architectural Features
CustomViT-Nano modernizes a small Vision Transformer design using several efficiency-focused architectural components.
| Component | Design in CustomViT-Nano |
|---|---|
| Patch Embedding | Multi-stage convolutional stem instead of single large patchify projection |
| Stem Activation | GELU |
| Token Layout | 14 × 14 patch tokens + CLS token |
| Normalization | Pre-RMSNorm inside Transformer blocks |
| Attention | Multi-head attention using PyTorch scaled dot-product attention |
| Position Encoding | 2D Rotary Positional Embeddings for image patch grids |
| MLP Block | SwiGLU gated feed-forward network |
| Classifier | CLS-token classification head |
ConvStem Design
Instead of directly projecting 16 × 16 image patches with one large-stride convolution, CustomViT-Nano uses a progressive convolutional stem:
224 × 224 × 3
↓
112 × 112 × 32
↓
56 × 56 × 64
↓
28 × 28 × 128
↓
14 × 14 × 224
This gives the model a stronger local visual inductive bias before global Transformer reasoning.
Benchmark & Evaluation
- Evaluation Dataset: ImageNet-1K validation set
- Total Evaluation Images: 50,000
- Input Resolution: 224 × 224
- Number of Classes: 1000
| Model | Parameters | Top-1 Accuracy | Top-5 Accuracy |
|---|---|---|---|
| CustomViT-Nano | 4.24M | 63.60% | 84.93% |
| Google ViT-B/16 | 86.6M | 80.31% | 95.49% |
Parameter Efficiency Comparison
CustomViT-Nano is approximately 20.4× smaller than the reference Google ViT-B/16 model.
Google ViT-B/16: 86.6M parameters
CustomViT-Nano: 4.24M parameters
Parameter reduction:
86.6M / 4.24M ≈ 20.4× smaller
Despite using only around 4.9% of the parameters of the 86.6M ViT baseline, CustomViT-Nano achieves:
- 63.60% Top-1 Accuracy
- 84.93% Top-5 Accuracy
on the ImageNet-1K validation set.
Target Use Cases & Applications
CustomViT-Nano is suitable for scenarios where a compact visual classifier is preferred over a large transformer model.
-
Compact Image Classification
Lightweight ImageNet-style classification with a transformer-based architecture. -
Edge & Resource-Constrained Vision
Useful for environments where model size and memory footprint are important constraints. -
Educational Vision Transformer Research
A compact architecture for studying ConvStem patch embeddings, 2D RoPE, SDPA attention, RMSNorm, and SwiGLU inside a full ImageNet-scale classifier. -
Backbone Experiments
Can be used as a small image encoder backbone for downstream classification or transfer-learning experiments. -
Efficient Model Baselines
Useful as a compact baseline for experiments involving distillation, pruning, quantization, or architecture search.
How to Use
Fast Inference with Hugging Face pipeline
from transformers import pipeline
classifier = pipeline(
"image-classification",
model="kd13/vit-nano-patch16-224",
trust_remote_code=True
)
results = classifier(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/hub/parrots.png"
)
for pred in results:
print(f"Label: {pred['label']} | Score: {pred['score']:.4f}")
Inference with PIL Image
from PIL import Image
from transformers import pipeline
classifier = pipeline(
"image-classification",
model="kd13/vit-nano-patch16-224",
trust_remote_code=True
)
image = Image.open("image.jpg").convert("RGB")
results = classifier(image)
for pred in results:
print(f"Label: {pred['label']} | Score: {pred['score']:.4f}")
Limitations
- The model is smaller than standard ViT-B models and therefore has lower absolute ImageNet accuracy.
- It is optimized for image classification, not detection, segmentation, captioning, or multimodal tasks.
- Performance may vary on images that differ significantly from ImageNet-style natural images.
- For maximum accuracy, larger models or teacher-distilled variants may perform better.
Disclaimer
This model is intended for research, experimentation, and efficient image-classification use cases. It should be evaluated carefully before use in production or safety-critical applications.