XMAG

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