Scaling Multi-Teacher Distillation for Digital Pathology

Abstract

Method

Vanilla multi-teacher distillation

Following previous work, a ViT student is trained to reconstruct the representations of a set of frozen teachers . For each teacher, linear layers and project the student CLS and patch tokens to the teacher embedding spaces:

where losses combine a cosine and a smooth L1 term, as done in previous work:

A single shared CLS token must thus simultaneously compress all teachers into one representation, which becomes increasingly bottlenecked as grows. We hypothesize that this is why vanilla MTD does not benefit from adding more teachers.

Teacher-specific tokens

We therefore abandon the single-token design and instead assign one independent token per teacher. Formally, we maintain a learnable token matrix , whose tokens are prepended to the patch tokens:

Structured attention masking

MoE blocks and training objective

As a complementary mechanism, the last transformer blocks (3 for OMNI-T/S/B, 5 for OMNI-L) replace the standard dense feed-forward networks with Mixture-of-Experts (MoE) layers, with 5 experts and top-2 routing. This increases the total model capacity while keeping inference cost constant, and a Switch Transformer load-balancing loss encourages balanced expert utilization. We also add a symmetric InfoNCE loss for each teacher, penalizing embeddings that collapse to similar representations across different samples. The full training objective is

with and . At downstream time, the teacher-specific tokens are aggregated with a simple mean-pooling operator.

Results

Scaling with the number of teachers

For a fair evaluation, we rank teachers by their individual performance on the THUNDER benchmark and add them from strongest to weakest. Increasing the number of teachers for vanilla MTD does not improve downstream performance, and even degrades it, while our method better scales. Similarly, the reconstruction quality of UNI2-h drops from 0.90 to 0.83 cosine similarity for vanilla MTD, whereas our approach retains 0.87 at 10 teachers.

Linear probing benchmarking

The full OMNI family demonstrates a consistent scaling trend. OMNI-B and OMNI-L are ranked second and first respectively, while being up to 10× more efficient in GFLOPs than the best teacher models. Notably, OMNI-L achieves the best tile-level F1 (87.8%) and slide-level BAcc (63.5%), outperforming all teachers including UNI2-h (681M) and H-optimus-1 (1.1B) despite being significantly smaller. A paired Wilcoxon signed-rank test shows that OMNI-L significantly outperforms the large majority of teachers at both tile and slide level.

Linear probing performance at tile level (F1, 13 datasets) and slide level (BAcc, 7 CPTAC datasets / 25 tasks), and their mean. Mean rank is computed per dataset then averaged. Bold: best; underline: second best. TCGA-derived THUNDER datasets are removed to avoid data leakage.

Ablation studies

Ablation study on the contribution of the structured attention, InfoNCE loss and MoE blocks (mean F1 on THUNDER linear probing, ViT-Base). When removing each contribution, performance drops, showing that all of them are necessary.

Analysis of per-teacher representations

A UMAP of the 10 teacher-specific tokens from OMNI-B on 3k images shows well-separated clusters, confirming that each token specializes toward its assigned teacher rather than collapsing into a common average embedding. Moreover, linearly interpolating two teacher tokens reveals a smooth and monotonic crossover, indicating that the student model has learned a geometrically well-structured latent space in which teacher representations are continuously interpolable. This might explain why teacher-specific tokens can be easily aggregated at downstream time.

Released models

We release OMNI, a collection of lightweight pathology vision encoders with different sizes, distilled from 10 state-of-the-art open-source foundation models: H-optimus-1, Virchow2, UNI2-h, Prov-GigaPath, Kaiko-ViT-B/8, CONCH 1.5, Hibou-L, H0-mini, KEEP and DINOv3-ViT-L/16. All models take 224 × 224 tiles with a patch size of 14 and output 10 teacher-specific tokens, which are mean-pooled into a single representation.

Pre-training dataset: 10M tiles extracted from 10K whole-slide images sourced from TCGA. Teacher embeddings are pre-computed and stored offline. Dataset and training code.

The pretrained weights can be loaded directly from the Hugging Face Hub:

from transformers import AutoModel

model = AutoModel.from_pretrained("sofieneb/omni-base", trust_remote_code=True).eval()
embedding = model(x)                            # (B, 768): average of the 10 teacher tokens
out = model(x, return_teacher_tokens=True)      # + out["teacher_tokens"]: (B, 10, 768)

BibTeX

@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}
}