README.md

---
license: cc-by-4.0
library_name: pytorch
pipeline_tag: feature-extraction
tags:
- node2vec
- graph-embedding
- metabolomics
- pytorch-geometric
- metabolights
metrics:
- roc_auc
---

# Node2Vec embeddings for the edge_ML metabolomics graph

128-dimensional Node2Vec embeddings for the 18,494 nodes of an undirected metabolite
co-response graph with 2,709,209 edges. Held-out link prediction reaches **AUC 0.988**,
against 0.920 for a degree-only baseline.

The graph, node properties and full pipeline are in the companion dataset repository.

## Files

| File | Contents | Size |
|---|---|---|
| `edge_ML_expected_ge5_n2v.pt` | `{embedding: [18494, 128] float32, node_id: [18494], args: {...}}` | 9.7 MB |
| `node2vec_model.py` | Model definition, training loop, embedding export | — |

## Using it

```python
import torch

ck = torch.load("edge_ML_expected_ge5_n2v.pt", weights_only=False)
z = ck["embedding"]                      # [18494, 128] float32
index = {nid: i for i, nid in enumerate(ck["node_id"])}
v = z[index["MTBLS1405_0002_00003332"]]  # one node's vector
```

`node_id[i]` is the original string ID for row `i`; the order is lexicographic over the
union of the graph's two endpoint columns, matching the dataset's graph object. Scores
were computed with **cosine** similarity, which is also the metric to use downstream.

## Training

| Parameter | Value |
|---|---|
| `embedding_dim` | 128 |
| `walk_length` | 20 |
| `context_size` | 10 |
| `walks_per_node` | 10 |
| `num_negative_samples` | 1 |
| `p`, `q` | 1.0, 1.0 (unbiased walks) |
| Batch size | 128 seed nodes, 145 batches per epoch |
| Optimiser | `SparseAdam`, lr 0.01 |
| Epochs | 20 |
| Parameters | 2,367,232 (18,494 × 128) |

Loss fell from 9.92 at initialisation to 0.880, flat from about epoch 14, at roughly
0.9 s/epoch on one H100. `sparse=True` on the model is what allows `SparseAdam`;
changing either requires changing the other.

```bash
uv run python node2vec_model.py --epochs 20
```

`Node2Vec` requires `pyg-lib >= 0.6.0` for its random-walk kernel, which is not on PyPI;
the dataset repository's `pyproject.toml` pins `pyg-lib` 0.9.0+pt214cu130 from
`data.pyg.org`.

## Evaluation

200,000 sampled positive edges against 200,000 non-edges verified absent from the full
edge set, scored by cosine similarity, AUC by the Mann-Whitney rank identity.

| Model | Scored edges | Score | AUC |
|---|---|---|---|
| 90/10 retrain | Held-out 10%, never seen | cosine | **0.9880** |
| 90/10 retrain | Its own training edges | cosine | 0.9892 |
| 90/10 retrain | Held-out 10%, never seen | degree product `d_u × d_v` | 0.9201 |
| Full graph (this release) | Its own training edges | cosine | 0.9892 |
| Full graph (this release) | Its own training edges | dot product | 0.9868 |

The held-out row is the one that matters: a second model was trained from scratch on 90%
of the edges and scored on the 10% it never saw. Held-out 0.9880 against in-sample 0.9892
is a gap of 0.001, so the model learns graph structure rather than memorising pairs. The
degree baseline matters because the graph is dense (median degree 90) — a high AUC that
merely reproduced the degree distribution would carry little information.

Other checks on the released embeddings:

- **Neighbourhood recovery** — of each node's 10 nearest embeddings, 49.9% are true graph
  neighbours against 1.6% expected by chance (31.7×); at top-50, 41.3% (26.2×).
- **Embedding health** — all finite; L2 norms 0.94 / 1.92 / 9.49 (min / median / max);
  per-dimension standard deviation 0.15–0.28, so no dead dimensions; mean cosine over
  200,000 random pairs is 0.0038, ruling out collapse.
- **Species purity** — 90.7% of all nodes have ten nearest embeddings sharing their
  species, rising above 98% for the three largest species and falling to 68–79% for
  species with a few hundred nodes.

## Limitations

- **Topology only.** The walks are unweighted, so neither the graph's `edge_attr`
  (`OddsRatio_log2`, `ChiTestsPValue`) nor its node features `x` influence these
  embeddings. Letting association strength steer the walks needs a weighted sampler or a
  pre-thresholded edge set; using the node features needs a message-passing model.
- **Transductive.** Node2Vec learns one vector per node in a fixed graph. There is no
  way to embed a node that was not present at training time.
- **Species and study are entangled.** Edges form mostly within a study and a study is
  normally one species, so the clean species separation partly reflects how the graph was
  assembled, not an independent biological signal.
- In the 90/10 evaluation split, 97 low-degree nodes were left isolated in the training
  graph and their vectors stay near initialisation. That affects only the held-out
  experiment; the released full-graph model has no isolated nodes.