| --- |
| license: cc-by-4.0 |
| library_name: pytorch |
| pipeline_tag: feature-extraction |
| tags: |
| - node2vec |
| - graph-embedding |
| - metabolomics |
| - pytorch-geometric |
| - metabolights |
| - metabolomics-workbench |
| metrics: |
| - roc_auc |
| --- |
| |
| # Node2Vec embeddings for cross-study metabolomics co-response graphs |
|
|
| 128-dimensional Node2Vec embeddings for two undirected graphs of metabolite co-response |
| edges that pass a chi-square validity filter (every expected frequency ≥ 5). Held-out |
| link prediction reaches **AUC 0.932** and **0.928**, against degree-only baselines of |
| 0.826 and 0.888. |
|
|
| The graphs, node properties and full pipeline are in the companion dataset repository: |
| [kozo2/metabolomics-edges-expected-ge5](https://huggingface.co/datasets/kozo2/metabolomics-edges-expected-ge5). |
|
|
| ## Files |
|
|
| | File | Contents | Size | |
| |---|---|---| |
| | `edge_ML_filtered_expected_ge5_n2v.pt` | `{embedding: [848, 128] float32, node_id: [848], args: {...}}` | 453 KB | |
| | `edge_MLvsMW_filtered_expected_ge5_n2v.pt` | `{embedding: [15758, 128] float32, node_id: [15758], args: {...}}` | 8.6 MB | |
| | `node2vec_model.py`, `config.py` | Model definition, training loop, embedding export | — | |
| | `heldout_check.py` | The held-out link-prediction evaluation reported below | — | |
|
|
| Each file records the exact hyperparameters that produced it under `args`, and |
| `node_id[i]` is the original study/assay/feature ID of embedding row `i`. |
|
|
| ## Using it |
|
|
| ```python |
| import torch |
| |
| ck = torch.load("edge_MLvsMW_filtered_expected_ge5_n2v.pt", weights_only=False) |
| z = torch.nn.functional.normalize(ck["embedding"], dim=1) # cosine space |
| idx = {node: i for i, node in enumerate(ck["node_id"])} |
| |
| q = z[idx["MTBLS311_0001_00001537"]] |
| top = (z @ q).topk(11).indices[1:] # drop self |
| print([ck["node_id"][i] for i in top]) |
| ``` |
|
|
| Scores are cosine similarities — that is the metric the embeddings were trained under |
| random walks and evaluated with, so use normalized vectors rather than raw dot products. |
|
|
| ## Training |
|
|
| | | `edge_ML` | `edge_MLvsMW` | |
| |---|---|---| |
| | graph | 848 nodes, 3,697 edges | 15,758 nodes, 89,277 edges | |
| | parameters | 108,544 | 2,017,024 | |
| | `embedding_dim` | 128 | 128 | |
| | `walk_length` / `context_size` | 20 / 10 | 20 / 10 | |
| | `walks_per_node` / negatives | 10 / 1 | 10 / 1 | |
| | `p` / `q` | 1.0 / 1.0 (unbiased) | 1.0 / 1.0 (unbiased) | |
| | `batch_size` / `lr` / optimizer | 128 / 0.01 / SparseAdam | 128 / 0.01 / SparseAdam | |
| | epochs / wall time | 200 / 31 s | 200 / 3 m 05 s | |
| | loss, first → last | 9.02 → 0.815 | 6.73 → 0.876 | |
|
|
| Trained on one NVIDIA H100 NVL. `sparse=True` makes the embedding table the only |
| parameter tensor, which is why the optimizer is `SparseAdam`. |
|
|
| ## Evaluation |
|
|
| `heldout_check.py` deduplicates undirected edges, holds out a random 10%, retrains from |
| scratch on the remaining 90%, then scores held-out edges against an equal number of |
| sampled non-edges by cosine similarity. The control is the preferential-attachment |
| baseline `d_u · d_v` on training degrees, which answers whether the embedding learned |
| anything beyond "popular nodes connect". |
|
|
| | | `edge_ML` | `edge_MLvsMW` | |
| |---|---|---| |
| | train / held-out edges | 3,328 / 369 | 80,350 / 8,927 | |
| | nodes isolated by the split | 42 | 413 | |
| | **held-out AUC (never seen)** | **0.932** | **0.928** | |
| | train AUC (in-sample) | 0.985 | 0.985 | |
| | degree baseline (held-out) | 0.826 | 0.888 | |
| | margin over baseline | +0.106 | +0.040 | |
|
|
| **Species neighbour purity** — of each node's 10 nearest embeddings by cosine, the |
| fraction sharing its species label — is 48.8% overall for `edge_ML` and 59.2% for |
| `edge_MLvsMW`, rising to 75–79% for the dominant host species. For the bipartite |
| `edge_MLvsMW` graph, the same measure on *source database* gives 67.5% overall (75.4% ST, |
| 47.7% MTBLS). |
|
|
| ## Caveats |
|
|
| - **Training budget matters more than it looks.** At 20 epochs, `edge_ML` scored a |
| held-out AUC of 0.791 — *below* its own 0.826 degree baseline. The embeddings only beat |
| preferential attachment once training converges. `--epochs` defaults to 200 for this |
| reason; do not shorten it without re-running `heldout_check.py`. |
| - On `edge_MLvsMW` the margin over the degree baseline is thin (+0.040). That graph is |
| bipartite with a few very high-degree studies, so much of its link structure is |
| explained by degree alone. |
| - `edge_MLvsMW` loss is flat from about epoch 19, so 200 epochs is more budget than that |
| graph needs. |
| - Embeddings are transductive: there is no way to embed a node that was not in the |
| training graph without retraining. |
|
|
