Ace3Z/graphist-v2

GrapHist v2 + VICReg: Edge-Informed Graph Self-Supervised Learning for Histopathology

GrapHist v2 extends GrapHist (arXiv:2603.00143) with edge-informed message passing: cell graphs carry a 75-dimensional feature vector per edge describing the inter-cellular region, and the encoder becomes an ACM-GINEConv that injects those edge features into aggregation.

The un-regularised v2 encoder dimensionally collapsed (pca_1 ≈ 0.5, effective dim ≈ 2). Adding a VICReg variance/covariance regulariser fixes it: this checkpoint ends collapse-free at pca_1 = 0.17, effective dim 11.9, and improves over vanilla GrapHist on transfer subtyping, cell-level phenotyping and survival.

Self-supervision follows GraphMAE (masked node-feature reconstruction, scaled cosine error) plus the VICReg term on the pooled graph embedding.

Repository Structure

graphist_V2.pt          # final v2+VICReg checkpoint - 126,597,010 B, md5 81a2e6b91cefff0bbbc13c6fd318ee78
models/
├── __init__.py         # build_model(args) factory
├── edcoder.py          # PreModel encoder-decoder wrapper
├── acm_gineconv.py     # ACM-GINEConv backbone (v2, edge-informed)
├── acm_gin.py          # ACM-GIN backbone (v1) - REQUIRED IMPORT, see note
└── utils.py            # activation helpers
graphist_utils.py       # NormalizeData + AddVirtualNode transforms, GraphDataset,
                        # filter_valid_graph_paths, checkpoint helpers - REQUIRED
rebase_graph_paths.py   # REQUIRED before using data/v2_graphs/ - see data/README.md
checkpoints/
└── graphist_v1.pt      # vanilla GrapHist v1 (ACM-GIN, 1-dim edges) - AdapterGNN starting point
data/                   # all datasets and embeddings - see data/README.md for the full
                        # artifact -> purpose -> branch -> paper-table map
├── v2_graphs/          #   75-dim edge cell graphs (BACH, BRACS, BreakHis)
├── v2_embeddings/      #   slide (4 cohorts) + cell (22 tars) embeddings from graphist_V2.pt
├── v1_graphs/          #   NuCLS + PanNuke 1-dim edge graphs (AdapterGNN, homophily)
├── v1_graphs_upstream/ #   byte-identical mirrors of the upstream v1 releases:
│                       #   BRACS + BreakHis tars, and tcga_brca/ (254 files, 271 GB)
├── homophily_tiles/    #   109,904 CellViT++ pseudo-labelled 224px tile graphs
├── baseline_embeddings/#   DINOv2, MAE, GrapHist-v1 TCGA-BRCA features (Table 3 baselines)
├── homophily_bach_cellvit/ # 400 per-image CellViT pseudo-label CSVs (BACH fat-tail)
├── preproc_final/      #   Table C.8 preprocessing-runtime run (19,200 BACH tiles)
└── survival/           #   GDC TCGA-BRCA clinical.tsv - required for Table 3
README.md

Which files do I need? For inference only: graphist_V2.pt, models/, graphist_utils.py. Everything under data/ is there to reproduce a specific published number — data/README.md says which one, for which branch.

⚠️ If you use data/v2_graphs/, you must run rebase_graph_paths.py first, or the loader drops every graph — non-fatally, with a single WARNING: line, then an unhelpful IndexError. data/README.md has the exact commands and the expected counts to check against.

acm_gin.py is required even though v2 never uses it. models/edcoder.py imports ACM_GIN_model and ACM_GINEConv_model at module top level; omitting it is an ImportError.

graphist_utils.py is required. Reproducing the published embeddings needs Compose([ToUndirected(), NormalizeData(...), AddVirtualNode(...)]), and neither transform lives in models/.

Which checkpoint this is

The run wrote four checkpoint files. All four contain identical weights — max |ΔW| = 0 across 332/332 tensors, same run_id and same best_loss; they differ only in optimizer state and in the epoch field (99 for three, 100 for _final), which is why their byte sizes differ. So the old "which of the four?" question is moot: any of them yields the same encoder.

This file is the best-validation-loss one (gineconv_vicreg_full_checkpoint.pt), and the one that produced every reported v2+VICReg number:

Field	Value
epoch	99
best_loss	0.019939
run_id	1udpmgxw
md5	81a2e6b91cefff0bbbc13c6fd318ee78

Parameters: encoder 7,980,649 (7.98 M); 9.29 M on the inference path, because embed() ends with rep = self.encoder_to_decoder(enc_rep) (edcoder.py:324), adding 1,310,720; 10,528,542 (10.53 M) for the full pretraining model (encoder + encoder_to_decoder + decoder 1,237,077 + mask token 96, counting each shared module once). A naive state_dict sum reports 19.35 M because acm_gineconv.py registers each channel MLP twice — 156 of the 332 keys are aliases of the same tensors (332 keys, 176 unique storages).

Requirements

torch >= 2.2
torch_geometric >= 2.5
numpy, pandas

Verified on torch 2.10.0+cu128 / PyG 2.7.0. CPU-only inference works.

Usage

1. Download

from huggingface_hub import snapshot_download
path = snapshot_download(repo_id="Ace3Z/graphist-v2")

2. Load

⚠️ Three arguments are load-critical — omit any one and you get either an AttributeError or a load_state_dict shape/key mismatch. Two are read with a getattr default that does not match this checkpoint; the third is a plain attribute read:

Argument	Must be	If omitted	Failure
edge_distance_in_proj	False	getattr default True (models/__init__.py:29)	edge_input_proj built as (512, 75); checkpoint has (512, 74)
encoder_norm	"layer"	getattr default "none" (models/__init__.py:22)	encoder.layer_norms.* missing from the model
concat_hidden	True	AttributeError — models/__init__.py:40 is args.concat_hidden, a plain attribute read with no default	if instead set to False: encoder_to_decoder built as (512, 512); checkpoint has (512, 2560)

(The generate_embs.py CLI defaults --concat_hidden to True, and False appears only as a PreModel.__init__ signature default that the Args flow below never reaches — so concat_hidden has no effective "silent wrong default", it simply must be present.)

import sys, torch
sys.path.insert(0, path)
from models import build_model

class Args:
    # --- architecture ---
    encoder = "acm_gineconv"
    decoder = "acm_gineconv"
    num_features = 96          # node features
    num_edge_features = 75     # edge features (projection sees 74; distance excluded)
    num_hidden = 512
    num_layers = 5
    concat_hidden = True       # load-critical
    encoder_norm = "layer"     # load-critical
    input_norm = "none"
    edge_distance_in_proj = False   # load-critical
    batchnorm = False
    activation = "prelu"
    norm = None
    residual = False
    # --- SSL objective ---
    loss_fn = "sce"
    alpha_l = 3
    mask_rate = 0.5
    replace_rate = 0.1
    drop_edge_rate = 0.0
    # --- VICReg (training only) ---
    vicreg_var_weight = 0.05
    vicreg_cov_weight = 0.002
    vicreg_gamma = 1.0
    # --- unused, vestigial from the GAT lineage, but read by build_model ---
    num_heads = 4
    num_out_heads = 1
    in_drop = 0.2
    attn_drop = 0.1
    negative_slope = 0.2

model = build_model(Args())
ckpt = torch.load(f"{path}/graphist_V2.pt", map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"], strict=True)   # note the key name
model.eval()

The checkpoint is a training dict — {epoch, model_state_dict, optimizer_state_dict, best_loss, run_id}. The weights are under model_state_dict, not model.

3. Inference

# x: [num_nodes, 96]   edge_index: [2, num_edges]   edge_attr: [num_edges, 75]
# batch: [num_nodes] graph assignment (zeros for a single graph)
with torch.no_grad():
    node_emb = model.embed(x, edge_index, edge_attr, batch)   # -> [num_nodes, 512]

Region/slide-level embeddings are the mean over node embeddings; slide-level is the mean over its tiles. For results matching the paper, apply the same transform pipeline used at training time (NormalizeData with the dataset's normalization.json, then AddVirtualNode) — see graphist_utils.py.

Verification

This repository was validated end-to-end before publishing: build_model(Args()) + load_state_dict(..., strict=True) returns 0 missing / 0 unexpected keys using only the files shipped here, and embed() returns a finite [n, 512] tensor.

Acknowledgements

Built on GrapHist (Ogut et al.), GraphMAE (Hou et al., 2022), ACM (Luan et al., 2022), GINEConv (Hu et al., 2020) and VICReg (Bardes et al., 2022). Developed at LTS4, EPFL.

Citation

The v2 preprint is not yet available. Please cite the original GrapHist paper:

@article{ogut2026graphist,
  title  = {GrapHist: Graph Self-Supervised Learning for Histopathology},
  author = {{\"O}{\u{g}}{\"u}t, Sevda and Vincent-Cuaz, C{\'e}dric and
            Dubljevic, Natalia and Hurtado, Carlos and Subramanian, Vaishnavi and
            Frossard, Pascal and Thanou, Dorina},
  year   = {2026},
  eprint = {2603.00143},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV}
}