Add ELF slide encoder weights and model card
Browse files- README.md +116 -0
- config.json +12 -0
- elf_slide_encoder.pth +3 -0
- modeling_elf.py +131 -0
- requirements.txt +4 -0
README.md
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: gpl-3.0
|
| 3 |
+
library_name: pytorch
|
| 4 |
+
pipeline_tag: feature-extraction
|
| 5 |
+
tags:
|
| 6 |
+
- pathology
|
| 7 |
+
- histopathology
|
| 8 |
+
- computational-pathology
|
| 9 |
+
- slide-encoder
|
| 10 |
+
- abmil
|
| 11 |
+
- foundation-model
|
| 12 |
+
- ensemble
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
# ELF slide encoder
|
| 16 |
+
|
| 17 |
+
**ELF** (Ensemble Learning of Foundation models) is a slide-level pathology encoder from the Li Lab at Stanford. It aggregates tile embeddings from five public foundation models into a unified whole-slide representation.
|
| 18 |
+
|
| 19 |
+
Paper: [Ensemble learning of pathology foundation models for precision oncology](https://arxiv.org/abs/2508.16085)
|
| 20 |
+
Code: [github.com/lilab-stanford/ELF](https://github.com/lilab-stanford/ELF)
|
| 21 |
+
Weights: [huggingface.co/luoxd96/ELF](https://huggingface.co/luoxd96/ELF)
|
| 22 |
+
|
| 23 |
+
This Hub repo contains the **inference-only** slide encoder: interpolate → LayerNorm → 8-head gated ABMIL (~151k parameters, ~0.6 MB). It is the trained `momentum_enc` forward used in the paper, not the MoCo training snapshot.
|
| 24 |
+
|
| 25 |
+
## Model details
|
| 26 |
+
|
| 27 |
+
| | |
|
| 28 |
+
|---|---|
|
| 29 |
+
| Developed by | Li Lab, Stanford University |
|
| 30 |
+
| Model type | Slide-level ABMIL aggregator |
|
| 31 |
+
| Inputs | Patch features \(X \in \mathbb{R}^{N \times C}\), \(C \in \{768, 1024, 1280, 1536\}\) |
|
| 32 |
+
| Outputs | `features_dim` \([1, C]\), `features` \([1, 768]\), `attention_weights` \([1, 1, N]\) |
|
| 33 |
+
| Pretraining | 53,699 WSIs, 20 anatomical sites |
|
| 34 |
+
| Tile encoders | UNI, CONCH v1.5, Prov-GigaPath, Virchow2, H-optimus-0 |
|
| 35 |
+
| Magnification | 10× patches (CLAM) |
|
| 36 |
+
| License | GPLv3, non-commercial academic use |
|
| 37 |
+
|
| 38 |
+
The paper ensemble is the **concatenation** of the five `features_dim` vectors (one per tile foundation model), not an average.
|
| 39 |
+
|
| 40 |
+
| Tile model | `--fm` | native `C` | `features_dim` |
|
| 41 |
+
|---|---|---|---|
|
| 42 |
+
| UNI | `uni` | 1024 | 1024 |
|
| 43 |
+
| CONCH v1.5 | `conch_v1_5` | 768 | 768 |
|
| 44 |
+
| Prov-GigaPath | `gigapath` | 1536 | 1536 |
|
| 45 |
+
| H-optimus-0 | `h0` | 1536 | 1536 |
|
| 46 |
+
| Virchow2 | `virchow2` | 2560 → 1280 (CLS+mean) | 1280 |
|
| 47 |
+
|
| 48 |
+
## How to use
|
| 49 |
+
|
| 50 |
+
```bash
|
| 51 |
+
pip install torch huggingface_hub h5py numpy
|
| 52 |
+
```
|
| 53 |
+
|
| 54 |
+
Download this file from the Hub (`modeling_elf.py`) or clone the [GitHub repo](https://github.com/lilab-stanford/ELF) and use `inference/model.py`.
|
| 55 |
+
|
| 56 |
+
```python
|
| 57 |
+
import torch
|
| 58 |
+
from huggingface_hub import hf_hub_download
|
| 59 |
+
|
| 60 |
+
# modeling_elf.py from this repo
|
| 61 |
+
from modeling_elf import ELFSlideEncoder, preprocess_patch_features
|
| 62 |
+
|
| 63 |
+
model = ELFSlideEncoder.from_pretrained("luoxd96/ELF", device="cuda")
|
| 64 |
+
|
| 65 |
+
# patches: [N, C] tile features for one slide and one foundation model
|
| 66 |
+
patches = torch.from_numpy(features).float()
|
| 67 |
+
patches = preprocess_patch_features(patches, foundation_model="uni") # virchow2 averages CLS+mean
|
| 68 |
+
|
| 69 |
+
with torch.inference_mode():
|
| 70 |
+
x = patches.unsqueeze(0).cuda()
|
| 71 |
+
features_dim, features_768, attn = model(x)
|
| 72 |
+
|
| 73 |
+
# features_dim: [1, C] — use this for the paper ensemble
|
| 74 |
+
# features_768: [1, 768]
|
| 75 |
+
# attn: [1, 1, N]
|
| 76 |
+
```
|
| 77 |
+
|
| 78 |
+
Batch extract from CLAM-style `h5` files (`dataset "features"`):
|
| 79 |
+
|
| 80 |
+
```bash
|
| 81 |
+
git clone https://github.com/lilab-stanford/ELF.git
|
| 82 |
+
cd ELF
|
| 83 |
+
|
| 84 |
+
python inference/extract_multiple_model_slide_embedding.py \
|
| 85 |
+
--checkpoint $(python -c "from huggingface_hub import hf_hub_download; print(hf_hub_download('luoxd96/ELF','elf_slide_encoder.pth'))") \
|
| 86 |
+
--input-dir /path/to/uni/<dataset>/h5_files \
|
| 87 |
+
--output-path /path/to/uni_elf \
|
| 88 |
+
--feature-models uni \
|
| 89 |
+
--gpu 0
|
| 90 |
+
```
|
| 91 |
+
|
| 92 |
+
Repeat for `conch_v1_5`, `gigapath`, `virchow2`, `h0`. Concatenate the five `features_dim` vectors for downstream linear probing, as in `evaluation/`.
|
| 93 |
+
|
| 94 |
+
## Intended use
|
| 95 |
+
|
| 96 |
+
Research feature extraction for computational pathology (classification, biomarker prediction, therapy-response studies). This is **not** a diagnostic device and should not be used for clinical decision-making without independent validation.
|
| 97 |
+
|
| 98 |
+
## Limitations
|
| 99 |
+
|
| 100 |
+
- Requires pre-extracted tile embeddings from the five foundation models above; it does not encode RGB tiles.
|
| 101 |
+
- Virchow2 inputs of dimension ≥ 2560 are averaged as CLS + mean → 1280, matching the paper.
|
| 102 |
+
- Licensed for non-commercial academic use.
|
| 103 |
+
|
| 104 |
+
## Citation
|
| 105 |
+
|
| 106 |
+
```bibtex
|
| 107 |
+
@misc{luo2026ensemblelearningpathologyfoundation,
|
| 108 |
+
title={Ensemble learning of pathology foundation models for precision oncology},
|
| 109 |
+
author={Xiangde Luo and Xiyue Wang and Feyisope Eweje and Xiaoming Zhang and Juan Luis Gomez Marti and Sarah Cascarino and Sen Yang and Yuchen Li and Ryan Quinton and Jinxi Xiang and Yuanfeng Ji and Zhe Li and Yijiang Chen and Colin Bergstrom and Ted Kim and Francesca Maria Olguin and Kelley Yuan and Matthew Abikenari and Andrew Heider and Sierra Willens and Sanjeeth Rajaram and Robert West and Joel Neal and Adam Schoenfeld and Maximilian Diehn and Chad Vanderbilt and Ruijiang Li},
|
| 110 |
+
year={2026},
|
| 111 |
+
eprint={2508.16085},
|
| 112 |
+
archivePrefix={arXiv},
|
| 113 |
+
primaryClass={cs.CV},
|
| 114 |
+
url={https://arxiv.org/abs/2508.16085},
|
| 115 |
+
}
|
| 116 |
+
```
|
config.json
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_type": "elf_slide_encoder",
|
| 3 |
+
"architectures": ["ELFSlideEncoder"],
|
| 4 |
+
"embed_dim": 768,
|
| 5 |
+
"num_heads": 8,
|
| 6 |
+
"n_params": 151304,
|
| 7 |
+
"encoder": "momentum_enc.norm+attn",
|
| 8 |
+
"weights": "elf_slide_encoder.pth",
|
| 9 |
+
"paper": "https://arxiv.org/abs/2508.16085",
|
| 10 |
+
"code": "https://github.com/lilab-stanford/ELF",
|
| 11 |
+
"hub": "https://huggingface.co/luoxd96/ELF"
|
| 12 |
+
}
|
elf_slide_encoder.pth
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:84464b4c83834f4b82df8e12b569c51af77e52a3d5010b8b4e2310431d5da23d
|
| 3 |
+
size 621629
|
modeling_elf.py
ADDED
|
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ELF slide encoder: interpolate → LayerNorm → 8-head ABMIL."""
|
| 2 |
+
from __future__ import annotations
|
| 3 |
+
|
| 4 |
+
from collections import OrderedDict
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import Optional, Union
|
| 7 |
+
|
| 8 |
+
import torch
|
| 9 |
+
import torch.nn as nn
|
| 10 |
+
import torch.nn.functional as F
|
| 11 |
+
|
| 12 |
+
DEFAULT_REPO_ID = "luoxd96/ELF"
|
| 13 |
+
WEIGHTS_FILE = "elf_slide_encoder.pth"
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class BatchedABMIL(nn.Module):
|
| 17 |
+
def __init__(self, dim: int):
|
| 18 |
+
super().__init__()
|
| 19 |
+
self.attention_a = nn.Sequential(nn.Linear(dim, dim), nn.Tanh())
|
| 20 |
+
self.attention_b = nn.Sequential(nn.Linear(dim, dim), nn.Sigmoid())
|
| 21 |
+
self.attention_c = nn.Linear(dim, 1)
|
| 22 |
+
|
| 23 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 24 |
+
return self.attention_c(self.attention_a(x) * self.attention_b(x))
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class ELFSlideEncoder(nn.Module):
|
| 28 |
+
def __init__(self, embed_dim: int = 768, num_heads: int = 8):
|
| 29 |
+
super().__init__()
|
| 30 |
+
if embed_dim % num_heads != 0:
|
| 31 |
+
raise ValueError(f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})")
|
| 32 |
+
self.embed_dim = embed_dim
|
| 33 |
+
self.num_heads = num_heads
|
| 34 |
+
self.norm = nn.LayerNorm(embed_dim)
|
| 35 |
+
self.attn = nn.ModuleList(
|
| 36 |
+
[BatchedABMIL(embed_dim // num_heads) for _ in range(num_heads)]
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
def forward(self, x: torch.Tensor, lens: Optional[torch.Tensor] = None):
|
| 40 |
+
"""
|
| 41 |
+
Args:
|
| 42 |
+
x: ``[B, N, C]`` patch features (``C`` in {768, 1024, 1280, 1536}).
|
| 43 |
+
lens: ``[B]`` native ``C`` per item; defaults to ``x.shape[-1]``.
|
| 44 |
+
|
| 45 |
+
Returns:
|
| 46 |
+
features_dim: ``[B, C]`` — ``softmax(ᾱ)ᵀ X``
|
| 47 |
+
features: ``[B, 768]`` — ``softmax(ᾱ)ᵀ X_768``
|
| 48 |
+
attention: ``[B, 1, N]``
|
| 49 |
+
"""
|
| 50 |
+
if x.ndim != 3:
|
| 51 |
+
raise ValueError(f"expected [B, N, C], got {tuple(x.shape)}")
|
| 52 |
+
batch, n_tiles, feat_dim = x.shape
|
| 53 |
+
if lens is None:
|
| 54 |
+
lens = torch.full((batch,), feat_dim, dtype=torch.long, device=x.device)
|
| 55 |
+
|
| 56 |
+
x768 = []
|
| 57 |
+
for i in range(batch):
|
| 58 |
+
c = int(lens[i].item())
|
| 59 |
+
x768.append(
|
| 60 |
+
F.interpolate(x[i, :, :c].unsqueeze(0), size=self.embed_dim, mode="linear", align_corners=True).squeeze(0)
|
| 61 |
+
)
|
| 62 |
+
x768 = self.norm(torch.stack(x768, dim=0))
|
| 63 |
+
|
| 64 |
+
head_dim = self.embed_dim // self.num_heads
|
| 65 |
+
heads = x768.view(batch, n_tiles, head_dim, self.num_heads)
|
| 66 |
+
logits = torch.stack([self.attn[h](heads[:, :, :, h]) for h in range(self.num_heads)], dim=-1)
|
| 67 |
+
attn = F.softmax(logits.mean(dim=-1).transpose(1, 2), dim=-1)
|
| 68 |
+
|
| 69 |
+
feat_768 = torch.bmm(attn, x768)[:, 0]
|
| 70 |
+
feat_native = torch.stack(
|
| 71 |
+
[torch.bmm(attn[i : i + 1], x[i : i + 1, :, : int(lens[i].item())])[0, 0] for i in range(batch)]
|
| 72 |
+
)
|
| 73 |
+
return feat_native, feat_768, attn
|
| 74 |
+
|
| 75 |
+
@classmethod
|
| 76 |
+
def from_pretrained(
|
| 77 |
+
cls,
|
| 78 |
+
repo_id: str = DEFAULT_REPO_ID,
|
| 79 |
+
filename: str = WEIGHTS_FILE,
|
| 80 |
+
device: Union[str, torch.device] = "cpu",
|
| 81 |
+
embed_dim: int = 768,
|
| 82 |
+
num_heads: int = 8,
|
| 83 |
+
) -> "ELFSlideEncoder":
|
| 84 |
+
from huggingface_hub import hf_hub_download
|
| 85 |
+
|
| 86 |
+
path = hf_hub_download(repo_id=repo_id, filename=filename)
|
| 87 |
+
return load_encoder(path, device=device, embed_dim=embed_dim, num_heads=num_heads)
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def preprocess_patch_features(features: torch.Tensor, foundation_model: Optional[str] = None) -> torch.Tensor:
|
| 91 |
+
x = features.float()
|
| 92 |
+
if (foundation_model or "").lower() == "virchow2" and x.shape[-1] >= 2560:
|
| 93 |
+
x = 0.5 * (x[..., :1280] + x[..., 1280:2560])
|
| 94 |
+
return x
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def _unwrap_state_dict(raw) -> OrderedDict:
|
| 98 |
+
if isinstance(raw, dict) and "state_dict" in raw:
|
| 99 |
+
raw = raw["state_dict"]
|
| 100 |
+
return OrderedDict((k[7:] if k.startswith("module.") else k, v) for k, v in raw.items())
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def extract_inference_weights(state_dict: dict) -> OrderedDict:
|
| 104 |
+
keys = list(state_dict.keys())
|
| 105 |
+
prefix = ""
|
| 106 |
+
if any(k.startswith("momentum_enc.") for k in keys):
|
| 107 |
+
prefix = "momentum_enc."
|
| 108 |
+
keep = ("norm.", "attn.")
|
| 109 |
+
out = OrderedDict(
|
| 110 |
+
(k[len(prefix) :], v)
|
| 111 |
+
for k, v in state_dict.items()
|
| 112 |
+
if k.startswith(prefix) and k[len(prefix) :].startswith(keep)
|
| 113 |
+
)
|
| 114 |
+
if not out:
|
| 115 |
+
raise KeyError(f"no norm/attn weights found; prefixes={sorted({k.split('.')[0] for k in keys})[:12]}")
|
| 116 |
+
return out
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def load_encoder(
|
| 120 |
+
checkpoint: Union[str, Path],
|
| 121 |
+
device: Union[str, torch.device] = "cpu",
|
| 122 |
+
embed_dim: int = 768,
|
| 123 |
+
num_heads: int = 8,
|
| 124 |
+
) -> ELFSlideEncoder:
|
| 125 |
+
ckpt = torch.load(str(checkpoint), map_location="cpu", weights_only=False)
|
| 126 |
+
weights = extract_inference_weights(_unwrap_state_dict(ckpt))
|
| 127 |
+
model = ELFSlideEncoder(embed_dim=embed_dim, num_heads=num_heads)
|
| 128 |
+
missing, unexpected = model.load_state_dict(weights, strict=True)
|
| 129 |
+
if missing or unexpected:
|
| 130 |
+
raise RuntimeError(f"load mismatch missing={missing} unexpected={unexpected}")
|
| 131 |
+
return model.to(device).eval()
|
requirements.txt
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.0
|
| 2 |
+
huggingface_hub
|
| 3 |
+
h5py
|
| 4 |
+
numpy
|