XMag
XMag is a pathology image encoder distilled from a high-magnification foundation model to a low-magnification student model.
- Teacher: frozen UNIv2
- Student: DINOv2 ViT-B/14
- Input: RGB pathology patch,
224 x 224@5x or 2um mpp - Output: CLS embedding and patch embeddings
- Training objective: global and local cosine feature distillation from high-magnification teacher features
Usage
import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoModel
eval_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])
model = AutoModel.from_pretrained("AI4PATH/XMAG", trust_remote_code=True)
model.eval()
image = Image.open("patch.png").convert("RGB")
pixel_values = eval_transform(image).unsqueeze(0)
with torch.no_grad():
outputs = model(pixel_values)
cls_embedding = outputs["cls_embedding"] # (1, 768)
patch_embeddings = outputs["patch_embeddings"] # (1, 256, 768)
pixel_values should be normalized RGB tensors with shape (B, 3, 224, 224).
The model does not apply preprocessing internally. Resize, ToTensor(), and ImageNet mean/std normalization should be done before calling the model.
Model Details
The student sees a low-magnification 224 x 224 view of the tissue region. During training, the frozen UNIv2 teacher sees the corresponding high-magnification 896 x 896 region split into 4 x 4 subpatches. The student is trained to match both:
- A global teacher representation, computed by averaging the 16 teacher local features.
- Local teacher representations, aligned to pooled student patch-token blocks.
This release contains the student EMA backbone weights only. Projection heads used during distillation are not included.
Requirements
pip install torch torchvision transformers huggingface_hub
The model code uses torch.hub to instantiate the DINOv2 ViT-B/14 backbone. The first load may need internet access to fetch the DINOv2 hub code, unless it is already cached.
Citation
If you use this model, please cite:
Su, Z., Akbar, A. R., & Niazi, M. K. K. (2025). Streamline pathology foundation model by cross-magnification distillation. arXiv preprint arXiv:2509.23097. Available at arXiv:2509.23097.
@article{su2025streamline,
title={Streamline pathology foundation model by cross-magnification distillation},
author={Su, Ziyu and Akbar, Abdul Rehman and Sajjad, Usama and Parwani, Anil V and Niazi, Muhammad Khalid Khan},
journal={arXiv preprint arXiv:2509.23097},
year={2025}
}