sofieneb/omni-large

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

Model Summary

OMNI-L is a 330M-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: 67.45 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:

Model	Params	Dim	GFLOPs	Tile LP F1 (%)	Slide LP BAcc (%)
OMNI-T	7.3M	192	1.39	83.9	63.2
OMNI-S	28.7M	384	5.48	85.3	63.1
OMNI-B	114M	768	21.8	87.1	63.2
OMNI-L (this model)	330M	1024	67.45	87.8	63.5

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

---

Model Details

Architecture	ViT-Large/14 with a MoE tail
Embedding dim	1024
Depth / heads	24 blocks / 16 heads
MoE	last 5 blocks, 5 experts, top-2 routing
Teacher tokens	10 (one per teacher)
Input	224 × 224 RGB, ImageNet normalization
Patch tokens	16 × 16 = 256

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:

Tile LP F1 (%)	Slide LP BAcc (%)	kNN F1 (%)	16-shot F1 (%)
87.8	63.5	84.8	80.7

---

Usage

import torch
from PIL import Image
from transformers import AutoModel

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

# Load OMNI-L
model = AutoModel.from_pretrained("sofieneb/omni-large", 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, 1024)

Teacher tokens and patch tokens

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

out["embedding"]       # (1, 1024)       average of the teacher tokens (same as model(x))
out["teacher_tokens"]  # (1, 10, 1024)   one token per teacher
out["patch_tokens"]    # (1, 256, 1024)  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, 1024)

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-large 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-large/<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-large --local-dir ./omni-large

Then in your offline job/script:

from transformers import AutoModel
model = AutoModel.from_pretrained(
    "./omni-large",        # 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}
  }