OMNI-B: Scaling Multi-Teacher Distillation for Digital Pathology

Model Summary

OMNI-B is a 114M-parameter pathology vision encoder distilled from 10 state-of-the-art pathology foundation models at once. Instead of compressing all teachers into a single CLS token (vanilla multi-teacher distillation), OMNI learns one token per teacher, together with a structured attention mask, Mixture-of-Experts (MoE) layers and a per-teacher contrastive loss. This lets distillation keep improving as more teachers are added.

By default the model returns the average of its 10 teacher-specific tokens, which is the embedding used for all results below. The individual teacher tokens and the patch tokens can also be returned with a flag.

• ⚡ Efficient: 21.8 GFLOPs per 224×224 tile
• 🧠 Distilled from 10 teachers: H-optimus-1, Virchow2, UNI2-h, Prov-GigaPath, Kaiko-ViT-B/8, CONCH v1.5, Hibou-L, H0-mini, KEEP, DINOv3-ViT-L/16
• 📚 Trained on 10M tiles from 10k TCGA whole-slide images

Other sizes of the collection:

📄 Paper: Scaling Multi-Teacher Distillation for Digital Pathology (NeurIPS 2026 Main Track)

Model Details

Results

Linear probing on THUNDER (tile level, 13 datasets, TCGA datasets removed) and PathoBench (slide level, 25 CPTAC tasks, mean pooling), plus kNN and 16-shot classification on THUNDER:

Usage

import torch
from PIL import Image
from transformers import AutoModel

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load OMNI-B
model = AutoModel.from_pretrained("sofieneb/omni-base", trust_remote_code=True).to(device).eval()
transform = model.get_transform()  # resize to 224 + ImageNet normalization

image = Image.open("tile.png").convert("RGB")
x = transform(image).unsqueeze(0).to(device)  # (1, 3, 224, 224)

with torch.inference_mode():
    embedding = model(x)  # average of the 10 teacher tokens

print(embedding.shape)  # (1, 768)

Teacher tokens and patch tokens

with torch.inference_mode():
    out = model(x, return_teacher_tokens=True, return_patch_tokens=True)

out["embedding"]       # (1, 768)       average of the teacher tokens (same as model(x))
out["teacher_tokens"]  # (1, 10, 768)   one token per teacher
out["patch_tokens"]    # (1, 256, 768)  16 x 16 patch tokens

# Teacher order of the tokens
print(model.teachers)

# Example: keep the token of a single teacher
idx = 0  # position of the teacher you want in model.teachers
single_teacher = out["teacher_tokens"][:, idx]  # (1, 768)

Teacher tokens are well separated and complementary, so the simple average works well out of the box. You can also learn a task-specific weighting of the 10 tokens (see Table 2 of the paper).

Mixed precision

with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    embedding = model(x)

Quick evaluation with THUNDER

eval_thunder.py wraps the model for the THUNDER benchmark.

Install THUNDER and set its data folder, where datasets, embeddings and results are stored:

pip install thunder-bench
export THUNDER_BASE_DATA_FOLDER=/path/to/thunder_data
hf download sofieneb/omni-base eval_thunder.py --local-dir .

Then run the benchmark, choosing the datasets with --datasets. The datasets are downloaded and their splits generated automatically before benchmarking (add --no-download to skip this once they are in place):

python eval_thunder.py --datasets <dataset> --tasks knn linear_probing

For example, on MHIST:

python eval_thunder.py --datasets mhist --tasks knn linear_probing

Results are written to $THUNDER_BASE_DATA_FOLDER/outputs/res/<dataset>/omni-base/<task>/. To collect them in one CSV:

thunder results-summary

Offline usage (HPC clusters without internet)

If compute nodes don't have internet, always run jobs with the offline flags:

export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1

Prepare the model in advance on a front-end/login node (with internet):

hf download sofieneb/omni-base --local-dir ./omni-base

Then in your offline job/script:

from transformers import AutoModel
model = AutoModel.from_pretrained(
    "./omni-base",        # local path instead of hub ID
    trust_remote_code=True,
    local_files_only=True,  # uses local files only
)

Citation

If our work contributes to your research, please consider citing our paper:

@inproceedings{boutaj2026omni, author = {Boutaj, Sofi{`e}ne and Marza, Pierre and Belagali, Varun and Samaras, Dimitris and Vakalopoulou, Maria and Christodoulidis, Stergios}, title = {Scaling Multi-Teacher Distillation for Digital Pathology}, booktitle = {Advances in Neural Information Processing Systems (NeurIPS)}, year = {2026} }