Update README.md

e560b96 verified 13 days ago

6.21 kB

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

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