Add model card: VICReg joint embedding with honest negative results

README.md

@@ -1,121 +1,117 @@
 ---
-language: en
 license: apache-2.0
 tags:
-  - marine-biology
-  - metagenomics
-  - environmental-modeling
-  - protein-domains
-  - tara-oceans
-  - vicreg
-  - joint-embedding
-  - self-supervised-learning
-  - pytorch
-library_name: pytorch
+- self-supervised-learning
+- vicreg
+- joint-embedding
+- marine-biology
+- metagenomics
+- environmental-modeling
+- protein-domains
+- tara-oceans
+- pfam
+- pytorch
 pipeline_tag: tabular-regression
 ---
 
 # TARA-WorldModel-VICReg
 
-Joint environment--genome embedding model for marine ecosystem productivity prediction, trained with Variance-Invariance-Covariance Regularization (VICReg) on TARA Oceans data.
+Joint environment–proteome embedding model using VICReg (Variance-Invariance-Covariance Regularization) self-supervised learning, applied to the TARA Oceans metagenomic dataset. This model aligns environmental and Pfam protein domain representations in a shared latent space.
 
 ## Model Description
 
-The World Model learns a shared latent space that aligns environmental context (satellite-derived variables) with microalgal protein domain composition (Pfam module abundances), then predicts marine productivity (chlorophyll-a, POC, NFLH) from the joint embedding.
-
 ### Architecture
+Two-layer MLP encoder with VICReg alignment loss:
 
 ```
-Training:
-  env (24 dims) → EncoderE(128 → 32) → z_env ─┐
-                                                 ├→ VICReg loss
-  pfam (20 dims) → EncoderP(256 → 128 → 32) → z_pfam
-                                                 └→ Predictor(64 → 3) → productivity
-
-Inference (environment-only):
-  env → EncoderE → z_env → Predictor → [chl-a, POC, NFLH]
+Environment branch: Input(env_dim) → Linear(hidden) → ReLU → Dropout(0.3) → Linear(32)
+PFAM branch:        Input(pfam_dim) → Linear(hidden) → ReLU → Dropout(0.3) → Linear(32)
 ```
 
-- **EncoderE**: Linear(24→128) + BN + ReLU + Dropout(0.3) → Linear(128→32) + BN + ReLU + Dropout(0.3)
-- **EncoderP**: Linear(20→256) + BN + ReLU + Dropout(0.3) → Linear(256→128) + BN + ReLU + Dropout(0.3) → Linear(128→32) + BN + ReLU + Dropout(0.3)
-- **Predictor**: Linear(32→64) + ReLU → Linear(64→3)
-- **Total parameters**: 53,187
-
-### VICReg Loss
+- **Latent dimension:** 32
+- **Parameters:** ~53K–64K depending on PFAM input dimensionality
+- **VICReg loss weights:** variance = 25.0, invariance = 25.0, covariance = 1.0
+- **Prediction head alpha:** 1.0
 
-Non-contrastive self-supervised alignment (Bardes et al., ICLR 2022):
-- **Invariance**: MSE between co-located env/pfam embeddings (λ=25)
-- **Variance**: Hinge loss preventing embedding collapse (λ=25)
-- **Covariance**: Off-diagonal penalty decorrelating dimensions (λ=1)
-- **Prediction**: MSE on productivity targets (α=1)
+### Training Data
+- **Source:** 1,151 samples with complete productivity data (chlorophyll-a, POC, NFLH) from 1,810 total TARA Oceans samples
+- **Environmental features:** Google Earth Engine oceanographic variables
+- **PFAM features:** CLR-transformed domain abundances reduced via PCA to 20, 32, or 64 dimensions
 
 ## Performance
 
-Joint embedding improves POC prediction (R² 0.422 → 0.532, 26% relative improvement) over environment-only baseline. Chlorophyll-a and NFLH are better predicted by environment alone (directly satellite-measured).
+This model represents an exploratory methodological approach that did not outperform the primary XGBoost bidirectional framework used in the ELF-NET study. Results are reported for transparency and reproducibility.
+
+### 6-Fold Leave-One-Basin-Out (LOBO) CV
 
-## Files
+| Target | Joint Model R² | Env-Only Baseline R² | Cohen's d | p-value |
+|--------|----------------|----------------------|-----------|---------|
+| POC | 0.532 | 0.422 | 0.026 | 0.38 |
+| Chl-a | 0.516 | 0.561 | — | — |
+| NFLH | 0.560 | 0.700 | — | — |
 
-### Fold Checkpoints (leave-one-basin-out spatial CV)
+The POC improvement (+0.110 R²) was not statistically significant.
 
-Two training runs are provided:
-- `world_model_fold_*_20260127_110243.pt` — Initial configuration (latent_dim=16)
-- `world_model_fold_*_20260127_111754.pt` — Best configuration from hyperparameter sweep (latent_dim=32)
+### 9-Fold Spatial Block CV (matching primary XGBoost design)
 
-Six folds per run: Arctic, Atlantic, Indian, Mediterranean, Pacific, Southern.
+| PFAM dim | XGB Baseline R² | VICReg R² | ΔR² |
+|----------|-----------------|-----------|-----|
+| pfam20 | 0.417 | −2.045 | −2.462 |
+| pfam32 | 0.417 | −4.217 | −4.634 |
+| pfam64 | 0.417 | −1.262 | −1.679 |
 
-### Configuration
+The MLP architecture produced catastrophically negative R² on spatially distinctive held-out basins (Mediterranean, mid-Pacific), where distribution shift defeats shallow neural networks. XGBoost's tree-based partitioning handles this regime far more effectively at N ≈ 1,100.
 
-- `phase2_best_config.json` — Hyperparameter sweep results (54 configurations, 3 seeds each)
+### Interpretation
 
-## Hyperparameters (Best Config)
+The poor performance under spatial CV is driven by the **architecture confound** (MLP vs. XGBoost for small tabular data), not necessarily by absence of the PFAM alignment signal. A fair comparison would require XGBoost-with-VICReg-embeddings, which was not evaluated. The XGBoost bidirectional framework was retained as the primary modeling approach.
 
-| Parameter | Value |
-|-----------|-------|
-| latent_dim | 32 |
-| dropout | 0.3 |
-| λ_invariance | 25.0 |
-| λ_variance | 25.0 |
-| λ_covariance | 1.0 |
-| pred_alpha | 1.0 |
-| learning_rate | 0.001 |
-| weight_decay | 1e-4 |
-| batch_size | 128 |
-| max_epochs | 300 |
-| patience | 30 |
-| grad_clip | 1.0 |
+## Repository Contents
+
+- Model checkpoints (`.pt` files) for each fold and PFAM dimensionality
+- Hyperparameter sweep results
+- Per-fold training curves and metrics
 
 ## Usage
 
 ```python
 import torch
 
-# Load fold checkpoint
-ckpt = torch.load("world_model_fold_Atlantic_20260127_111754.pt", map_location="cpu")
-
-# ckpt contains model_state_dict for the full WorldModel
-# Requires WorldModel class from the training codebase
+# Load a VICReg checkpoint
+checkpoint = torch.load(
+    "path/to/vicreg_checkpoint.pt",
+    map_location="cpu",
+    weights_only=False
+)
+state_dict = checkpoint["model_state_dict"]
 ```
 
-## Dataset
+## Related Models
 
-- **1,810 ocean samples** with co-located environment and Pfam profiles
-- **24 environmental variables** (GEE oceanographic/atmospheric)
-- **20 Pfam module features** (aggregated from 9,466 domains via co-occurrence clustering)
-- **3 productivity targets** (chlorophyll-a, POC, NFLH)
-- **Spatial cross-validation**: Leave-one-basin-out (6 ocean basins)
+- [TARA-XGBoost-Bidirectional](https://huggingface.co/GreenGenomicsLab/TARA-XGBoost-Bidirectional) — Primary bidirectional models used in the ELF-NET study
+- [algaGPT](https://huggingface.co/GreenGenomicsLab/algaGPT) — Protein classification model used for proteome extraction
 
-## Related Models
+## References
+
+- Bardes, A., Ponce, J., & LeCun, Y. (2022). VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning. *ICLR 2022*.
+
+## Authors
 
-- [GreenGenomicsLab/algaGPT](https://huggingface.co/GreenGenomicsLab/algaGPT) — AlgaGPT protein classification
-- [GreenGenomicsLab/LA4SR-Pythia70m-b-ckpt-55000](https://huggingface.co/GreenGenomicsLab/LA4SR-Pythia70m-b-ckpt-55000) — LA4SR-Pythia classification
-- [GreenGenomicsLab/TARA-ELF-NET](https://huggingface.co/GreenGenomicsLab/TARA-ELF-NET) — Deep bidirectional env↔pfam models
-- [GreenGenomicsLab/TARA-XGBoost-Bidirectional](https://huggingface.co/GreenGenomicsLab/TARA-XGBoost-Bidirectional) — XGBoost bidirectional models
+David R. Nelson, Kourosh Salehi-Ashtiani
+
+Green Genomics Lab, New York University Abu Dhabi
 
 ## Citation
 
-LA4SR classification models:
-> Nelson DR, Jaiswal AK, Ismail NS, Mystikou A, Salehi-Ashtiani K. *Patterns*. 2024;6(11).
+```bibtex
+@article{elfnet2026,
+  title={ELF-NET: Environment-Linked Functional Network for marine microalgal domain-environment coupling},
+  author={Nelson, David R. and Salehi-Ashtiani, Kourosh},
+  journal={Forthcoming},
+  year={2026}
+}
+```
 
-## License
+## Contact
 
-Apache 2.0
+Kourosh Salehi-Ashtiani — ksa3@nyu.edu