AugMent — contrastive alignment transfers proteomic predictive signal into metabolomics-only cohorts

Plasma proteomics predicts disease well but is measured in few people. NMR metabolomics is cheap and measured in everyone. AugMent learns the metabolome→proteome alignment on the ~35,000 UK Biobank participants who have both, then applies the frozen encoder to cohorts that have only metabolomics — improving disease risk prediction in populations that were never assayed for proteins.

This is a transfer-learning result: paired samples are the source, metabolomics-only cohorts are the target, and the transferred representation costs nothing extra to measure.

Code for Contrastive alignment transfers proteomic predictive signals to metabolomics data (Hu et al.). Results, supplementary tables and figures: Zenodo record 21861685.

---

The central claim

The frozen CLIP encoder, trained only on paired UKB data, was applied to 442,559 UKB participants with metabolomics but no proteomics (test set N = 88,512, 597 analyzable diseases). Risk prediction improved significantly (q < 0.05) for 34 diseases — 14 inside the proteomics-advantaged set and 20 outside it — while overall performance across the remaining spectrum was preserved.

Setting	Cohort	Improved	Significant (q<0.05)	Significantly worse	Median ΔC-index
Proteomics-advantaged diseases (n=88)	UKB metabolomics-only (N=88,512)	68	14	0	0.011 (max 0.083, SLE)
Remaining disease spectrum (n=509)	UKB metabolomics-only	295	20	3	0.010 (max 0.159, toxic liver dis.)
External cohort, 24 phenotypes	Human Phenotype Project (N=1,442)	—	9 (ROC-AUC, q<0.05)	6	up to ΔROC-AUC 0.133

On the 88 proteomics-advantaged diseases no disease significantly favored raw metabolomics on any of the three metrics (C-index, detection rate, likelihood ratio). Gains concentrated in renal, inflammatory and respiratory conditions: chronic nephritic syndrome (C-index 0.89, Δ 0.067), Crohn's disease (0.68, Δ 0.052), hypertensive renal disease (0.91, Δ 0.044).

Transfer holds across biobanks: an alignment learned entirely in UKB improved cross-sectional prediction in an Israeli cohort profiled on the same NMR platform — and across a change of biofluid, EDTA plasma to serum, without retraining. Largest ROC-AUC gains: migraines, gallstones, hypertension, allergic rhinitis (Δ 0.112–0.133).

Secondary findings

1. Alignment beats reconstruction. Of five architectures (ridge, MLP, Transformer, FedCoder, CLIP), only the two embedding-alignment models improved over raw metabolomics across the full disease spectrum. The three translation models did not; the Transformer, the most expressive, was the worst. On the paired test set (N = 8,706) CLIP improved 70 of the 88 proteomics-advantaged diseases, median ΔC-index 0.029, and discriminated more diseases at C-index > 0.75 than any other model in both regimes (18 vs. 12 for raw metabolomics; 47 vs. 39 on the remaining set).
2. Effective transfer does not require per-protein recovery. CLIP recovers only 135 of 2,923 proteins at r > 0.3 — far fewer than the Transformer (1,923) — yet predicts best, and the reconstruction models recover those same 135 proteins comparably or better. A temperature ablation confirms the objective drives the trade-off: weakening the contrastive term raises recovery (median r 0.07 → 0.25) and lowers prediction. The recovered minority is canonical lipoprotein machinery (LDLR, PCSK9, ANGPTL3/4, LCAT, PLTP, APOA1, APOC1, PLA2G7, PLA2G10), the biological interface the two modalities share.
3. The gains are not covariate information. Age, sex and BMI alone are a strong comparator (mean ΔC-index +0.037 and +0.064 over the residualized-metabolomics reference), but the augmented representation exceeds the covariates alone by 0.038 and 0.025 (rank-biserial r = 0.92 and 0.60, both P < 10⁻¹³). Both omics layers are residualized on the covariates before use, so that margin is not demographic signal reintroduced by the encoder.

---

Repository layout

Path	What it is
src/augment_omics/	The package — models, preprocessing, evaluation, CLI
src/augment_omics/model/	Alignment architectures (align.py), autoencoders (embed.py), layers (nn.py), datamodules (data.py)
src/augment_omics/prep/	DAE imputation (impute.py), covariate residualization (residualizer.py)
src/augment_omics/eval/	Cox model (cox_model.py), survival metrics (survival.py), linear probes (linear_eval.py)
src/augment_omics/configs/	Hydra configs per subcommand (align/, embed/, eval/, gen_data/)
scripts/	Analysis and figure generation — readsresults/data/, writes results/main/

results/ is not tracked in git (figures, per-disease metric CSVs and supplementary tables are large binaries). Download it from the Zenodo record and unpack it at the repo root as results/ before running the analysis scripts; the archive ships its own results/README.md file-by-file guide.

Environment

Python ≥ 3.11, managed with uv. Locked for Linux x86_64 and macOS arm64.

git clone https://github.com/RasmussenLab/AugMent && cd AugMent
uv sync          # installs from uv.lock, creates .venv
source .venv/bin/activate
augment --help   # CLI entry point

Core stack: PyTorch + PyTorch Lightning (alignment models), scikit-learn (ridge, logistic regression), torchsurv (Cox partial likelihood at biobank scale), fast-cindex (Numba-JIT Harrell's C, needed for 1,000-iteration bootstraps), SHAP, Hydra.

Models were trained on the UKB Research Analysis Platform. GPU is recommended for the alignment step; everything downstream runs on CPU.

Data

Neither cohort can be redistributed.

• UK Biobank (application 32683) — Olink Explore 3072 EDTA plasma proteomics (2,923 proteins) and Nightingale NMR metabolomics (170 absolute concentrations), under an approved UKB application. After preprocessing: 43,527 paired participants (26,115 train / 8,706 val / 8,706 test) and 442,559 metabolomics-only (354,047 train / 88,512 test).
• Human Phenotype Project — external validation, N = 1,442, Nightingale NMR serum metabolomics, via the HPP data access process. No proteomics exists in HPP.

To exercise the pipeline without either, generate synthetic data with matched dimensionality:

augment gen_data -- output_dir=data/synthetic m_features=170 p_features=2923 n_samples=60000

Reproducing the pipeline

The CLI is Hydra-backed: augment <command> --config-name <name> -- key=value ... (note the -- before overrides).

Step	Command	Notes
1. Impute + residualize	augment_omics.prep.impute, .residualizer	DAE imputation (20% masking, masked MSE), then OLS residualization against age, sex, BMI
2. Train alignment	augment align --config-name clip	Alsoridge, translation (MLP/Transformer), fedcoder. Configs in configs/align/
3. Encode a new cohort	augment infer --model clip --checkpoint <ckpt> --omics-data <csv> --save-path <out>	Frozen encoder → proteome-augmented representations
4. Per-disease Cox + metrics	augment eval --config-name default -- eval.task_types=[test_survival] ...	1,000-iteration bootstrap; writes the CSVs in results/data/survival_metrics/

Key hyperparameters (from Methods): CLIP uses hidden dims [256, 1024, 2048], lr 1e-4, 90 epochs, learnable inverse temperature initialized at 1/τ₀ (τ₀ = 0.01, initial multiplier 100) and clamped at 100; the temperature ablation re-initializes at τ₀ = 0.1, 1, 10, 100 with everything else held fixed. Cox models are a single linear layer, lr 1e-3, ≤20 epochs, early stopping patience 5, Efron ties, partial likelihood evaluated within mini-batches. Disease endpoints are three-character ICD-10 codes with ≥80 incident cases, 10-year follow-up cap, prevalent cases and 180-day incident cases excluded.

Regenerating the figures

Every analysis script reads the numeric outputs under results/data/ and writes into results/main/ — no cohort access required once the Zenodo archive is unpacked. Run from the repo root:

python scripts/model_selection.py                        # Fig. 1c-h: architecture benchmark + ablation
python scripts/threshold_sensitivity.py                  # C-index threshold sweep (0.65-0.85), CLIP vs FedCoder
python scripts/unpaired_benchmark.py                     # Fig. 2 & 3: transfer into metabolomics-only UKB
python scripts/unpaired_covariate_ablation.py            # Fig. 4g,h: covariate specifications
python scripts/unpaired_model_delta_vs_metabolomics.py   # per-model deltas on the metabolomics-only cohort
python scripts/hpp_benchmark.py                          # Fig. 5: external validation in HPP
python scripts/mechanisms.py                             # Fig. 4a-d: per-protein recovery
python scripts/enrichment_figure.py                      # Fig. 4e: GO enrichment of recovered proteins
python scripts/shap_plot.py                              # Fig. 4f-h: per-disease SHAP
python scripts/cindex_stats.py                           # results/main/models_stats.txt
python scripts/gain_vs_baseline.py                       # gain vs. baseline-difficulty summary
python scripts/build_supplementary_tables.py             # Supplementary_Tables.xlsx

Input and output paths are set at the bottom of each script. Two scripts take arguments instead: compute_protein_correlations.py and knn_overlap.py need model checkpoints, passed on the command line.

Conventions used throughout

• Proteomics-advantaged set — 88 diseases where a Cox model on raw proteomics beats one on raw metabolomics by ΔC-index > 0.01, defined on the paired test set (349 diseases). Remaining set — the other 261 (paired) or 509 (population-scale, including 248 diseases that only reach the case threshold at population scale).
• Metrics — C-index (Harrell's), DR (detection rate at FPR ≤ 10%), LR (= DR/FPR). All bootstrapped 1,000×; Benjamini-Hochberg applied within each disease set as a separate family.
• _with_cov / _no_cov — whether age, sex, BMI enter the Cox model.
• In the metrics CSVs, _1_ is the model arm, _2_ the baseline, _diff_ = arm 1 − arm 2.

Limitations

Effect sizes are modest (median ΔC-index ~0.01) — meaningful for population-scale stratification, not for individual patient decisions. The count of significant diseases is conservative by construction: with 1,000 bootstrap resamples the smallest attainable two-sided p-value is 0.002, so under Benjamini-Hochberg no disease can be called significant in the 509-disease family unless ≥21 reach that floor together; significance is collective rather than individual at this resolution. HPP tests whether the alignment transfers, not whether the UKB effect sizes replicate — it is cross-sectional, rests on 11–61 cases per disease, and resamples one cohort 50×. The transfer is bounded by biology: proteomic signals with no metabolomic reflection cannot cross the alignment, and phenotypes defined by a directly measured metabolite (hypercholesterolemia, pre-diabetes) are diluted by projection into proteomic space. Both cohorts are predominantly of European ancestry.

Citation

Hu, D., Rohrer, C., Pielies Avellí, M., Merino, J., Jensen, L. J., Rasmussen, S. Contrastive alignment transfers proteomic predictive signals to metabolomics data.

License

MIT — see LICENSE.