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.

---

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

---

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.

1. Compact Image Classification
   Lightweight ImageNet-style classification with a transformer-based architecture.

2. Edge & Resource-Constrained Vision
   Useful for environments where model size and memory footprint are important constraints.

3. 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.

4. Backbone Experiments
   Can be used as a small image encoder backbone for downstream classification or transfer-learning experiments.

5. 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.