Ace3Z
/

graphist-v2

@@ -9,206 +9,154 @@ tags:
 - pytorch-geometric
 - graph-representation-learning
 - edge-features
-- vicreg
 ---
-# GrapHist v2 + VICReg: Edge-Informed Graph Self-Supervised Learning for Histopathology
-GrapHist v2 extends [GrapHist](https://huggingface.co/ogutsevda/graphist)
-([arXiv:2603.00143](https://arxiv.org/abs/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`](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
 ```python
 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.)
 ```python
 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
-```python
-# 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](https://huggingface.co/ogutsevda/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](https://www.epfl.ch/labs/lts4/).
-## Citation
-The v2 preprint is not yet available. Please cite the original GrapHist paper:
 ```bibtex
 @article{ogut2026graphist,
@@ -222,3 +170,6 @@ The v2 preprint is not yet available. Please cite the original GrapHist paper:
   primaryClass  = {cs.CV}
 }
 ```

 - pytorch-geometric
 - graph-representation-learning
 - edge-features
 ---
+# GrapHist v2 — Edge-Informed Graph Self-Supervised Learning for Histopathology
+GrapHist v2 encodes a histopathology tile as a **cell graph** — one node per nucleus, edges
+from a Delaunay triangulation pruned at 100 µm — and embeds it with a self-supervised
+**ACM-GINEConv** encoder. It extends [GrapHist](https://huggingface.co/ogutsevda/graphist)
+([arXiv:2603.00143](https://arxiv.org/abs/2603.00143)) by making the **edges carry
+information**: each edge holds a 75-dimensional descriptor of the tissue between two cells,
+and the encoder injects those features into message passing instead of treating edges as mere
+connectivity.
+**Trained on 11,149,499 cell graphs from TCGA-BRCA**, 100 epochs, batch 2048, no labels. The
+objective is GraphMAE-style masked node-feature reconstruction (mask 50 %, scaled cosine
+error, α = 3) plus a **VICReg** variance/covariance term on the pooled graph embedding, which
+keeps all embedding dimensions in use (final `pca_1` 0.17, effective dimension 11.9).
+7.98 M-parameter encoder · 512-dim embeddings · 0.093 ms/patch inference.
+## Results
+Slide-level MIL, transfer to unseen cohorts (test macro-F1 %, best of three MIL heads):
+| | BRACS | BreakHis | BACH |
+|---|:--:|:--:|:--:|
+| GrapHist (v1) | 60.30 | 89.37 | **69.16** |
+| **GrapHist v2** | **69.00** | **95.53** | 68.98 |
+Survival on TCGA-BRCA (Cox PH, C-index): **0.793**, vs 0.763 for v1, 0.724 MAE, 0.632 DINOv2.
+Cell-type identification and patch-level subtyping are in [`DATA.md`](DATA.md), which maps
+every artifact here to the result it reproduces.
+## What's in this repo
 ```
+graphist_v2.pt        the model — 126,597,010 B, md5 81a2e6b91cefff0bbbc13c6fd318ee78
+modeling/             code to load and run it, plus helper scripts
+graphs/               cell graphs: TCGA-BRCA (11.1 M, pre-training) · BACH · BRACS ·
+                      BreakHis · SPIDER-breast
+embeddings/           embeddings from this model: slide-level (4 cohorts), cell-level (22 sets)
+baselines/            DINOv2, MAE and GrapHist v1 features for the comparison rows
+labels/               TCGA-BRCA clinical table and slide labels
+studies/              side analyses: homophily, AdapterGNN, batch effect, preprocessing runtime
+upstream_v1/          mirror of the original GrapHist v1 releases
 ```
+For inference you need only `graphist_v2.pt` + `modeling/`. Everything else exists to reproduce
+a specific published number — see [`DATA.md`](DATA.md).
+> ⚠️ **If you use `graphs/`, run `modeling/rebase_graph_paths.py` first.** The label CSVs store
+> bare filenames and the loader resolves them against the working directory, so without this
+> step every graph is silently dropped. `DATA.md` has the commands and the expected counts.
+## Usage
 ```python
 from huggingface_hub import snapshot_download
 path = snapshot_download(repo_id="Ace3Z/graphist-v2")
 ```
+Three arguments are **load-critical** — `build_model` reads two with defaults that do not match
+this checkpoint, and the third has no default at all:
+| Argument | Must be | If wrong |
+|---|---|---|
+| `edge_distance_in_proj` | `False` | `edge_input_proj` built as (512, **75**); checkpoint has (512, **74**) |
+| `encoder_norm` | `"layer"` | `encoder.layer_norms.*` missing from the model |
+| `concat_hidden` | `True` | `AttributeError` if omitted; (512, 512) instead of (512, **2560**) if `False` |
 ```python
 import sys, torch
+sys.path.insert(0, f"{path}/modeling")
 from models import build_model
 class Args:
+    encoder = decoder = "acm_gineconv"
+    num_features = 96           # per-cell features
+    num_edge_features = 75      # per-edge features (projection sees 74; distance excluded)
     num_hidden = 512
     num_layers = 5
+    concat_hidden = True             # load-critical
+    encoder_norm = "layer"           # load-critical
+    edge_distance_in_proj = False    # load-critical
     input_norm = "none"
     batchnorm = False
     activation = "prelu"
     norm = None
     residual = False
     loss_fn = "sce"
     alpha_l = 3
     mask_rate = 0.5
     replace_rate = 0.1
     drop_edge_rate = 0.0
+    vicreg_var_weight = 0.05         # training only
     vicreg_cov_weight = 0.002
     vicreg_gamma = 1.0
+    num_heads = 4                    # vestigial, but read by build_model
     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()
+# 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- and slide-level embeddings are means over node embeddings. To match the published
+numbers, apply the transforms used at training time — `NormalizeData` with the dataset's
+`normalization.json`, then `AddVirtualNode` (both in `modeling/graphist_utils.py`).
+Requires `torch >= 2.2`, `torch_geometric >= 2.5`, `numpy`, `pandas`. Verified on torch
+2.10.0+cu128 / PyG 2.7.0; CPU-only inference works.
+> `modeling/models/acm_gin.py` is required even though this model never uses it —
+> `edcoder.py` imports it at module level.
+## Graph format
+| Tensor | Shape | Meaning |
+|---|---|---|
+| `x` | `[num_nodes, 96]` | per-cell morphology, texture and colour features |
+| `edge_index` | `[2, num_edges]` | Delaunay edges, pruned at 100 µm |
+| `edge_attr` | `[num_edges, 75]` | column 0 = centroid distance (µm); 1–74 describe the inter-cellular region |
+The v1 releases use a 1-dimensional `edge_attr` (distance only) and are **not** interchangeable
+with this model in either direction.
+## Verification
+`build_model` + `load_state_dict(strict=True)` returns 0 missing / 0 unexpected keys using only
+the files shipped here. All 397 BACH slide embeddings regenerate from this checkpoint at
+minimum cosine similarity **0.9999999997** against the published set (a different checkpoint
+gives ≈ 0), confirming these weights produced the released embeddings.
+## Licence and citation
+Released `cc-by-nc-sa-4.0`. Derived from public cohorts — BACH/ICIAR 2018, BreakHis, BRACS,
+TCGA-BRCA, PanNuke, NuCLS and [SPIDER-breast](https://huggingface.co/datasets/histai/SPIDER-breast)
+(`cc-by-nc-4.0`, research use only). These are derived cell-level features, not images; the
+source cohorts keep their own terms, and `labels/tcga_brca_clinical.tsv` is the open-access GDC
+clinical export. Cite the original cohort papers alongside GrapHist.
 ```bibtex
 @article{ogut2026graphist,
   primaryClass  = {cs.CV}
 }
 ```
+Built on 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](https://www.epfl.ch/labs/lts4/).