Multi-Modal Foundation · MahmoodLab
UNI2-h
Vision transformer generating 1,536-dimensional embeddings from histology patches for pathology analysis.
What UNI2-h is
UNI2-h is a second-generation pathology foundation model built by the Mahmood Lab AI for Pathology at Harvard and Brigham and Women's Hospital. It functions as a general-purpose vision backbone designed to extract feature representations from digitized histopathology images across diverse biological and medical contexts.
Computational pathology often struggles with high morphological heterogeneity across tissues and staining variations. UNI2-h addresses this by pre-training a large vision transformer directly on vast amounts of pathology image tiles, yielding general representations that capture nuanced tissue architecture and cellular detail.
The model provides pre-extracted image features for regions of interest, enabling researchers to perform downstream analysis such as disease classification, search retrieval, segmentation, and slide-level inference without training large neural networks from scratch.
How it works
UNI2-h is based on a custom Vision Transformer Giant architecture (ViT-H/14) containing 681 million parameters. It utilises a patch size of 14 by 14 pixels, an embedding dimension of 1,536 across 24 attention heads, eight register tokens, and a SwiGLU feed-forward network with SiLU activations. The architecture supports dynamic image sizing, natively processing inputs such as 224 by 224 pixel patches.
The model was trained self-supervised using the DINOv2 framework, combining DINO multi-crop self-distillation loss, iBOT masked-image modeling, and KoLeo regularization applied to class tokens. Training was conducted in bfloat16 mixed precision using PyTorch Fully Sharded Data Parallel (PyTorch-FSDP) across Nvidia A100 80GB GPUs.
Key facts
- Model architecture
- Custom ViT-H/14 with SwiGLU FFN
- Parameters
- 681 million
- Embedding dimension
- 1536
- Patch size
- 14
- Input image size
- 224 x 224 pixels (dynamic)
- Training objective
- DINOv2 (self-distillation, iBOT masked-image modeling, KoLeo regularization)
- Training hardware
- Nvidia A100 80GB GPUs
What you can use it for
- Region-of-interest classification
- Extracted patch representations can be categorised via frozen features using linear probing, logistic regression, k-nearest neighbours, or nearest centroid approaches like SimpleShot.
- Whole-slide classification
- Embeddings from bags of class tokens extracted across whole-slide images can be aggregated using multiple instance learning (MIL) frameworks like CLAM.
- Image retrieval
- Morphologically similar regions of interest can be searched and retrieved using nearest-neighbour matching across pre-computed feature representations.
- Dense prediction and tissue segmentation
- The backbone can be fine-tuned using dense prediction frameworks such as ViTDet or ViT-Adapter with Mask2Former to segment cells and tissue structures.
Run it here
We list this model but cannot run it yet: excluded (gated): Gated and CC-BY-NC-ND-4.0..
Everything on this page still applies. Running it needs a loader for this repository, not a bigger machine.
Strengths and limits
Good at
- Pre-trained on over 200 million histology tiles covering both H&E and immunohistochemistry preparations.
- Produces competitive performance on region-of-interest and slide-level classification tasks using frozen features without parameter fine-tuning.
- Employs a custom ViT-H architecture equipped with register tokens and SwiGLU layers to enhance representation quality.
- Offers pre-computed feature sets for major public benchmarks including TCGA, CPTAC, and PANDA.
Watch out for
- Dense prediction tasks such as segmentation generally require fine-tuning with external frameworks rather than using frozen features directly.
- Operates as a vision-only feature extractor and does not natively parse or output textual diagnoses or multimodal descriptions.
- High parameter count (681M) demands substantial memory and computational resources for downstream model fine-tuning.
Inputs and outputs
You provide
Image
PNG, JPEG or TIFF. Sent to the model as-is.
You get back
Mean-pooled representation, one vector per input.
Licence
Check the licence before commercial use · cc-by-nc-nd-4.0
This model is gated by its authors — its licence has to be accepted on the source repository before it will run.
Questions
Can I use UNI2-h commercially?
Check the licence before commercial use. This model is gated by its authors — its licence has to be accepted on the source repository before it will run.
What staining protocols are represented in the pre-training data?
The model was trained on image tiles sampled from whole-slide images prepared with both haematoxylin and eosin (H&E) and immunohistochemistry (IHC) staining from Mass General Brigham.
What is the output shape of the feature extractor?
When evaluating an input patch of 224 by 224 pixels, the model outputs a 1,536-dimensional feature vector per patch.
Does UNI2-h require fine-tuning for all applications?
No. The model can be used with frozen features for region-of-interest classification, image retrieval, and whole-slide multiple instance learning, though fine-tuning is recommended for segmentation tasks.
Which software libraries can run the model?
The model can be defined and loaded using PyTorch and the timm library (version 0.9.8 or higher).
How to cite it
Chen, R.J., Ding, T., Lu, M.Y., Williamson, D.F.K., et al. Towards a general-purpose foundation model for computational pathology. Nat Med (2024). https://doi.org/10.1038/s41591-024-02857-3
Written from UNI2-h's own documentation (model card). Statements citing a figure that does not appear there are removed before publication. Licence and running details come from our registry.