SingleBicycle/cnn-mesh-xattn-fusion

CNN×Mesh Cross-Attention Fusion — w3b_xmod_ftnceq (best single model)

A 3D-CNN T1 autoencoder fused with a frozen cortical-surface GCN embedding at the 128-d bottleneck, trained to maximize the heritability of the learned embedding. This is the best single model from a 12-model lever study. Built so a collaborator can run inference to get a 128-d embedding per subject, then do loci counting / GWAS / genetics eval.

Download (private repo)

from huggingface_hub import snapshot_download
path = snapshot_download("SingleBicycle/cnn-mesh-xattn-fusion", repo_type="model")  # needs HF login with access

Fastest path to genetics eval: skip inference and use phenotypes_for_loci/*.csv directly (GWAS-ready PC128 traits for KING n=2119). Run inference (below) only for a new cohort.

Honest framing (please keep in any writeup): the single-model differences in this study are within ~1 SE (bootstrap SE on Sum h² ≈ 2.06). The defensible claim is "a recipe of nceq backbone + 4–6 epochs + warmup reproducibly reaches ~+6.2% over CNN-only"; w3b_xmod_ftnceq (+6.38%) is the point-leader of a top cluster (with warm +6.25%, long6 +6.09%) that are statistically tied. Do not report +6.38% as a significant SOTA over the others.

Model card

Name	w3b_xmod_ftnceq — cross-attention fusion, fine-tuned from the InfoNCE+queue ("nceq") CNN backbone + mesh distillation
Result	Sum h² = 37.777 (+6.38% vs CNN-only baseline 35.512) on KING discovery, n=2119
Architecture	FusionAE_v2, fusion_type = cross_mod (4 heads), hidden_dim = 128, mesh_dim = 128
Forward	z_fused = z_cnn + CrossAttn(z_cnn, mesh_adapter(z_mesh)); recon = CNN.decode(z_fused)
Fusion (CrossModBlock)	stack [z_cnn, z_mesh'] as 2 tokens → LayerNorm→MHSA(4 heads)+residual → LayerNorm→MLP(128→256→128)+residual → concat(256)→Linear→delta (zero-init)
Inputs	(1) preprocessed T1 volume (1×182×218×182), (2) z_mesh (128-d cortical-mesh GCN embedding)
Output	128-d z_fused per subject (the embedding; the Sum-h² "phenotype" is a downstream residualize→PCA-128)
Checkpoint meta	best epoch=5, val_loss≈0.2702, fusion_type=cross_mod, distill_weight=10, NCE=SID=EMA=0
Trained on	UK Biobank–derived T1 + cortical meshes (n=6067 train / 606 val); evaluated on KING discovery n=2119

Why "cross-attention + nceq backbone"

A controlled study showed the heritable signal is carried by the backbone, not the fusion operator: the same cross-attention from a fresh CNN scores only +2.18%, but from the contrastively-pretrained nceq backbone it scores +6.38%. Given that backbone, cross-attention ties/edges a plain residual-MLP fusion. See code/make_xmod_nceq_init.py for how the cross_mod model inherits the nceq CNN + mesh-adapter (strict=False, 156 inherited keys, fresh attention).

Files

weights/
  w3b_xmod_ftnceq_ckpt_best.pt    (528 MB)  the model: CNN + cross-attn fusion + mesh_adapter
  w3b_xmod_ftnceq_train.log                  exact training args (provenance)
  mesh_gcn_..._best_recon.pt      (260 KB)  frozen mesh GCN (needed to make z_mesh for NEW subjects)
code/
  v2_fusion_models.py   FusionAE_v2 + CrossModBlock (+ the other fusion blocks)
  cnn_engine.py         CNNAutoencoder (3D-CNN AE)
  datasets.py           VolumeMeshCachedDataset + the FIXED brain z-score T1 normalization
  load_t1_variants.py   T1 normalizations (use `fixed`)
  mesh_branch.py        mesh-GCN wrapper (to compute z_mesh)
  precompute_z.py       offline z_mesh / z_cnn extractor
  v2_extract.py         reference extractor
  v2_train.py           training script (to reproduce the model)
  v2_losses.py          recon + mesh-distillation losses
  make_xmod_nceq_init.py  builds the cross_mod init that inherits the nceq backbone
run_inference_full_cohort.py   << run this to get embeddings
phenotypes_for_loci/           << ready-made quantitative traits for GWAS (n=2119 KING cohort)
  w3b_xmod_ftnceq_pc128.csv     2119 × 128 residualized PCs  (Sum h² = 37.777)
  cnn_only_zcnn_pc128.csv       2119 × 128  baseline (Sum h² = 35.512) — count loci for BOTH
  covariates.csv                genetic PCs + age + sex + UKB imaging fields
  README.txt                    GWAS usage + ID-bridge caveats
README.md  MANIFEST.txt

Environment

Tested in a CUDA PyTorch image with: torch, monai, numpy, pandas, nibabel, tqdm. (The lab uses docker image cortex-fusion-monai:cu128.)

How to run inference on your cohort

Step A — get z_mesh (128-d) per subject

z_fused needs a per-subject cortical-mesh embedding. Two options:

• Reuse cached z_mesh (for subjects already processed): <sid>.npy (128-d) per subject.
• New subjects: run the frozen mesh GCN (weights/mesh_gcn_*.pt via code/mesh_branch.py
  • code/precompute_z.py) on each subject's LH+RH cortical surface (thickness, curvature, coords) to produce a 128-d <sid>.npy. Requires the cortex_fusion mesh stack.

Step B — extract z_fused

python run_inference_full_cohort.py \
  --volume_manifest /path/to/manifest.csv   # columns: sid,t1_path \
  --z_mesh_dir      /path/to/z_mesh \        # one <sid>.npy (128-d) per subject \
  --out_dir        /path/to/out \
  --write_csv                                # also writes embeddings.csv (FID,IID,emb000..emb127)
# (--ckpt defaults to weights/w3b_xmod_ftnceq_ckpt_best.pt; fusion type auto-detected)

Output: one 128-d z_fused <sid>.npy per subject (+ optional embeddings.csv).

Step C — genetics eval (loci / GWAS)

Either use the ready-made phenotypes_for_loci/*.csv (KING n=2119) directly, or build the same from your extracted z_fused: residualize covariates → PCA top-128 → per-PC GWAS → union loci. Count loci for BOTH the model and the baseline so the increase has a reference. See phenotypes_for_loci/README.txt for the ID-bridge and "already-residualized" caveats.

Reproduce training

python code/make_xmod_nceq_init.py     # build cross_mod init from the nceq backbone
python code/v2_train.py --fusion_type cross_mod --mode all \
  --init_ckpt <xmod_init>/ckpt_init.pt --batch_size 2 \
  --use_cosine_lr --lr_proj 5e-4 --lr_branch 1e-5 \
  --epochs 6 --warmup_epochs 1 --patience 6 --distill_weight 10 --save_every_epoch --seed 42 \
  --out_dir runs/w3b_xmod_ftnceq

License / data

Model weights and PCs are derived from UK Biobank data; use is subject to the UK Biobank Material Transfer / data-access terms. IDs are anonymized SUBJ_*. Not for clinical use.