File size: 4,825 Bytes
d538c61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | ---
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.
|