Add GrapHist v2 + VICReg final checkpoint (best val loss, epoch 99) with loader files and validated model card

README.md

@@ -0,0 +1,183 @@
+---
+license: cc-by-nc-sa-4.0
+library_name: pytorch
+pipeline_tag: graph-ml
+tags:
+- graph-neural-networks
+- histopathology
+- self-supervised-learning
+- 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 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 - REQUIRED for inference
+README.md
+```
+
+> **`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 training run produced four checkpoints. **This is the best-validation-loss one**, 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.98 M** (what you use for inference); **10.53 M** for the full
+pretraining model (encoder + projection + decoder + mask token, 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.
+
+## 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.** `build_model` reads them via `getattr` with
+defaults that do **not** match this checkpoint — omit any one and `load_state_dict` fails
+on a shape or key mismatch:
+
+| Argument | Must be | Default if omitted | Failure |
+|---|---|---|---|
+| `edge_distance_in_proj` | `False` | `True` | `edge_input_proj` built as (512, **75**); checkpoint has (512, **74**) |
+| `encoder_norm` | `"layer"` | `"none"` | `encoder.layer_norms.*` missing from the model |
+| `concat_hidden` | `True` | `False` | `encoder_to_decoder` built as (512, 512); checkpoint has (512, **2560**) |
+
+```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{graphist2026,
+  title  = {GrapHist: Graph Self-Supervised Learning for Histopathology},
+  year   = {2026},
+  eprint = {2603.00143},
+  archivePrefix = {arXiv}
+}
+```

graphist_V2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4be24e445e5065a01cdb6a157db84bbbda2c44f4dbf263ebd9d61ad26e375dfb
+size 126597010

graphist_utils.py

@@ -0,0 +1,276 @@
+import os
+import stat as stat_module
+import torch
+import signal
+import random
+import numpy as np
+import pandas as pd
+from typing import List
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from torch_geometric.data import Dataset, Data
+from torch_geometric.transforms import BaseTransform
+
+
+def _is_valid_graph_path(p):
+    """A single os.stat (regular file, non-empty). One syscall instead of two."""
+    try:
+        st = os.stat(p)
+        return stat_module.S_ISREG(st.st_mode) and st.st_size > 0
+    except OSError:
+        return False
+
+
+def filter_valid_graph_paths(graph_paths: list, num_workers: int = 32) -> list:
+    """Filter out graph paths that don't exist or have zero size (truncated).
+
+    Parallelized across threads: this is I/O-bound stat() traffic on a network
+    filesystem, so a thread pool gives near-linear speedup. On the full dataset
+    (~8.9M paths) the serial two-stat version took tens of minutes; this is
+    ~1-2 min. Order is preserved (map keeps input order) and the result is
+    identical to the serial check.
+    """
+    with ThreadPoolExecutor(max_workers=num_workers) as ex:
+        flags = list(ex.map(_is_valid_graph_path, graph_paths, chunksize=2000))
+    valid = [p for p, ok in zip(graph_paths, flags) if ok]
+    skipped = len(graph_paths) - len(valid)
+    if skipped:
+        print(f"WARNING: Filtered out {skipped} missing/empty graph files")
+    return valid
+
+
+class GraphDataset(Dataset):
+    def __init__(self, graph_paths: list, transform=None, pre_transform=None):
+        super().__init__(None, transform, pre_transform)
+        self.graph_paths = filter_valid_graph_paths(graph_paths)
+
+    @property
+    def processed_file_names(self):
+        return self.graph_paths
+
+    def len(self):
+        return len(self.graph_paths)
+
+    def get(self, idx):
+        graph_path = self.graph_paths[idx]
+        graph = torch.load(graph_path, weights_only=False)
+        graph.graph_path = graph_path
+        return graph
+
+
+class NormalizeData(BaseTransform):
+    def __init__(
+        self,
+        scale_dict: dict,
+        attrs: List[str] = ["x", "edge_attr"],
+        edge_attr_skip_cols: List[int] | None = None,
+        min_std: float | None = None,
+        clip: float | None = None,
+    ):
+        tensor_dict = defaultdict(lambda: defaultdict())  # convert the data to tensors
+        for key, value in scale_dict.items():
+            tensor_dict[key]["mean"] = torch.tensor(value["mean"])
+            tensor_dict[key]["std"] = torch.tensor(value["std"])
+
+        # Optional std floor: clamp tiny stds upward to prevent near-constant
+        # dims (which have intrinsically tiny variance from the feature pipeline,
+        # e.g. high-freq Fourier descriptors, GLCM moments) from inflating
+        # outlier raw values by 1/std factors of 10^3-10^4.
+        if min_std is not None:
+            for key in tensor_dict:
+                tensor_dict[key]["std"] = torch.clamp(
+                    tensor_dict[key]["std"], min=min_std
+                )
+
+        # For specified edge_attr columns (e.g. column 0 = spatial distance),
+        # set mean=0, std=1 so they pass through untouched.
+        if edge_attr_skip_cols and "edge_attr" in tensor_dict:
+            for col in edge_attr_skip_cols:
+                tensor_dict["edge_attr"]["mean"][col] = 0.0
+                tensor_dict["edge_attr"]["std"][col] = 1.0
+
+        self.tensor_dict = tensor_dict
+        self.attrs = attrs
+        self.clip = clip
+        self.edge_attr_skip_cols = edge_attr_skip_cols or []
+
+    def forward(self, data: Data) -> Data:
+        for store in data.stores:
+            for key, value in store.items(*self.attrs):
+                if value.numel() > 0:
+                    mean = self.tensor_dict[key]["mean"]
+                    std = self.tensor_dict[key]["std"]
+                    value = (value - mean) / std
+                    value = torch.nan_to_num(value, nan=0.0)
+                    # Hard bound on standardized values to stop residual
+                    # cross-dataset explosion (a feature far from the TCGA mean).
+                    # Skip edge_attr columns left raw (e.g. col 0 = distance).
+                    if self.clip is not None:
+                        if key == "edge_attr" and self.edge_attr_skip_cols:
+                            keep = value[:, self.edge_attr_skip_cols].clone()
+                            value = value.clamp(-self.clip, self.clip)
+                            value[:, self.edge_attr_skip_cols] = keep
+                        else:
+                            value = value.clamp(-self.clip, self.clip)
+                    store[key] = value
+        return data
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+
+class AddVirtualNode(BaseTransform):
+    """Add a virtual node connected to every real node.
+
+    Args:
+        mean_edge_distance: Raw (un-normalised) mean spatial distance from
+            the training set.  Used as column 0 of the virtual-node edge
+            features so that the degree-normalised message weighting remains
+            strictly positive.  Remaining columns (visual + relational
+            features) are set to 0.0 because those columns *are* normalised,
+            and 0 represents their mean.
+    """
+
+    def __init__(self, mean_edge_distance: float):
+        self.mean_edge_distance = mean_edge_distance
+
+    def forward(self, data: Data) -> Data:
+        num_nodes = data.num_nodes
+        device = data.x.device if data.x is not None else "cpu"
+
+        virtual_node_feat = torch.zeros((1, data.x.size(-1)), device=device)
+        data.x = torch.cat([data.x, virtual_node_feat], dim=0)
+
+        # Keep any node-level label(s) aligned with the added virtual node: the
+        # cell-level embedding path slices batch.label with the VN-inclusive ptr,
+        # so a per-node label must gain one entry for the VN (placeholder -1).
+        # Scalar / graph-level labels (size != num_nodes) are left untouched.
+        for _lab in ("label", "labels"):
+            _v = getattr(data, _lab, None)
+            if torch.is_tensor(_v) and _v.dim() >= 1 and _v.size(0) == num_nodes:
+                _pad = torch.full(
+                    (1, *_v.shape[1:]), -1, dtype=_v.dtype, device=_v.device
+                )
+                setattr(data, _lab, torch.cat([_v, _pad], dim=0))
+
+        row = torch.arange(num_nodes, device=device)
+        col = torch.full((num_nodes,), num_nodes, device=device)
+        new_edges = torch.stack([torch.cat([row, col]), torch.cat([col, row])], dim=0)
+        data.edge_index = torch.cat([data.edge_index, new_edges], dim=1)
+
+        data.virtual_node_index = torch.tensor([num_nodes], device=device)
+
+        if data.edge_attr is not None:
+            edge_dim = data.edge_attr.size(-1)
+            new_edge_attr = torch.zeros(
+                (2 * num_nodes, edge_dim), device=device
+            )
+            # Column 0 = raw spatial distance (not normalised); use training-
+            # set mean so virtual-node edges have a typical positive weight.
+            new_edge_attr[:, 0] = self.mean_edge_distance
+            data.edge_attr = torch.cat([data.edge_attr, new_edge_attr], dim=0)
+
+        return data
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(mean_edge_distance={self.mean_edge_distance})"
+
+
+class GracefulKiller:
+    kill_now = False
+
+    def __init__(self):
+        signal.signal(signal.SIGINT, self.exit_gracefully)
+        signal.signal(signal.SIGTERM, self.exit_gracefully)
+
+    def exit_gracefully(self, signum, frame):
+        self.kill_now = True
+
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def seed_worker(worker_id):
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+
+def split_data(data_df: pd.DataFrame, split_df: pd.DataFrame):
+
+    train_samples = split_df.loc[split_df["split"] == "train", "sample_id"].unique()
+    val_samples = split_df.loc[split_df["split"] == "val", "sample_id"].unique()
+    test_samples = split_df.loc[split_df["split"] == "test", "sample_id"].unique()
+
+    train_mask = data_df["sample_id"].isin(train_samples)
+    val_mask = data_df["sample_id"].isin(val_samples)
+    test_mask = data_df["sample_id"].isin(test_samples)
+
+    return train_mask, val_mask, test_mask
+
+
+def get_current_lr(optimizer):
+    return optimizer.state_dict()["param_groups"][0]["lr"]
+
+
+def create_optimizer(
+    opt, model, lr, weight_decay, get_num_layer=None, get_layer_scale=None
+):
+    opt_lower = opt.lower()
+
+    parameters = model.parameters()
+    opt_args = dict(lr=lr, weight_decay=weight_decay)
+
+    opt_split = opt_lower.split("_")
+    opt_lower = opt_split[-1]
+    if opt_lower == "adam":
+        optimizer = torch.optim.Adam(parameters, **opt_args)
+    elif opt_lower == "adamw":
+        optimizer = torch.optim.AdamW(parameters, **opt_args)
+    elif opt_lower == "adadelta":
+        optimizer = torch.optim.Adadelta(parameters, **opt_args)
+    elif opt_lower == "radam":
+        optimizer = torch.optim.RAdam(parameters, **opt_args)
+    elif opt_lower == "sgd":
+        opt_args["momentum"] = 0.9
+        return torch.optim.SGD(parameters, **opt_args)
+    else:
+        assert False and "Invalid optimizer"
+
+    return optimizer
+
+
+def load_checkpoint(checkpoint_fpath, model, optimizer, just_model=False):
+    checkpoint = torch.load(checkpoint_fpath, weights_only=False)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    if just_model:
+        return model
+    else:
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        return (
+            model,
+            optimizer,
+            checkpoint["epoch"],
+            checkpoint["best_loss"],
+            checkpoint["run_id"],
+        )
+
+
+def save_checkpoint(checkpoint_fpath, model, optimizer, epoch, best_loss, run_id):
+    torch.save(
+        {
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "best_loss": best_loss,
+            "run_id": run_id,
+        },
+        checkpoint_fpath,
+    )
+    return

models/__init__.py

@@ -0,0 +1,73 @@
+from .edcoder import PreModel
+
+
+def build_model(args):
+    num_heads = args.num_heads
+    num_out_heads = args.num_out_heads
+    num_hidden = args.num_hidden
+    num_layers = args.num_layers
+    residual = args.residual
+    attn_drop = args.attn_drop
+    in_drop = args.in_drop
+    norm = args.norm
+    negative_slope = args.negative_slope
+    encoder_type = args.encoder
+    decoder_type = args.decoder
+    mask_rate = args.mask_rate
+    drop_edge_rate = args.drop_edge_rate
+    replace_rate = args.replace_rate
+    batchnorm = args.batchnorm
+    # Encoder between-layer norm: none | layer | graph. Back-compat: the older
+    # boolean --encoder_layernorm maps to "layer" when --encoder_norm is unset.
+    encoder_norm = getattr(args, "encoder_norm", "none")
+    if encoder_norm == "none" and getattr(args, "encoder_layernorm", False):
+        encoder_norm = "layer"
+    # Learnable normalization of the projected input embedding (encoder stem).
+    encoder_input_norm = getattr(args, "input_norm", "none")
+    # When False, the edge distance (col 0) is used only as the aggregation
+    # weight, not fed into the edge feature projection.
+    edge_distance_in_proj = getattr(args, "edge_distance_in_proj", True)
+    # VICReg-style anti-collapse regularizer weights (0 = off). Applied on the
+    # graph-level pooled embedding during training. Absent on eval-time args
+    # (generate_embs) -> getattr defaults keep those code paths unaffected.
+    vicreg_var_weight = getattr(args, "vicreg_var_weight", 0.0)
+    vicreg_cov_weight = getattr(args, "vicreg_cov_weight", 0.0)
+    vicreg_gamma = getattr(args, "vicreg_gamma", 1.0)
+
+    activation = args.activation
+    loss_fn = args.loss_fn
+    alpha_l = args.alpha_l
+    concat_hidden = args.concat_hidden
+    num_features = args.num_features
+    num_edge_features = args.num_edge_features
+
+    model = PreModel(
+        in_dim=int(num_features),
+        edge_in_dim=int(num_edge_features),
+        num_hidden=int(num_hidden),
+        num_layers=num_layers,
+        nhead=num_heads,
+        nhead_out=num_out_heads,
+        activation=activation,
+        feat_drop=in_drop,
+        attn_drop=attn_drop,
+        negative_slope=negative_slope,
+        residual=residual,
+        encoder_type=encoder_type,
+        decoder_type=decoder_type,
+        mask_rate=mask_rate,
+        norm=norm,
+        loss_fn=loss_fn,
+        drop_edge_rate=drop_edge_rate,
+        replace_rate=replace_rate,
+        alpha_l=alpha_l,
+        concat_hidden=concat_hidden,
+        batchnorm=batchnorm,
+        encoder_norm=encoder_norm,
+        encoder_input_norm=encoder_input_norm,
+        edge_distance_in_proj=edge_distance_in_proj,
+        vicreg_var_weight=vicreg_var_weight,
+        vicreg_cov_weight=vicreg_cov_weight,
+        vicreg_gamma=vicreg_gamma,
+    )
+    return model

models/acm_gin.py

@@ -0,0 +1,337 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch import Tensor
+from typing import Union
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.inits import reset
+from torch_geometric.typing import OptPairTensor, Size
+from torch_geometric.utils import scatter
+
+from .utils import create_activation
+
+
+class ACM_GIN(MessagePassing):
+    """Single ACM-GIN convolution layer with edge-aware message passing.
+
+    The message from node j to node i incorporates the intermediate edge
+    feature e_ij_prime (precomputed by the outer model) alongside the
+    scalar spatial weight a_ij used for degree normalization:
+
+        m_ij = a_ij * ReLU(H_j + e_ij_prime)
+    """
+
+    def __init__(
+        self,
+        nn_lowpass: torch.nn.Module,
+        nn_highpass: torch.nn.Module,
+        nn_fullpass: torch.nn.Module,
+        nn_lowpass_proj: torch.nn.Module,
+        nn_highpass_proj: torch.nn.Module,
+        nn_fullpass_proj: torch.nn.Module,
+        nn_mix: torch.nn.Module,
+        T: float = 3.0,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+        self.nn_lowpass = nn_lowpass
+        self.nn_highpass = nn_highpass
+        self.nn_fullpass = nn_fullpass
+        self.nn_lowpass_proj = nn_lowpass_proj
+        self.nn_highpass_proj = nn_highpass_proj
+        self.nn_fullpass_proj = nn_fullpass_proj
+        self.nn_mix = nn_mix
+        self.sigmoid = torch.nn.Sigmoid()
+        self.softmax = torch.nn.Softmax(dim=1)
+        self.T = T
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn_lowpass)
+        reset(self.nn_highpass)
+        reset(self.nn_fullpass)
+        reset(self.nn_lowpass_proj)
+        reset(self.nn_highpass_proj)
+        reset(self.nn_fullpass_proj)
+        reset(self.nn_mix)
+
+    def forward(
+        self,
+        x: Union[Tensor, OptPairTensor],
+        edge_index: Tensor,
+        edge_weight: Tensor,
+        edge_feat: Tensor,
+        size: Size = None,
+    ) -> Tensor:
+        """Forward pass of a single ACM-GIN layer.
+
+        Args:
+            x: Node features [N, hidden_dim] or (x_src, x_dst) pair.
+            edge_index: Edge indices [2, E].
+            edge_weight: Scalar spatial distance per edge [E] (first column
+                of the original edge_attr, used for degree normalization).
+            edge_feat: Intermediate edge features [E, hidden_dim], i.e.
+                e_ij_prime precomputed by the edge MLP in the outer model.
+            size: Optional bipartite graph size.
+        """
+        if isinstance(x, Tensor):
+            x: OptPairTensor = (x, x)
+
+        # propagate_type: (x: OptPairTensor, edge_weight: Tensor, edge_feat: Tensor)
+        out = self.propagate(
+            edge_index, x=x, edge_weight=edge_weight, edge_feat=edge_feat, size=size
+        )
+
+        # TODO: This computes degree as the sum of edge weights, not the count of
+        # neighbors. Consider whether unweighted degree (standard GIN) would be more
+        # appropriate, especially if edge weights can sum to zero for some nodes.
+        deg = scatter(edge_weight, edge_index[1], 0, out.size(0), reduce="sum")
+        deg_inv = 1.0 / deg
+        deg_inv.masked_fill_(deg_inv == float("inf"), 0)
+        out = deg_inv.view(-1, 1) * out
+
+        x_r = x[1]
+        assert x_r is not None, (
+            "Target node features (x_r) must not be None for ACM_GIN"
+        )
+        out_lowpass = (x_r + out) / 2.0
+        out_highpass = (x_r - out) / 2.0
+
+        # compute embeddings for each filter
+        out_lowpass = self.nn_lowpass(out_lowpass)
+        out_highpass = self.nn_highpass(out_highpass)
+        out_fullpass = self.nn_fullpass(x_r)
+        # compute importance weights per filter
+        alpha_lowpass = self.sigmoid(self.nn_lowpass_proj(out_lowpass))
+        alpha_highpass = self.sigmoid(self.nn_highpass_proj(out_highpass))
+        alpha_fullpass = self.sigmoid(self.nn_fullpass_proj(out_fullpass))
+        alpha_cat = torch.concat([alpha_lowpass, alpha_highpass, alpha_fullpass], dim=1)
+        alpha_cat = self.softmax(self.nn_mix(alpha_cat / self.T))
+
+        out = alpha_cat[:, 0].view(-1, 1) * out_lowpass
+        out = out + alpha_cat[:, 1].view(-1, 1) * out_highpass
+        out = out + alpha_cat[:, 2].view(-1, 1) * out_fullpass
+
+        return out
+
+    def message(self, x_j: Tensor, edge_weight: Tensor, edge_feat: Tensor) -> Tensor:
+        """Edge-aware message: m_ij = a_ij * ReLU(H_j + e_ij_prime)."""
+        return edge_weight.view(-1, 1) * F.relu(x_j + edge_feat)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"nn_lowpass={self.nn_lowpass}, "
+            f"nn_highpass={self.nn_highpass}, "
+            f"nn_fullpass={self.nn_fullpass})"
+        )
+
+
+class ACM_GIN_model(nn.Module):
+    """Multi-layer ACM-GIN model with edge-aware message passing.
+
+    Both node and edge features are projected into hidden_dim at the start
+    (via ``node_input_proj`` and ``edge_input_proj``).  This ensures
+    uniform dimensions throughout, so every edge MLP receives 3 * hidden_dim
+    and the edge residual connection is valid from layer 0 onward.
+
+    At each layer k the model:
+      1. Computes intermediate edge features via an edge MLP:
+         e_ij_prime = MLP_edge(H_i || H_j || E_ij)
+      2. Updates edge state with a residual connection:
+         E_ij^(k) = E_ij^(k-1) + e_ij_prime
+      3. Passes messages using the scalar spatial weight and the
+         intermediate edge features:
+         m_ij = a_ij * ReLU(H_j + e_ij_prime)
+      4. Applies ACM channel mixing (low/high/full-pass) on the
+         aggregated messages.
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        num_layers,
+        hidden_dim,
+        edge_in_dim,
+        batchnorm,
+        activation="relu",
+    ):
+        super(ACM_GIN_model, self).__init__()
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.gnn_batchnorm = batchnorm
+        self.out_dim = out_dim
+
+        # Project raw node and edge features into hidden_dim so that both
+        # live in the same space from the very start.  This enables the
+        # edge residual connection at every layer (including layer 0) and
+        # keeps all edge MLP input dimensions uniform at 3 * hidden_dim.
+        self.node_input_proj = nn.Linear(in_dim, hidden_dim)
+        self.edge_input_proj = nn.Linear(edge_in_dim, hidden_dim)
+
+        self.ACM_convs = nn.ModuleList()
+        self.nns_lowpass = nn.ModuleList()
+        self.nns_highpass = nn.ModuleList()
+        self.nns_fullpass = nn.ModuleList()
+        self.nns_lowpass_proj = nn.ModuleList()
+        self.nns_highpass_proj = nn.ModuleList()
+        self.nns_fullpass_proj = nn.ModuleList()
+        self.nns_mix = nn.ModuleList()
+        self.edge_mlps = nn.ModuleList()
+
+        self.activation_name = activation
+
+        for i in range(self.num_layers):
+            # --- Edge MLP for this layer ---
+            # Both nodes and edges have been projected to hidden_dim before
+            # the loop, so the input is always (src + dst + edge_state) =
+            # 3 * hidden_dim for every layer.
+            edge_mlp_in = 3 * hidden_dim
+
+            if self.gnn_batchnorm:
+                self.edge_mlps.append(
+                    nn.Sequential(
+                        nn.Linear(edge_mlp_in, hidden_dim),
+                        nn.BatchNorm1d(hidden_dim),
+                        create_activation(activation),
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.BatchNorm1d(hidden_dim),
+                        create_activation(activation),
+                    )
+                )
+            else:
+                self.edge_mlps.append(
+                    nn.Sequential(
+                        nn.Linear(edge_mlp_in, hidden_dim),
+                        create_activation(activation),
+                        nn.Linear(hidden_dim, hidden_dim),
+                        create_activation(activation),
+                    )
+                )
+
+            # --- Projection modules to compute importance weights ---
+            for channel_proj_module in [
+                self.nns_lowpass_proj,
+                self.nns_highpass_proj,
+                self.nns_fullpass_proj,
+            ]:
+                if i == self.num_layers - 1:
+                    channel_proj_module.append(nn.Linear(self.out_dim, 1))
+                else:
+                    channel_proj_module.append(nn.Linear(self.hidden_dim, 1))
+
+            # --- Weights mixing module as attention mechanism ---
+            self.nns_mix.append(nn.Linear(3, 3))
+
+            # --- GIN channel MLPs ---
+            # After node_input_proj, all nodes are hidden_dim, so
+            # local_input_dim is always hidden_dim.
+            local_input_dim = self.hidden_dim
+
+            if i == self.num_layers - 1:
+                local_out_dim = self.out_dim
+            else:
+                local_out_dim = self.hidden_dim
+
+            for channel_module in [
+                self.nns_lowpass,
+                self.nns_highpass,
+                self.nns_fullpass,
+            ]:
+                if self.gnn_batchnorm:
+                    sequential = nn.Sequential(
+                        nn.Linear(local_input_dim, self.hidden_dim),
+                        nn.BatchNorm1d(self.hidden_dim),
+                        create_activation(self.activation_name),
+                        nn.Linear(self.hidden_dim, local_out_dim),
+                        nn.BatchNorm1d(local_out_dim),
+                        create_activation(self.activation_name),
+                    )
+                else:
+                    sequential = nn.Sequential(
+                        nn.Linear(local_input_dim, self.hidden_dim),
+                        create_activation(self.activation_name),
+                        nn.Linear(self.hidden_dim, local_out_dim),
+                        create_activation(self.activation_name),
+                    )
+
+                channel_module.append(sequential)
+
+            self.ACM_convs.append(
+                ACM_GIN(
+                    nn_lowpass=self.nns_lowpass[i],
+                    nn_highpass=self.nns_highpass[i],
+                    nn_fullpass=self.nns_fullpass[i],
+                    nn_lowpass_proj=self.nns_lowpass_proj[i],
+                    nn_highpass_proj=self.nns_highpass_proj[i],
+                    nn_fullpass_proj=self.nns_fullpass_proj[i],
+                    nn_mix=self.nns_mix[i],
+                )
+            )
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.reset_parameters()
+            elif isinstance(m, nn.BatchNorm1d):
+                m.reset_parameters()
+
+    def forward(self, x, edge_index, edge_attr, batch=None, return_hidden=False):
+        """Forward pass through all ACM-GIN layers with edge updates.
+
+        `batch` is accepted for API parity with ACM_GINEConv_model (GraphNorm);
+        it is unused here.
+
+        Args:
+            x: Node features [N, in_dim].
+            edge_index: Edge indices [2, E].
+            edge_attr: Edge features [E, edge_in_dim]. The first column
+                (index 0) is the scalar spatial distance used for degree
+                normalization; the full vector evolves through layers via
+                the edge MLPs.
+            return_hidden: If True, also return all intermediate node states.
+
+        Returns:
+            x: Final node embeddings [N, out_dim].
+            outs: (optional) List of node states after each layer.
+        """
+        # Extract scalar spatial distance for degree normalization
+        # (stays fixed across layers)
+        edge_weight = edge_attr[:, 0]
+
+        # Project node and edge features from raw dims to hidden_dim
+        x = self.node_input_proj(x)
+        edge_state = self.edge_input_proj(edge_attr)
+
+        outs = []
+        for i in range(self.num_layers):
+            # Step 1: Compute intermediate edge features
+            src, dst = edge_index
+            edge_mlp_input = torch.cat([x[src], x[dst], edge_state], dim=-1)
+            e_ij_prime = self.edge_mlps[i](edge_mlp_input)
+
+            # Step 2: Update edge state with residual (safe at every layer
+            # because both edge_state and e_ij_prime are hidden_dim)
+            edge_state = edge_state + e_ij_prime
+
+            # Step 3-4: Edge-aware ACM message passing
+            x = self.ACM_convs[i](
+                x=x,
+                edge_index=edge_index,
+                edge_weight=edge_weight,
+                edge_feat=e_ij_prime,
+            )
+            outs.append(x)
+
+        if return_hidden:
+            return x, outs
+        else:
+            return x
+
+
+if __name__ == "__main__":
+    acm_gin = ACM_GIN_model(46, 46, 2, 256, 74, True)
+    print(sum(p.numel() for p in acm_gin.parameters() if p.requires_grad))

models/acm_gineconv.py

@@ -0,0 +1,385 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch import Tensor
+from typing import Union
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.inits import reset
+from torch_geometric.typing import OptPairTensor, Size
+from torch_geometric.utils import scatter
+
+from .utils import create_activation
+
+
+class GraphNorm(nn.Module):
+    """GraphNorm (Cai et al., ICML 2021): per-graph normalization with a
+    learnable mean-shift, applied to node states between GNN layers.
+
+    Faithful to the official implementation (lsj2408/GraphNorm):
+        sub = x - alpha * mean_g(x)          # per-graph mean, learnable shift
+        out = gamma * sub / sqrt(mean_g(sub^2) + eps) + beta
+
+    where mean_g is computed within each graph of the batch. Unlike BatchNorm
+    (cross-graph batch statistics -> noisy on heterogeneous graphs), GraphNorm
+    normalizes each graph independently, which is the principled choice for
+    batches of heterogeneous cell-graphs.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.gamma = nn.Parameter(torch.ones(dim))   # scale (weight)
+        self.beta = nn.Parameter(torch.zeros(dim))    # shift (bias)
+        self.alpha = nn.Parameter(torch.ones(dim))    # learnable mean-shift (mean_scale)
+
+    def reset_parameters(self):
+        nn.init.ones_(self.gamma)
+        nn.init.zeros_(self.beta)
+        nn.init.ones_(self.alpha)
+
+    def forward(self, x: Tensor, batch: Tensor = None) -> Tensor:
+        if batch is None:
+            batch = torch.zeros(x.size(0), dtype=torch.long, device=x.device)
+        num_graphs = int(batch.max().item()) + 1
+        # Per-graph mean, broadcast back to nodes.
+        mean = scatter(x, batch, 0, num_graphs, reduce="mean")[batch]
+        sub = x - self.alpha * mean
+        var = scatter(sub * sub, batch, 0, num_graphs, reduce="mean")[batch]
+        std = (var + self.eps).sqrt()
+        return self.gamma * sub / std + self.beta
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(dim={self.dim})"
+
+
+class ACM_GINEConv(MessagePassing):
+    """Single ACM-GINEConv layer: GINEConv-style message passing + ACM filtering.
+
+    Messages incorporate static edge features via addition before ReLU
+    (following the GINEConv formulation), then aggregated messages are
+    decomposed into low-pass, high-pass, and identity channels with
+    learned adaptive mixing.
+
+    Message:  m_ij = a_ij * ReLU(H_j + E_ji)
+    where E_ji are static encoded edge features (no per-layer edge MLP).
+    """
+
+    def __init__(
+        self,
+        nn_lowpass: torch.nn.Module,
+        nn_highpass: torch.nn.Module,
+        nn_fullpass: torch.nn.Module,
+        nn_lowpass_proj: torch.nn.Module,
+        nn_highpass_proj: torch.nn.Module,
+        nn_fullpass_proj: torch.nn.Module,
+        nn_mix: torch.nn.Module,
+        T: float = 3.0,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+        self.nn_lowpass = nn_lowpass
+        self.nn_highpass = nn_highpass
+        self.nn_fullpass = nn_fullpass
+        self.nn_lowpass_proj = nn_lowpass_proj
+        self.nn_highpass_proj = nn_highpass_proj
+        self.nn_fullpass_proj = nn_fullpass_proj
+        self.nn_mix = nn_mix
+        self.sigmoid = torch.nn.Sigmoid()
+        self.softmax = torch.nn.Softmax(dim=1)
+        self.T = T
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn_lowpass)
+        reset(self.nn_highpass)
+        reset(self.nn_fullpass)
+        reset(self.nn_lowpass_proj)
+        reset(self.nn_highpass_proj)
+        reset(self.nn_fullpass_proj)
+        reset(self.nn_mix)
+
+    def forward(
+        self,
+        x: Union[Tensor, OptPairTensor],
+        edge_index: Tensor,
+        edge_weight: Tensor,
+        edge_feat: Tensor,
+        size: Size = None,
+    ) -> Tensor:
+        """Forward pass of a single ACM-GINEConv layer.
+
+        Args:
+            x: Node features [N, hidden_dim] or (x_src, x_dst) pair.
+            edge_index: Edge indices [2, E].
+            edge_weight: Scalar spatial distance per edge [E] (used for
+                degree normalization).
+            edge_feat: Static edge features [E, hidden_dim] encoded once
+                at the start; same features reused at every layer.
+            size: Optional bipartite graph size.
+        """
+        if isinstance(x, Tensor):
+            x: OptPairTensor = (x, x)
+
+        # propagate_type: (x: OptPairTensor, edge_weight: Tensor, edge_feat: Tensor)
+        out = self.propagate(
+            edge_index, x=x, edge_weight=edge_weight, edge_feat=edge_feat, size=size
+        )
+
+        # Degree normalization using sum of edge weights
+        deg = scatter(edge_weight, edge_index[1], 0, out.size(0), reduce="sum")
+        deg_inv = 1.0 / deg
+        deg_inv.masked_fill_(deg_inv == float("inf"), 0)
+        out = deg_inv.view(-1, 1) * out
+
+        x_r = x[1]
+        assert x_r is not None, (
+            "Target node features (x_r) must not be None for ACM_GINEConv"
+        )
+
+        # ACM frequency decomposition
+        out_lowpass = (x_r + out) / 2.0
+        out_highpass = (x_r - out) / 2.0
+
+        # Compute embeddings for each filter
+        out_lowpass = self.nn_lowpass(out_lowpass)
+        out_highpass = self.nn_highpass(out_highpass)
+        out_fullpass = self.nn_fullpass(x_r)
+
+        # Compute importance weights per filter
+        alpha_lowpass = self.sigmoid(self.nn_lowpass_proj(out_lowpass))
+        alpha_highpass = self.sigmoid(self.nn_highpass_proj(out_highpass))
+        alpha_fullpass = self.sigmoid(self.nn_fullpass_proj(out_fullpass))
+        alpha_cat = torch.concat([alpha_lowpass, alpha_highpass, alpha_fullpass], dim=1)
+        alpha_cat = self.softmax(self.nn_mix(alpha_cat / self.T))
+
+        # Adaptive mixing
+        out = alpha_cat[:, 0].view(-1, 1) * out_lowpass
+        out = out + alpha_cat[:, 1].view(-1, 1) * out_highpass
+        out = out + alpha_cat[:, 2].view(-1, 1) * out_fullpass
+
+        return out
+
+    def message(self, x_j: Tensor, edge_weight: Tensor, edge_feat: Tensor) -> Tensor:
+        """GINEConv-style message: m_ij = a_ij * ReLU(H_j + E_ji)."""
+        return edge_weight.view(-1, 1) * F.relu(x_j + edge_feat)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"nn_lowpass={self.nn_lowpass}, "
+            f"nn_highpass={self.nn_highpass}, "
+            f"nn_fullpass={self.nn_fullpass})"
+        )
+
+
+class ACM_GINEConv_model(nn.Module):
+    """Multi-layer ACM-GINEConv model with static edge feature integration.
+
+    Unlike ACM_GIN_model (which learns per-layer edge MLPs), this model
+    encodes edge features once via a linear projection and reuses them
+    at every GNN layer. This follows the GINEConv philosophy:
+      - Edge features are incorporated into messages via ReLU(H_j + E_ji)
+      - No per-layer edge update MLP (significant memory savings)
+      - ACM filtering and adaptive mixing remain unchanged
+
+    This is a simpler, more memory-efficient alternative suitable for
+    pretraining on a single GPU where the full edge-update model may OOM.
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        num_layers,
+        hidden_dim,
+        edge_in_dim,
+        batchnorm,
+        activation="relu",
+        norm_type="none",
+        input_norm="none",
+        edge_distance_in_proj=True,
+    ):
+        super(ACM_GINEConv_model, self).__init__()
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.gnn_batchnorm = batchnorm
+        if norm_type not in ("none", "layer", "graph"):
+            raise ValueError(f"Unknown norm_type {norm_type!r}; expected none/layer/graph")
+        self.norm_type = norm_type
+        if input_norm not in ("none", "batch", "layer"):
+            raise ValueError(f"Unknown input_norm {input_norm!r}; expected none/batch/layer")
+        self.input_norm_type = input_norm
+        self.out_dim = out_dim
+
+        # Column 0 of edge_attr is the spatial distance. When False, exclude it
+        # from the edge feature projection: the distance is then used ONLY as the
+        # per-edge aggregation weight a_ij (and degree norm), not mixed into the
+        # additive edge feature E_ji. This also removes the raw-distance scale
+        # from the projection input.
+        self.edge_distance_in_proj = edge_distance_in_proj
+        edge_proj_in = edge_in_dim if edge_distance_in_proj else edge_in_dim - 1
+
+        # Project raw node and edge features into hidden_dim once.
+        # Edge features remain static after this projection.
+        self.node_input_proj = nn.Linear(in_dim, hidden_dim)
+        self.edge_input_proj = nn.Linear(edge_proj_in, hidden_dim)
+
+        # Optional LEARNABLE normalization of the projected input embedding
+        # (applied right after node_input_proj). Makes the input conditioning
+        # adaptive; the static json NormalizeData still defines the fixed SCE
+        # target. "batch" = online per-feature stats; "layer" = per-node.
+        if input_norm == "batch":
+            self.input_norm = nn.BatchNorm1d(hidden_dim)
+        elif input_norm == "layer":
+            self.input_norm = nn.LayerNorm(hidden_dim)
+        else:
+            self.input_norm = None
+
+        # Optional per-layer normalization on node states BETWEEN conv layers
+        # (the research-conventional spot for LayerNorm/GraphNorm). Anti-collapse
+        # measure: stops variance concentrating onto a single PCA direction.
+        #   "layer" -> nn.LayerNorm  (per-node, across features)
+        #   "graph" -> GraphNorm     (per-graph, across nodes; needs batch)
+        # (BatchNorm lives inside the channel MLPs via `batchnorm`, the GIN spot.)
+        self.layer_norms = nn.ModuleList() if norm_type in ("layer", "graph") else None
+
+        self.ACM_convs = nn.ModuleList()
+        self.nns_lowpass = nn.ModuleList()
+        self.nns_highpass = nn.ModuleList()
+        self.nns_fullpass = nn.ModuleList()
+        self.nns_lowpass_proj = nn.ModuleList()
+        self.nns_highpass_proj = nn.ModuleList()
+        self.nns_fullpass_proj = nn.ModuleList()
+        self.nns_mix = nn.ModuleList()
+
+        self.activation_name = activation
+
+        for i in range(self.num_layers):
+            # --- Projection modules to compute importance weights ---
+            for channel_proj_module in [
+                self.nns_lowpass_proj,
+                self.nns_highpass_proj,
+                self.nns_fullpass_proj,
+            ]:
+                if i == self.num_layers - 1:
+                    channel_proj_module.append(nn.Linear(self.out_dim, 1))
+                else:
+                    channel_proj_module.append(nn.Linear(self.hidden_dim, 1))
+
+            # --- Weights mixing module as attention mechanism ---
+            self.nns_mix.append(nn.Linear(3, 3))
+
+            # --- ACM channel MLPs ---
+            local_input_dim = self.hidden_dim
+
+            if i == self.num_layers - 1:
+                local_out_dim = self.out_dim
+            else:
+                local_out_dim = self.hidden_dim
+
+            for channel_module in [
+                self.nns_lowpass,
+                self.nns_highpass,
+                self.nns_fullpass,
+            ]:
+                if self.gnn_batchnorm:
+                    sequential = nn.Sequential(
+                        nn.Linear(local_input_dim, self.hidden_dim),
+                        nn.BatchNorm1d(self.hidden_dim),
+                        create_activation(self.activation_name),
+                        nn.Linear(self.hidden_dim, local_out_dim),
+                        nn.BatchNorm1d(local_out_dim),
+                        create_activation(self.activation_name),
+                    )
+                else:
+                    sequential = nn.Sequential(
+                        nn.Linear(local_input_dim, self.hidden_dim),
+                        create_activation(self.activation_name),
+                        nn.Linear(self.hidden_dim, local_out_dim),
+                        create_activation(self.activation_name),
+                    )
+
+                channel_module.append(sequential)
+
+            self.ACM_convs.append(
+                ACM_GINEConv(
+                    nn_lowpass=self.nns_lowpass[i],
+                    nn_highpass=self.nns_highpass[i],
+                    nn_fullpass=self.nns_fullpass[i],
+                    nn_lowpass_proj=self.nns_lowpass_proj[i],
+                    nn_highpass_proj=self.nns_highpass_proj[i],
+                    nn_fullpass_proj=self.nns_fullpass_proj[i],
+                    nn_mix=self.nns_mix[i],
+                )
+            )
+
+            if self.norm_type == "layer":
+                self.layer_norms.append(nn.LayerNorm(local_out_dim))
+            elif self.norm_type == "graph":
+                self.layer_norms.append(GraphNorm(local_out_dim))
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.reset_parameters()
+            elif isinstance(m, (nn.BatchNorm1d, nn.LayerNorm, GraphNorm)):
+                m.reset_parameters()
+
+    def forward(self, x, edge_index, edge_attr, batch=None, return_hidden=False):
+        """Forward pass through all ACM-GINEConv layers.
+
+        Args:
+            x: Node features [N, in_dim].
+            edge_index: Edge indices [2, E].
+            edge_attr: Edge features [E, edge_in_dim]. The first column
+                (index 0) is the scalar spatial distance used for degree
+                normalization.
+            batch: Node->graph assignment [N]; required only when
+                norm_type == "graph" (GraphNorm normalizes per graph).
+            return_hidden: If True, also return all intermediate node states.
+
+        Returns:
+            x: Final node embeddings [N, out_dim].
+            outs: (optional) List of node states after each layer.
+        """
+        # Extract scalar spatial distance for degree normalization
+        edge_weight = edge_attr[:, 0]
+
+        # Project node and edge features into hidden_dim ONCE
+        x = self.node_input_proj(x)
+        if self.input_norm is not None:
+            x = self.input_norm(x)  # learnable input-embedding normalization
+        # Distance (col 0) stays as edge_weight only when excluded from the proj.
+        edge_proj_input = edge_attr if self.edge_distance_in_proj else edge_attr[:, 1:]
+        edge_feat = self.edge_input_proj(edge_proj_input)  # Static: reused every layer
+
+        outs = []
+        for i in range(self.num_layers):
+            # GINEConv-style message passing with ACM filtering
+            # Edge features are static — no per-layer edge update
+            x = self.ACM_convs[i](
+                x=x,
+                edge_index=edge_index,
+                edge_weight=edge_weight,
+                edge_feat=edge_feat,
+            )
+            if self.layer_norms is not None:
+                if self.norm_type == "graph":
+                    x = self.layer_norms[i](x, batch)
+                else:
+                    x = self.layer_norms[i](x)
+            outs.append(x)
+
+        if return_hidden:
+            return x, outs
+        else:
+            return x
+
+
+if __name__ == "__main__":
+    model = ACM_GINEConv_model(46, 46, 2, 256, 74, True)
+    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"ACM_GINEConv_model parameters: {num_params:,}")

models/edcoder.py

@@ -0,0 +1,333 @@
+from typing import Optional
+from itertools import chain
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .acm_gin import ACM_GIN_model
+from .acm_gineconv import ACM_GINEConv_model
+from torch_geometric.utils import dropout_edge
+from torch_geometric.utils import add_self_loops
+from torch_geometric.nn import global_mean_pool
+
+
+def sce_loss(x, y, alpha=3):
+    x = F.normalize(x, p=2, dim=-1)
+    y = F.normalize(y, p=2, dim=-1)
+
+    loss = (1 - (x * y).sum(dim=-1)).pow_(alpha)
+    loss = loss.mean()
+
+    return loss
+
+
+def setup_module(
+    m_type,
+    in_dim,
+    out_dim,
+    num_hidden,
+    num_layers,
+    activation,
+    batchnorm,
+    edge_in_dim,
+    norm_type="none",
+    input_norm="none",
+    edge_distance_in_proj=True,
+) -> nn.Module:
+
+    if m_type == "acm_gin":
+        mod = ACM_GIN_model(
+            int(in_dim),
+            int(out_dim),
+            num_layers,
+            int(num_hidden),
+            edge_in_dim,
+            batchnorm,
+            activation=activation,
+        )
+    elif m_type == "acm_gineconv":
+        mod = ACM_GINEConv_model(
+            int(in_dim),
+            int(out_dim),
+            num_layers,
+            int(num_hidden),
+            edge_in_dim,
+            batchnorm,
+            activation=activation,
+            norm_type=norm_type,
+            input_norm=input_norm,
+            edge_distance_in_proj=edge_distance_in_proj,
+        )
+    else:
+        raise NotImplementedError
+
+    return mod
+
+
+class PreModel(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        edge_in_dim: int,  # Dimensionality of raw edge features (e.g. 74)
+        num_hidden: int,
+        num_layers: int,
+        nhead: int,  # TODO: Check if nhead is actually needed; currently unused (vestigial from GAT?)
+        nhead_out: int,  # TODO: Check if nhead_out is actually needed; currently unused (vestigial from GAT?)
+        activation: str,
+        feat_drop: float,  # TODO: Check if feat_drop is actually needed; currently unused
+        attn_drop: float,  # TODO: Check if attn_drop is actually needed; currently unused (vestigial from GAT?)
+        negative_slope: float,  # TODO: Check if negative_slope is actually needed; currently unused (vestigial from GAT?)
+        residual: bool,  # TODO: Check if residual is actually needed; currently unused
+        norm: Optional[str],  # TODO: Check if norm is actually needed; currently unused
+        mask_rate: float = 0.3,
+        encoder_type: str = "gat",
+        decoder_type: str = "gat",
+        loss_fn: str = "sce",
+        drop_edge_rate: float = 0.0,
+        replace_rate: float = 0.1,
+        alpha_l: float = 2,
+        concat_hidden: bool = False,
+        batchnorm=False,
+        encoder_norm="none",
+        encoder_input_norm="none",
+        edge_distance_in_proj=True,
+        vicreg_var_weight: float = 0.0,
+        vicreg_cov_weight: float = 0.0,
+        vicreg_gamma: float = 1.0,
+    ):
+        super(PreModel, self).__init__()
+        self._mask_rate = mask_rate
+        self._encoder_type = encoder_type
+        self._decoder_type = decoder_type
+        self._drop_edge_rate = drop_edge_rate
+        self._output_hidden_size = num_hidden
+        self._concat_hidden = concat_hidden
+
+        # VICReg-style anti-collapse regularizer (0 = off); see _vicreg_terms.
+        self._vicreg_var_weight = vicreg_var_weight
+        self._vicreg_cov_weight = vicreg_cov_weight
+        self._vicreg_gamma = vicreg_gamma
+        self.last_loss_components = {}
+
+        self._replace_rate = replace_rate
+        self._mask_token_rate = 1 - self._replace_rate
+
+        assert num_hidden % nhead == 0
+        assert num_hidden % nhead_out == 0
+
+        enc_num_hidden = num_hidden
+        enc_nhead = 1
+
+        dec_in_dim = num_hidden
+        dec_num_hidden = num_hidden
+
+        # Build encoder
+        self.encoder = setup_module(
+            m_type=encoder_type,
+            in_dim=in_dim,
+            out_dim=enc_num_hidden,
+            num_hidden=enc_num_hidden,
+            num_layers=num_layers,
+            activation=activation,
+            batchnorm=batchnorm,
+            edge_in_dim=edge_in_dim,
+            norm_type=encoder_norm,
+            input_norm=encoder_input_norm,
+            edge_distance_in_proj=edge_distance_in_proj,
+        )
+
+        # Build decoder for attribute prediction
+        self.decoder = setup_module(
+            m_type=decoder_type,
+            in_dim=dec_in_dim,
+            out_dim=in_dim,
+            num_hidden=dec_num_hidden,
+            num_layers=1,
+            activation=activation,
+            batchnorm=batchnorm,
+            edge_in_dim=edge_in_dim,
+            edge_distance_in_proj=edge_distance_in_proj,
+        )
+
+        self.enc_mask_token = nn.Parameter(torch.zeros(1, in_dim))
+        if concat_hidden:
+            self.encoder_to_decoder = nn.Linear(
+                dec_in_dim * num_layers, dec_in_dim, bias=False
+            )
+        else:
+            self.encoder_to_decoder = nn.Linear(dec_in_dim, dec_in_dim, bias=False)
+
+        # Setup loss function
+        self.criterion = self.setup_loss_fn(loss_fn, alpha_l)
+
+    @property
+    def output_hidden_dim(self):
+        return self._output_hidden_size
+
+    def setup_loss_fn(self, loss_fn, alpha_l):
+        if loss_fn == "mse":
+            criterion = nn.MSELoss()
+        elif loss_fn == "sce":
+            criterion = partial(sce_loss, alpha=alpha_l)
+        else:
+            raise NotImplementedError
+        return criterion
+
+    def encoding_mask_noise(self, x, mask_rate=0.3, virtual_node_index=None):
+        num_nodes = x.shape[0]
+        all_indices = torch.arange(num_nodes, device=x.device)
+
+        # Remove virtual node index from masking candidates
+        if virtual_node_index is not None:
+            all_indices = all_indices[~torch.isin(all_indices, virtual_node_index)]
+
+        perm = all_indices[torch.randperm(len(all_indices), device=x.device)]
+
+        # random masking
+        num_mask_nodes = int(mask_rate * len(perm))
+        mask_nodes = perm[:num_mask_nodes]
+        keep_nodes = perm[num_mask_nodes:]
+
+        out_x = x.clone()
+
+        if self._replace_rate > 0:
+            num_noise_nodes = int(self._replace_rate * num_mask_nodes)
+            perm_mask = torch.randperm(num_mask_nodes, device=x.device)
+            token_nodes = mask_nodes[
+                perm_mask[: int(self._mask_token_rate * num_mask_nodes)]
+            ]
+            noise_nodes = mask_nodes[
+                perm_mask[-int(self._replace_rate * num_mask_nodes) :]
+            ]
+            noise_to_be_chosen = torch.randperm(len(perm), device=x.device)[
+                :num_noise_nodes
+            ]
+            noise_to_be_chosen = all_indices[noise_to_be_chosen]
+
+            out_x[token_nodes] = 0.0
+            out_x[noise_nodes] = x[noise_to_be_chosen]
+        else:
+            token_nodes = mask_nodes
+            out_x[mask_nodes] = 0.0
+
+        out_x[token_nodes] += self.enc_mask_token
+
+        return out_x, (mask_nodes, keep_nodes)
+
+    def forward(self, batch):
+        # ---- attribute reconstruction ----
+        x, edge_index, edge_attr, virtual_node_index, batch = (
+            batch.x,
+            batch.edge_index,
+            batch.edge_attr,
+            getattr(batch, "virtual_node_index", None),
+            batch.batch,
+        )
+        loss = self.mask_attr_prediction(
+            x, edge_index, edge_attr, batch, virtual_node_index
+        )
+        return loss
+
+    def mask_attr_prediction(self, x, edge_index, edge_attr, batch, virtual_node_index):
+
+        use_x, (mask_nodes, keep_nodes) = self.encoding_mask_noise(
+            x,
+            self._mask_rate,
+            virtual_node_index,
+        )
+
+        if self._drop_edge_rate > 0:
+            use_edge_index, masked_edges = dropout_edge(
+                edge_index, self._drop_edge_rate
+            )
+            use_edge_attr = edge_attr[masked_edges]
+            use_edge_index, use_edge_attr = add_self_loops(
+                use_edge_index, use_edge_attr, fill_value="min"
+            )
+        else:
+            use_edge_index = edge_index
+            use_edge_attr = edge_attr
+
+        enc_rep, all_hidden = self.encoder(
+            use_x, use_edge_index, use_edge_attr, batch=batch, return_hidden=True
+        )
+        if self._concat_hidden:
+            enc_rep = torch.cat(all_hidden, dim=1)
+
+        # ---- attribute reconstruction ----
+        rep = self.encoder_to_decoder(enc_rep)
+
+        # VICReg anti-collapse regularization on the graph-level pooled embedding,
+        # computed BEFORE the decoder remask below, and only while training so that
+        # val_loss stays pure reconstruction and the collapse metric is independent.
+        if self.training and (
+            self._vicreg_var_weight > 0 or self._vicreg_cov_weight > 0
+        ):
+            reg_loss, reg_comps = self._vicreg_terms(rep, batch)
+        else:
+            reg_loss, reg_comps = rep.new_zeros(()), {}
+
+        if self._decoder_type not in ("mlp", "linear"):
+            # * remask, re-mask
+            rep[mask_nodes] = 0
+
+        if self._decoder_type in ("mlp", "linear"):
+            recon = self.decoder(rep)
+        else:
+            recon = self.decoder(rep, use_edge_index, use_edge_attr)
+
+        x_init = x[mask_nodes]
+        x_rec = recon[mask_nodes]
+
+        loss = self.criterion(x_rec, x_init)
+        self.last_loss_components = {"sce": float(loss.detach()), **reg_comps}
+
+        return loss + reg_loss
+
+    def _vicreg_terms(self, rep, batch):
+        """VICReg-style anti-collapse terms on the graph-level pooled embedding.
+
+        Returns ``(reg_loss, components)``. Combines a variance hinge (keep each
+        embedding dim's per-batch std >= gamma) and a covariance penalty
+        (decorrelate dims). Applied on ``global_mean_pool(rep)`` -- the same
+        representation the PCA-collapse metric is computed on -- so it directly
+        opposes the dimensional collapse observed once reconstruction saturates.
+        """
+        z = global_mean_pool(rep, batch)  # [G, D]
+        G = z.size(0)
+        D = z.size(1)
+        reg = rep.new_zeros(())
+        comps = {}
+        if self._vicreg_var_weight > 0:
+            std = torch.sqrt(z.var(dim=0, unbiased=False) + 1e-4)
+            var_term = torch.clamp(self._vicreg_gamma - std, min=0).mean()
+            reg = reg + self._vicreg_var_weight * var_term
+            comps["vic_var"] = float(var_term.detach())
+        if self._vicreg_cov_weight > 0 and G > 1:
+            zc = z - z.mean(dim=0, keepdim=True)
+            cov = (zc.t() @ zc) / (G - 1)  # [D, D]
+            off_diag_sq = cov.pow(2).sum() - cov.diagonal().pow(2).sum()
+            cov_term = off_diag_sq / D
+            reg = reg + self._vicreg_cov_weight * cov_term
+            comps["vic_cov"] = float(cov_term.detach())
+        return reg, comps
+
+    def embed(self, x, edge_index, edge_attr, batch):
+        if self._concat_hidden:
+            enc_rep, all_hidden = self.encoder(
+                x, edge_index, edge_attr, batch=batch, return_hidden=True
+            )
+            enc_rep = torch.cat(all_hidden, dim=1)
+        else:
+            enc_rep = self.encoder(x, edge_index, edge_attr, batch=batch)
+        rep = self.encoder_to_decoder(enc_rep)
+        return rep
+
+    @property
+    def enc_params(self):
+        return self.encoder.parameters()
+
+    @property
+    def dec_params(self):
+        return chain(*[self.encoder_to_decoder.parameters(), self.decoder.parameters()])

models/utils.py

@@ -0,0 +1,75 @@
+import torch.nn as nn
+import torch
+from functools import partial
+
+
+def create_activation(name):
+    if name == "relu":
+        return nn.ReLU()
+    elif name == "gelu":
+        return nn.GELU()
+    elif name == "prelu":
+        return nn.PReLU()
+    elif name is None:
+        return nn.Identity()
+    elif name == "elu":
+        return nn.ELU()
+    else:
+        raise NotImplementedError(f"{name} is not implemented.")
+
+
+def create_norm(name):
+    if name == "layernorm":
+        return nn.LayerNorm
+    elif name == "batchnorm":
+        return nn.BatchNorm1d
+    elif name == "graphnorm":
+        return partial(NormLayer, norm_type="groupnorm")
+    else:
+        return nn.Identity
+
+
+class NormLayer(nn.Module):
+    def __init__(self, hidden_dim, norm_type):
+        super().__init__()
+        if norm_type == "batchnorm":
+            self.norm = nn.BatchNorm1d(hidden_dim)
+        elif norm_type == "layernorm":
+            self.norm = nn.LayerNorm(hidden_dim)
+        elif norm_type == "graphnorm":
+            self.norm = norm_type
+            self.weight = nn.Parameter(torch.ones(hidden_dim))
+            self.bias = nn.Parameter(torch.zeros(hidden_dim))
+
+            self.mean_scale = nn.Parameter(torch.ones(hidden_dim))
+        else:
+            raise NotImplementedError
+
+    def forward(self, graph, x):
+        tensor = x
+        if self.norm is not None and type(self.norm) != str:
+            return self.norm(tensor)
+        elif self.norm is None:
+            return tensor
+
+        batch_list = graph.batch_num_nodes
+        batch_size = len(batch_list)
+        batch_list = torch.Tensor(batch_list).long().to(tensor.device)
+        batch_index = (
+            torch.arange(batch_size).to(tensor.device).repeat_interleave(batch_list)
+        )
+        batch_index = batch_index.view((-1,) + (1,) * (tensor.dim() - 1)).expand_as(
+            tensor
+        )
+        mean = torch.zeros(batch_size, *tensor.shape[1:]).to(tensor.device)
+        mean = mean.scatter_add_(0, batch_index, tensor)
+        mean = (mean.T / batch_list).T
+        mean = mean.repeat_interleave(batch_list, dim=0)
+
+        sub = tensor - mean * self.mean_scale
+
+        std = torch.zeros(batch_size, *tensor.shape[1:]).to(tensor.device)
+        std = std.scatter_add_(0, batch_index, sub.pow(2))
+        std = ((std.T / batch_list).T + 1e-6).sqrt()
+        std = std.repeat_interleave(batch_list, dim=0)
+        return self.weight * sub / std + self.bias