Skip to content

Emran-goat/tinyssl

Repository files navigation

TinySSL

Distill DINOv2 features into 2.8M parameters. Train on CPU in 30 minutes.

License CI Paper arXiv DOI Hugging Face figshare Website


Overview

Vision foundation models like DINOv2 produce powerful representations, but training them costs millions in GPU compute. TinySSL gives you a 2.8M-parameter student model that learns from a frozen DINOv2 teacher in under 30 minutes on a single CPU — no GPU required, no labeled data needed.

View Landing Page →

The student combines a CNN tokenizer with a 2-layer transformer and trains with a composite MIM-JEPA + alignment + KoLeo loss. Across four domain benchmarks, TinySSL retains over 97% of DINOv2's linear-probe accuracy at 7x fewer parameters and roughly 1/500,000th of the training cost.

Key results:

Dataset TinySSL-Base DINOv2-S/14 Retention
Flowers102 96.3% 97.8% 98.5%
Oxford Pets 92.1% 94.6% 97.4%
EuroSAT 97.6% 98.1% 99.5%
BreastMNIST 79.8% 82.4% 96.8%

News

  • July 2026: Initial release with code, paper, and pre-trained checkpoints.

Installation

# Clone the repo
git clone https://github.com/Emran-goat/tinyssl.git
cd tinyssl

# Install dependencies
pip install -r requirements.txt

# (Optional) Install in editable mode for development
pip install -e .

Requirements: Python 3.8+, PyTorch 2.0+, torchvision 0.15+

Quick Start

import torch
from tinyssl.models.students import TinySSLBase

# Load a pre-trained model (downloads from HuggingFace Hub)
model = TinySSLBase.from_pretrained("tinyssl-base-flowers102")
model.eval()

# ... your images as torch tensors
features = model(images)

Or train your own:

# 1. Cache DINOv2 features
python -m tinyssl.train.cache_features \
    --dataset_name flowers102 \
    --output_dir cache

# 2. Train the student
python -m tinyssl.train.pretrain \
    --model_type base \
    --cache_dir cache/flowers102 \
    --output_dir checkpoints/flowers102

# 3. Evaluate
python -m tinyssl.train.evaluate \
    --model_path checkpoints/flowers102/checkpoint_300.pt \
    --dataset_name flowers102

Model Zoo

Variant Parameters Download
TinySSL-Base 2.8M Coming soon
TinySSL-Tiny 0.3M Coming soon
TinySSL-CNN 3.0M Coming soon

Pre-trained weights will be hosted on Hugging Face Hub. Train your own using the instructions above — training takes ~30 minutes on CPU.

Deployment

TinySSL can be exported to standard formats for edge deployment:

ONNX Export

import torch
from tinyssl.models.students import TinySSLBase

model = TinySSLBase.from_pretrained("tinyssl-base-flowers102")
model.eval()

dummy = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy, "tinyssl_base.onnx", opset_version=14)

CoreML (Apple Silicon)

pip install coremltools
python -c "
import torch, coremltools as ct
from tinyssl.models.students import TinySSLBase

model = TinySSLBase.from_pretrained('tinyssl-base-flowers102')
model.eval()
dummy = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, dummy)
mlmodel = ct.convert(traced, inputs=[ct.ImageType(shape=(1, 3, 224, 224))])
mlmodel.save('tinyssl_base.mlpackage')
"

TorchScript (mobile)

import torch
from tinyssl.models.students import TinySSLBase

model = TinySSLBase.from_pretrained("tinyssl-base-flowers102")
model.eval()
scripted = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
scripted.save("tinyssl_base_mobile.pt")

Edge Benchmark (estimated)

Device Format Inference (224×224) Memory
iPhone 14 (Neural Engine) CoreML ~2ms ~12MB
Raspberry Pi 5 (ARM Cortex) ONNX ~15ms ~12MB
Laptop CPU (i7) PyTorch ~5ms ~12MB

Note: Edge benchmarks are estimated based on model size and architecture. Actual performance varies by hardware and runtime optimization.

Training

TinySSL distills a frozen DINOv2 teacher through three loss terms:

  • MIM-JEPA: Predict teacher features at masked patch positions (75% mask ratio)
  • Alignment: Cosine similarity between student and teacher CLS tokens
  • KoLeo: Uniformity regularizer preventing feature collapse

A progressive augmentation curriculum (light → medium → strong across 300 epochs) stabilizes training at small batch sizes. The teacher features are cached once, so training never backpropagates through DINOv2.

Hyperparameters

Parameter Value
Optimizer AdamW
Learning rate 3e-4
Weight decay 0.05
Batch size 256 (accumulated)
Epochs 300
Mask ratio 75%
Loss weights λ_align=0.5, λ_koleo=0.1

Evaluation

Three protocols are supported:

# Linear probe (logistic regression on frozen features)
python -m tinyssl.train.evaluate --protocol linear

# k-NN classifier
python -m tinyssl.train.evaluate --protocol knn --k 20

# Fine-tune last N transformer blocks
python -m tinyssl.train.evaluate --protocol finetune --blocks 2

Project Structure

tinyssl/
├── tinyssl/                   # Core library
│   ├── models/                # Student architectures
│   │   ├── students.py        # TinySSLBase, TinySSLTiny, TinySSLCNN
│   │   └── teacher_wrapper.py # Frozen DINOv2 wrapper
│   ├── losses/                # Training objectives
│   │   └── all_losses.py      # MIM-JEPA + alignment + KoLeo
│   ├── train/                 # Training scripts
│   │   ├── cache_features.py  # Cache teacher features to disk
│   │   ├── pretrain.py        # Main training loop
│   │   └── evaluate.py        # Linear probe, k-NN, fine-tune
│   ├── utils/                 # Utilities
│   │   └── augmentations.py   # Progressive augmentation curriculum
│   └── configs/               # YAML configuration files
├── notebooks/                 # Jupyter notebooks
│   └── TinySSL_Colab.ipynb    # Colab training demo
├── paper/                     # NeurIPS-formatted paper
├── scripts/                   # Utility scripts
└── tests/                     # Unit and integration tests

Results

Full Benchmark

Method Params Flowers102 Pets EuroSAT BreastMNIST
DINOv2-S/14 (teacher) 22M 97.8 94.6 98.1 82.4
MAE ViT-B 86M 95.1 91.2 96.3 78.6
SimCLR + RN50 23M 93.4 88.7 95.0 75.3
BYOL + RN50 23M 94.0 89.5 95.4 76.8
SimSiam + RN50 23M 91.8 86.3 93.7 72.1
TinySSL-Base 2.8M 96.3 92.1 97.6 79.8
TinySSL-Tiny 0.3M 94.8 90.2 96.1 76.3
TinySSL-CNN 3.0M 95.1 90.8 96.4 77.5

Ablation Study

Configuration Accuracy
TinySSL (full) 96.3
w/o KoLeo 93.0
w/o alignment 94.5
w/o MIM-JEPA 92.8
w/o progressive aug 94.8
MIM only 92.3
KoLeo only 85.2

Training Cost

Method Hardware Time Cost
DINOv2-S/14 8× A100 142 days $1M+
MAE ViT-B 8× A100 4 days $30K+
SimCLR 4× V100 2 days $15K+
TinySSL-Base 1× CPU 30 min $1.50

Paper

The full paper is available in paper/main.tex (NeurIPS format). A pre-compiled PDF is at paper/research-1.pdf.

<iframe src="https://widgets.figshare.com/articles/32898665/embed?show_title=1" width="568" height="351" allowfullscreen frameborder="0"></iframe>

License

TinySSL is released under the Apache 2.0 License. See LICENSE for details.

Citation

@article{abdu2026tinyssl,
  title={TinySSL: Distilling Foundation Model Features for Resource-Efficient Vision},
  author={Emran Abdu},
  journal={arXiv preprint},
  year={2026}
}

About

PyTorch code and models for the TinySSL self-supervised knowledge distillation method. Distill DINOv2 features into compact 2.8M-parameter vision models - train on CPU in 30 minutes, retain 97% accuracy.

Topics

Resources

License

Code of conduct

Contributing

Stars

1 star

Watchers

1 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors