---
license: cc-by-nc-4.0
library_name: transformers
pipeline_tag: image-feature-extraction
tags:
  - histopathology
  - digital-pathology
  - foundation-model
  - vision-transformer
  - mixture-of-experts
  - knowledge-distillation
  - custom_code
---

# [NeurIPS 2026] OMNI-T: Scaling Multi-Teacher Distillation for Digital Pathology

## Model Summary

**OMNI-T** is a 7.3M-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**: 1.39 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 (this model)** | 7.3M | 192 | 1.39 | 83.9 | 63.2 |
| [OMNI-S](https://huggingface.co/sofieneb/omni-small) | 28.7M | 384 | 5.48 | 85.3 | 63.1 |
| [OMNI-B](https://huggingface.co/sofieneb/omni-base) | 114M | 768 | 21.8 | 87.1 | 63.2 |
| [OMNI-L](https://huggingface.co/sofieneb/omni-large) | 330M | 1024 | 67.45 | 87.8 | 63.5 |

 🌐 **Project page**: [*Scaling Multi-Teacher Distillation for Digital Pathology*](https://sofiene-boutaj.github.io/omni/) (NeurIPS 2026 Main Track)

---

## Model Details

| | |
| --- | --- |
| Architecture | ViT-Tiny/14 with a MoE tail |
| Embedding dim | 192 |
| Depth / heads | 12 blocks / 3 heads |
| MoE | last 3 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 (%) |
| -------------- | ----------------- | ---------- | -------------- |
| 83.9 | 63.2 | 81.5 | 77.5 |

---

## Usage

```python
import torch
from PIL import Image
from transformers import AutoModel

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

# Load OMNI-T
model = AutoModel.from_pretrained("sofieneb/omni-tiny", 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, 192)
```

### Teacher tokens and patch tokens

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

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

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

```python
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](https://github.com/MICS-Lab/thunder).

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

```bash
pip install thunder-bench
export THUNDER_BASE_DATA_FOLDER=/path/to/thunder_data
hf download sofieneb/omni-tiny 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):

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

For example, on MHIST:

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

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

```bash
thunder results-summary
```
---

## Offline usage (HPC clusters without internet)

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

```bash
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
```

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

```bash
hf download sofieneb/omni-tiny --local-dir ./omni-tiny
```

Then in your offline job/script:

```python
from transformers import AutoModel
model = AutoModel.from_pretrained(
    "./omni-tiny",        # 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:

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