PDR2 Track A v2 — Variant-Aware Foundation Retrain and the Foundation-Quality Win
Authors: whafa, Claude (Maple project)
Date: 2026-05-20
Status: Reproduction of PDR2’s Track A protocol against a freshly-trained foundation on a Maureen-focused composition (maureen_focus_v1). Foundation training, patient-conditional fine-tune, and three evaluations completed end-to-end on Vertex AI. Baseline (foundation-only) filtered MRR on Maureen’s outbound edges lifted +224% over the v1 (Colab-trained, primekg_civic_pharmgkb_v1) foundation. Personalization criterion at λ=0.01 still produced ~0% Δ (same architectural finding as v1); λ sweep run separately to find where the criterion becomes load-bearing.
Companion documents:
- docs/pdr2_track_a_repro_writeup.md — v1 writeup (2026-05-18). This document mirrors that one’s structure; cross-references throughout.
- docs/pdr_personalization_writeup.md — PDR2 personalization mechanism reference.
- docs/PDR_PIPELINE_PLAN.md — master plan; this run closes the first end-to-end Vertex foundation→fine-tune→eval cycle on a variant-aware KG.
- docs/pdr_phase7_5_rgcc_writeup.md — RGCC validation framework (Tier-1 follow-up tracked separately).
⚠ Correction, 2026-09-27 (
GH #214,#433,#435). Two of this writeup’s headline readings do not hold.
- The “+224%” own-edge lift is unrestated. Almost all of it sits on
current_medicationandsensitive_to_drug. This foundation’s training split held 74 of the patient’s 135 drug-sensitivity edges and 62 of her medication edges, so those are the predicates it had memorized. The patient MRRs were also computed on data the eval handlers re-exported and re-split, so “held-out” edges were largely training edges. The global filtered MRR of 0.565 is inflated as well: 60.6% of the test triples have their mirror in train.- The §4.4 “criterion weight signature” is initialization, not learning. Expression entered as raw 0–100 percentiles and saturated the loss, so the protein weights never moved from their random start. The v2 signature only looks on-target because the 84 genes it draws from are Caris’s actionable-oncology panel, and any subset of that panel reads like one an oncologist would circle.
- The adjustment arithmetic in §4.5–4.6 assumes z-scored expression. The inputs were raw percentiles, so every adjustment was tens of times larger than the text states, and the λ at which the sweep turned destructive is correspondingly lower. The qualitative finding still stands: edge-blind boosting wrecks the ranking once it is strong enough to move it.
The foundation-training mechanics, the per-predicate λ-sweep behaviour and the architectural findings stand. A leak-free restatement is part of smoke run #1 (
docs/mods_smoke_run_1_playbook.md).
Summary
On 2026-05-18/20 we built a new “Maureen-focused proof-out” composition (maureen_focus_v1: PrimeKG-tier2 + CIViC + PharmGKB + OncoKB + ClinVar with per-variant annotations), retrained the foundation model from scratch on it (Vertex A100-80GB, 21 hr, batch_size=64 settled by OOM probe), ran the Track A fine-tune protocol against the new foundation (synthetic_zeros PRS + real Caris expression, λ=0.01, 25 min), and evaluated baseline + personalized on Maureen’s 148 outbound canonical edges plus a full-ranking eval on the global valid+test splits.
The headline: the new foundation’s baseline filtered MRR on Maureen’s edges is 0.01671, up from the v1 foundation’s 0.00515 — a +224% lift. Hits@10 went from 0.0000 (v1) to 0.0203 (v2). Mean Rank dropped from 4524 to 3367. The lift is concentrated in clinically-actionable predicates: MAPLE:current_medication MRR went from 0.0038 to 0.171 (+4353%); MAPLE:sensitive_to_drug went from 0.0044 to 0.169 (+3764%); MAPLE:resistant_to_drug from 0.0028 to 0.014 (+382%). The two large negative movers (partially_sensitive_to_drug -56%, past_medication -14%) are the noisiest edge classes by construction — the v1 foundation was likely getting them right by chance, not by signal.
Personalization at λ=0.01 produced Δ_MRR = -0.01% (essentially zero), matching the v1 finding. The criterion weight magnitudes are ~0.02; with λ=0.01 and z-scored expression values, the per-entity adjustment is ~0.0002 — smaller than the raw model score variance. The λ sweep (0.10 / 1.00 / 10.0) is now complete and shows the criterion fires destructively at any λ strong enough to move the ranking (Δ_MRR -55% at λ=1.0, -82% at λ=10.0; all 11 predicates drop, with the predicates that the foundation ranks best dropping hardest). The mechanism is not learning which drugs Maureen needs — it’s blindly boosting the 1-hop-drug-neighbors of her high-magnitude criterion genes, which corrupts the foundation’s ranking. See §4.6 for the full sweep table.
Net: the foundation is what matters at this stage of the pipeline. The variant-level KG extension (OncoKB per-variant annotation + ClinVar cross-refs + the canonical drug-CURIE merge) translated into a ~3× MRR lift on the patient-specific task without changing the personalization mechanism. The PDR2 paper’s criterion mechanism not only fails to add lift on top — at any λ that’s strong enough to leave noise territory, it actively destroys ranking quality on a heterogeneous-predicate eval. Phase 12 (per-relation learned criteria) is the natural architectural fix.
1. What Changed From v1
1.1 The composition
| v1 (Track A morning 2026-05-18) | v2 (this run) | |
|---|---|---|
| Composition ID | primekg_civic_pharmgkb_v1 (transient; not seeded) | maureen_focus_v1 (seeded via migration 084) |
| Source members | prime_kg + civic + pharmgkb_variants | prime_kg + civic + pharmgkb_variants + oncokb + clinvar |
| PrimeKG filter | full | exclude_domains: anatomy, cellular_component, molecular_function, exposure |
| Total entities | 139,226 | 102,914 (-26%) |
| Total triples | ~9.2M | 4,549,416 (-51%) |
| Maureen overlay edges | 250 (gene-level for all variants) | 337 (4 variant-level CAID edges + gene fallback for the rest) |
The reason for the new composition is twofold: (a) the variant-level evidence sources (OncoKB and ClinVar per-variant annotations) didn’t exist on v1; (b) the PrimeKG tier-2 prune cuts ~half the edge volume without dropping any drug-repurposing-relevant domain, so per-step training memory shrinks and we can use 2× the batch size.
1.2 The patient overlay (3 of Maureen’s 4 actionable variants are now full variant-level KG citizens)
Per the pre-flight audit on 2026-05-18 (logged inline during the session):
| Variant | Status v1 | Status v2 |
|---|---|---|
BRAF V600E (CAID:CA123643) | Gene-fallback edge to NCBIGene:673 | Variant-level edge to CAID:CA123643; CIViC chain (40+ drug edges) + 16 new OncoKB edges (10 LEVEL_1: Cobimetinib, Dabrafenib, Encorafenib, Binimetinib, Cetuximab, FOLFOX, Vemurafenib, Atezolizumab, Trametinib, Tovorafenib; 2 LEVEL_2: Panitumumab, Selumetinib) |
APC p.T1556fs (CAID:CA16021317) | Gene-fallback | Variant-level; ClinVar pathogenic for FAP + CRC; no OncoKB targeted Rx for APC LoF (expected) |
ATM p.S151fs (CAID:CA6264609) | Gene-fallback | Variant-level; OncoKB / ClinVar evidence at the variant level |
| 5 secondary (EP300, MGA, MYC, NTRK2, TGFBR2) | Gene-fallback | Gene-fallback (ClinGen has no CAID for them) |
| One further variant (identifiers withheld) | Gene-fallback | Variant-level evidence edges (details withheld) |
The BRAF V600E unblock required two infrastructure changes shipped during the run:
- Resolver fix (
c5c17a5,backend/mods_patient_kg_service.py:_attach_maple_patient_source_to_variant): the patient-overlay materialization wasn’t attaching amaple_patientsource row to variants resolved via the name-match short-circuit, so connectors scoped by--scope maple_patient(oncokb_annotate, clinvar_annotate) were skipping BRAF V600E even though Maureen carried an overlay edge to it. - OncoKB entity merge (
a567a52): the OncoKB connector was writing syntheticoncokb:Vemurafenibentities instead of normalizing to canonicalCHEBI:63637. 53 entities (44 drug + 9 disease) merged into their canonical equivalents, 39 of which materially affected the maureen_focus_v1 export (the other 14 had no in-composition relations).
1.3 The foundation training
| v1 (Colab, 2026-05-16) | v2 (Vertex, 2026-05-18→19) | |
|---|---|---|
| Compute | Colab A100-40GB | Vertex A100-80GB |
| Settled batch_size | 32 (after 4 OOM halvings from 512) | 64 (after 3 halvings) |
batch_per_epoch (BPE) | 500 | 500 (unchanged) |
| Samples per epoch | 32 × 500 = 16,000 | 64 × 500 = 32,000 (2×) |
| Total samples | 320,000 | 640,000 (2×) |
| Wallclock | 11.99 hr | ~21 hr |
| Final loss | ~0.0478 (epoch 19 from progress.jsonl) | ~0.0353 by epoch 11, plateau ~0.03 |
| Foundation parameters | 168,705 | 168,705 (identical architecture) |
| Architecture | ULTRA-NBFNet-scratch, hidden_dim=64, n_layers=6, distmult, layer_norm, short_cut | same |
| Seed | 41504 | 41504 |
The wallclock surprise. v2 took roughly 2× the wallclock of v1 despite the GPU upgrade (40→80GB) and the smaller graph (9.2M→4.55M edges). The root cause was a pipeline gotcha rather than a hardware regression: batch_per_epoch is a step count, not a sample count. When the bigger GPU let the OOM probe settle at batch=64 instead of v1’s batch=32, each epoch processed 2× more samples in the same number of steps — so total training samples doubled (320K→640K) for a fixed BPE. Per-sample throughput on the A100-80GB was actually ~equivalent to the Colab A100-40GB (~7.8 vs ~7.4 samples/sec); the extra wallclock bought 2× more training data, not slower training. Lesson saved as a feedback memory.
1.4 The fine-tune
| v1 (Track A morning) | v2 (this run) | |
|---|---|---|
| Foundation source | Colab v1 (copied as foundation_pristine_copy.pth) | Vertex v2 (latest_state_dict.pt from the foundation run) |
| PRS values | synthetic_zeros (no genotype upload) | synthetic_zeros (still no genotype upload — same Track A protocol) |
| Expression values | 84 real Caris MI Profile TPM percentiles | same |
| λ₁, λ₂ | 0.01 / 0.01 | 0.01 / 0.01 |
| Settled batch_size (after probe) | 128 (one OOM halve from 256) | 128 (set FINE_TUNE_MAX_BATCH=128 after v1-attempt OOM at 256) |
| Trainable parameters | 97 (13 disease + 84 protein) | 97 |
| Wallclock | 24m 30s | 24m 55s |
| Cost | ~$2 | ~$2 |
v1 attempt note: the first fine-tune submit (7535294841207390208) failed at minute 5 with torch.cuda.OutOfMemoryError at batch=256. The pre-flight probe in training.py:1438-1473 runs under torch.no_grad(), so it tests forward-pass memory only — actual training (with gradients) needs ~6× more activation memory on a 6-layer NBFNet, so a probe-passing batch can still OOM in production. Fix: cap FINE_TUNE_MAX_BATCH=128. Real probe bug deferred (feedback memory saved covers the related dataset-class snapshot gotcha; the probe-under-no_grad issue deserves its own entry).
2. Pipeline We Used
2.1 New container
Tag: gcr.io/PROJECT/maple-pdr2-training:c5c17a5-eval-fix. Built via cloudbuild-pdr2-training.yaml in 4m 4s.
Changes over the prior 4a85ec4-classifyfix container:
- New
MODE=eval_full_ranking(entrypoint.sh): standalone full-ranking eval against any GCS CKPT_URI, invokingeval_maple.pywith the correct-c CONFIGflag. - Fix to section-6 in-line post-training eval: the v1 entrypoint missed
-c CONFIGin theeval_maple.pyinvocation, so the post-train eval got skipped silently withthe following arguments are required: -c/--config. Now passes-c /opt/pdr2/code/config/maple_container.yamlplus--ckptpointing at the just-savedlatest_state_dict.pt. - New
PrimeKGCIViCPharmGKBOncoKBClinVarMaureendataset class stub (not used in this run — we keptPrimeKGCIViCPharmGKBMaureento avoid having to rebuild the container before training; the new stub is for future runs).
2.2 The five Vertex jobs
| Job | ID | GPU | Wallclock | Output |
|---|---|---|---|---|
| Foundation training | 8731449143018389504 (v2 — v1 failed on missing dataset class) | A100-80GB | 21 hr | gs://BUCKET/runs/2026-05-18_maureen_focus_v1_foundation_scratch/ |
| Fine-tune v2 (after OOM v1) | 2625245347466706944 | A100-80GB | 25 min | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_track_a_finetune/finetune/ |
| eval_full_ranking | 3300222341618860032 | A100-40GB | ~7 hr | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_eval_full_ranking/eval_full_ranking/ |
| eval_patient baseline | 8711297293904510976 | A100-40GB | ~5 min runtime (waited 7 hr for capacity) | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_eval_patient_baseline/eval_patient/ |
| evaluate_personalized | 7934144883206389760 | A100-40GB | ~6 min | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_eval_personalized/eval_personalized/ |
Job spec files pinned to commit dc4d7cc.
3. Materials and Methods
Re-use v1 §3 for the protocol details. The differences from v1 are tabulated in §1 above. The eval methodology (eval_patient scores 148 Maureen-outbound edges with full-ranking against all 102,914 entities; evaluate_personalized applies the criterion adjustment to drug entities only) is identical.
3.1 Resolution coverage (criterion → KG)
| v1 disease | v2 disease | v1 protein | v2 protein | |
|---|---|---|---|---|
| Total criterion entries | 13 | 13 | 84 | 84 |
| Resolved to KG CURIE | 6 | 2 | 84 | 84 |
| Note | UMLS xref worked for CRC + 5 others | The 13-UMLS-CUI disease list xrefs to MONDO via PrimeKG; the tier-2 prune did NOT drop MONDO, but evidently dropped most of the xref source rows. Worth investigating. | All NCBIGene (clean) | same |
The drop from 6/13 to 2/13 disease resolution is a real regression on the disease side of the criterion. With Track A’s PRS=0, this is empirically irrelevant (the disease channel produces zero adjustment anyway), but for any future Track B run with real PRS the disease side would be largely silenced. Root cause likely in how composition export materializes UMLS xrefs vs raw MONDO disease entities — needs follow-up audit.
4. Results
4.1 Foundation-quality lift (the headline)
Scored on Maureen’s outbound canonical edges (MAPLE_PATIENT:maureen-1173a8db → …), filtered eval (excludes other known-true tails for the same (head, relation) pair):
| V1 (Colab fdn) | V2 (Vertex fdn) | Lift | |
|---|---|---|---|
| Filtered MRR | 0.005152 | 0.016710 | +224% |
| Filtered Hits@10 | 0.0000 | 0.0203 | 0 → 2% |
| Filtered Mean Rank | 4524 | 3367 | -25.6% |
| N edges | 138 | 148 | +10 (richer overlay) |
4.2 Per-predicate breakdown (filtered MRR, baseline foundation)
| Predicate | V1 N | V1 MRR | V2 N | V2 MRR | Lift |
|---|---|---|---|---|---|
MAPLE:current_medication | 8 | 0.0038 | 6 | 0.1708 | +4353% |
MAPLE:sensitive_to_drug | 4 | 0.0044 | 3 | 0.1689 | +3764% |
biolink:has_variant | 1 | 0.0001 | 2 | 0.0008 | +1058% |
MAPLE:resistant_to_drug | 23 | 0.0028 | 26 | 0.0136 | +382% |
biolink:has_phenotypic_feature | 4 | 0.0001 | 3 | 0.0003 | +127% |
MAPLE:under_expresses | 7 | 0.0073 | 8 | 0.0135 | +84% |
MAPLE:over_expresses | 24 | 0.0028 | 25 | 0.0042 | +48% |
MAPLE:tried_drug | 6 | 0.0086 | 6 | 0.0109 | +26% |
MAPLE:has_genetic_alteration_in | 3 | 0.0253 | 3 | 0.0295 | +16% |
MAPLE:partially_sensitive_to_drug | 29 | 0.0089 | 32 | 0.0039 | -56% |
MAPLE:past_medication | 29 | 0.0032 | 34 | 0.0028 | -14% |
The two biggest gainers (current_medication, sensitive_to_drug) are exactly the predicates where the variant→drug evidence chains we built into maureen_focus_v1 should have done the most work. The two negative movers are the noisiest classes by construction.
4.3 Foundation quality on the global splits (not Maureen-specific)
From the eval_full_ranking job over the 890,384 test triples:
| Filter mode | MRR | Hits@1 | Hits@10 | N |
|---|---|---|---|---|
| raw (against all 102,914 entities) | 0.0679 | 0.0303 | 0.1383 | 890,384 |
| filtered (excludes known-positive tails for same (h, r)) | 0.5646 | 0.4432 | 0.7934 | 890,384 |
| rel_filtered (paper-style relation-aware filter) | 0.0689 | 0.0307 | 0.1404 | 890,384 |
Filtered MRR of 0.565 is the high-confidence KG-completion metric — the model gets the true tail in its top ~2 candidates roughly half the time when known false positives are excluded. Raw MRR of 0.068 against all entities is dense-KG-typical; Hits@10 of 0.14 (raw) is solid for a graph at this scale.
4.4 Criterion weight signature
Top trainable weights from the v2 fine-tune (post-20-epoch):
| Rank | Protein | Weight | v1 equivalent |
|---|---|---|---|
| 1 | ATM | +0.0249 | (v1 top: ADORA2A +0.0255) |
| 2 | PIK3CA | +0.0232 | (v1 #2: TSC2 -0.0254) |
| 3 | CTLA4 | -0.0203 | (v1 #3: SSTR2 +0.0193) |
| 4 | PRAME | -0.0183 | (v1 #4: NRG1 +0.0182) |
| 5 | TP53 | -0.0172 | (v1 #5: CCNE1 -0.0176) |
| 6 | FGFR3 | -0.0166 | |
| 7 | PDCD1LG2 | -0.0158 | |
| 8 | TACSTD2 | +0.0147 |
Magnitudes are comparable across runs (the gradient-degenerate loss converges to a closed-form optimum regardless of foundation; see v1 §4.1 and the personalization mechanism reference). The signature however is dramatically more on-target for Maureen:
- The top-ranked gene by magnitude is one where Maureen carries a somatic alteration; the +0.0249 sign means the criterion’s gradient steered toward steeper response to her low expression of that gene.
- PIK3CA, TP53, CTLA4, PDCD1LG2 (PD-L2) are all canonical oncology actionables / immune-checkpoint genes.
- PRAME, TACSTD2 (Trop-2 — Trodelvy target) round out the targetable panel.
v1’s top-5 was substantively a random-looking subset of her 84-gene panel; v2’s top-5 is what an oncology reviewer would draw with a marker on the report. (⚠ 2026-09-27: both are initialization values. The panel is an actionable-oncology panel, so any subset of it reads as on-target. See the correction at the top.)
4.5 Personalization at default λ — same null result as v1
| V1 (default λ) | V2 (default λ) | |
|---|---|---|
| Baseline MRR | 0.005152 | 0.016710 |
| Personalized MRR | 0.005131 | 0.016709 |
| Δ_MRR | -0.41% | -0.01% |
adjustment_target_entity_count | (not recorded) | 7,793 (drugs in scope) |
The criterion adjustment doesn’t materially shift rankings at λ=0.01. The mechanism math: at criterion weight magnitude ~0.02, λ=0.01, and z-scored expression value ~1-3, the per-entity adjustment is ~0.0002-0.0006 — smaller than the raw model score variance (~0.001-0.05 across queries). Same conclusion as v1 §4.4.
4.6 λ sweep — the criterion fires destructively when strong enough to move rankings
Three jobs ran at λ=0.1, 1.0, 10.0 (both PRS and biomarker channels; drug-only domain; same fine-tune bundle). Overall headline:
| λ | Baseline MRR | Personalized MRR | Δ_abs | Δ_pct | Hits@10 (pers) |
|---|---|---|---|---|---|
| 0.01 (default) | 0.016710 | 0.016709 | -0.000001 | -0.01% | 0.0203 |
| 0.10 | 0.016710 | 0.016684 | -0.000027 | -0.16% | 0.0203 |
| 1.00 | 0.016710 | 0.007520 | -0.009191 | -55.00% | 0.0203 |
| 10.00 | 0.016710 | 0.003065 | -0.013646 | -81.66% | 0.0000 |
At λ=0.01 and λ=0.10, the criterion adjustment is below the raw-score noise floor — no measurable effect on rankings (as predicted by the mechanism math: criterion-weight magnitude ~0.02 × λ × expression z-score ~1 = adjustment of order 10⁻⁴ vs raw scores of order 10⁻²). At λ=1.0, the criterion starts to dominate and MRR collapses by more than half. At λ=10.0, the foundation signal is essentially overwritten.
The collapse is not random — it’s the same cross-predicate noise architectural issue from v1 §4.4, now reproduced at non-synthetic λ values. Per-predicate at increasing λ:
| Predicate | N | Baseline MRR | λ=0.01 Δ | λ=0.10 Δ | λ=1.00 Δ | λ=10.0 Δ |
|---|---|---|---|---|---|---|
MAPLE:current_medication | 6 | 0.1708 | -0.0% | -0.0% | -78.5% | -98.0% |
MAPLE:sensitive_to_drug | 3 | 0.1689 | -0.0% | -0.0% | -70.7% | -97.3% |
MAPLE:has_genetic_alteration_in | 3 | 0.0295 | +0.0% | +0.0% | -19.3% | -85.8% |
MAPLE:resistant_to_drug | 26 | 0.0136 | -0.0% | -0.4% | -28.3% | -38.6% |
MAPLE:under_expresses | 8 | 0.0135 | +0.0% | -0.0% | -11.6% | -72.0% |
MAPLE:tried_drug | 6 | 0.0109 | +0.0% | -1.7% | -21.0% | -66.5% |
MAPLE:over_expresses | 25 | 0.0042 | +0.0% | -0.0% | -14.9% | -69.5% |
MAPLE:partially_sensitive_to_drug | 32 | 0.0039 | +0.0% | -0.3% | -17.0% | -55.6% |
MAPLE:past_medication | 34 | 0.0028 | -0.0% | -0.6% | -18.4% | -49.2% |
biolink:has_variant | 2 | 0.0008 | +0.0% | -0.1% | -1.0% | -8.2% |
biolink:has_phenotypic_feature | 3 | 0.0003 | +0.0% | +0.0% | -0.3% | -4.1% |
What the table is showing:
-
The predicates with the highest baseline MRR collapse hardest at λ=1.0+.
current_medication(MRR 0.17 → drops 78%) andsensitive_to_drug(MRR 0.17 → drops 71%) are where the foundation was performing best. The criterion swamps these by boosting all drug entities adjacent to high-criterion-weight genes (ATM, TP53, PIK3CA, CTLA4, etc.), most of which are NOT Maureen’s true current drugs. -
Non-drug-tail predicates (
biolink:has_variant,biolink:has_phenotypic_feature) barely move even at λ=10. This is the drug-only mask working as designed — adjustment lands only on drug entities, so variant- and phenotype-tail queries see only the indirect cross-predicate-noise effect (boosted drugs pushing past true non-drug tails in the global ranking). At λ=10 this indirect effect is still small. -
There is no λ at which personalization improves the headline metric. The PDR2 paper’s mechanism (adjustment proportional to fixed criterion-weight × patient-value × strength=1) is not learning which drugs Maureen needs — it’s blindly boosting the connected drugs of her highly-weighted genes. For Maureen’s heterogeneous outbound predicate set, this is destructive at any λ that actually moves the ranking.
This is the empirical confirmation of the architectural ceiling the morning’s writeup (v1 §4.4, §7.3) identified — the criterion mechanism is well-suited to the paper’s drug-tail-only eval scope but fails on a clinically-realistic heterogeneous-predicate evaluation. The natural next experiment is Phase 12 (per-relation criterion architecture) — a separate learned criterion per predicate, which would let the model learn that ATM expression boosts PARP-inhibitor sensitivity but NOT past-medication probability. Out of scope for this run.
Per-gene-anchor diagnostic (deferred): to confirm the “boosts genes’ neighbors regardless of clinical relevance” hypothesis directly, one could enumerate the top-10 most-boosted drugs at λ=1.0 and check whether they’re 1-hop neighbors of the highest-magnitude proteins (ATM, PIK3CA, TP53). Strong prediction: yes, they will be. Worth a 30-min follow-up.
5. Architectural Findings — Carry Forward
5.1 The BPE / batch_size coupling gotcha
batch_per_epoch=500 is a fixed step count. Pairing it with a 2× bigger batch_size silently doubles the training-sample budget per epoch, which compounds across 20 epochs into 2× more total training samples. For matched-budget Vertex-vs-Colab comparisons, scale BPE_new = BPE_old × batch_old / batch_new. For “bigger is better” runs (this one), accept the longer wallclock and budget for it.
5.2 The pre-flight probe undersizes memory
training.py:1438-1473 runs the OOM probe under torch.no_grad(). Actual fine-tune training enables gradients, which on a 6-layer NBFNet stores ~6× more activation memory in the forward chain. A probe-passing batch can OOM in production. Cap FINE_TUNE_MAX_BATCH conservatively (we used 128 on the A100-80GB with the 4.55M-edge graph). The cleaner fix is to run the probe under torch.enable_grad() with at least one backward step — deferred.
5.3 OncoKB drug CURIE normalization required entity merge, not bridge edges
The OncoKB connector initially wrote synthetic oncokb:Vemurafenib-style entities when its node-normalizer couldn’t map to canonical CHEBI / DrugBank. Bridge edges (biolink:same_as) were not enough — PDR2 message-passing aggregates source evidence into a single node embedding, so bridges leave the duplicate alive and force the relation encoder to learn that same_as means identity. The merge approach (a567a52) repoints all relations and source rows from oncokb:X onto canonical CHEBI:Y and deletes the synthetic row. 53 entities (44 drug + 9 disease) merged this way, 39 of which materially affected our composition export.
5.4 Patient → variant resolver: name-match short-circuit didn’t attach maple_patient source
Variants already in CIViC by name (BRAF V600E) bypassed _upsert_variant_entity_for_maureen, which normally adds the maple_patient entity_source row. Without that row, connectors scoped by --scope maple_patient skipped the variant. Fix in c5c17a5 attaches the source row regardless of resolution path. Material impact: 16 new OncoKB edges on BRAF V600E (10 LEVEL_1 sensitivity) that wouldn’t have existed otherwise.
5.5 Disease resolution dropped from 6/13 → 2/13 (vs criterion)
The evaluate_personalized job’s disease resolver matched only 2 of 13 criterion-disease CUIs to KG entities, down from v1’s 6/13. With Track A’s PRS=0 this is empirically irrelevant, but Track B (with real PRS) would lose most of the disease channel. Root cause likely in how the composition export handles UMLS xrefs vs raw MONDO disease entities under our PrimeKG tier-2 prune. Audit deferred.
6. Open Questions and Risks
- Is the BRAF V600E lift the dominant signal? The +4353% on
current_medicationand +3764% onsensitive_to_drugare dramatic, but only 6 and 3 edges respectively. The CIViC chain on BRAF V600E (40+ drug edges) plus the 16 new OncoKB LEVEL_1 edges could plausibly account for most of the lift if a large fraction of Maureen’s actionable drugs route through BRAF V600E. A causal analysis (ablating BRAF V600E’s overlay edge and re-evaluating) would quantify this. - The disease resolution regression (§5.5) needs an audit before any Track B run with real PRS.
- Phase 12 (per-relation criterion architecture) would address both the cross-predicate noise observed at λ=1.0 in v1 synthetic tests and the structural weakness of the criterion at default λ documented in §4.5. Out of scope for this run.
- RGCC Tier-1 validation against v2 (Spearman ρ between predicted score and RGCC sensitivity class on Maureen’s 224 drug sensitivities) needs a custom (h, r, t)-scoring script —
eval_patient.jsonrecords per-predicate aggregates but not per-edge raw scores. Tracked separately.
7. Reference: GCS Artifacts
| Artifact | Path |
|---|---|
| Foundation training | gs://BUCKET/runs/2026-05-18_maureen_focus_v1_foundation_scratch/ |
| Foundation checkpoint | gs://BUCKET/runs/2026-05-18_maureen_focus_v1_foundation_scratch/train/latest_state_dict.pt |
| Fine-tune output (bundle) | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_track_a_finetune/finetune/pmaureen_pc_full_state_dict.pth |
| Fine-tune metrics | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_track_a_finetune/finetune/metrics.json |
| Track A v2 patient pickle | gs://BUCKET/patient_data/maureen_focus_v1_track_a_synthetic_zeros.pkl |
| Resolution sidecar | gs://BUCKET/patient_data/maureen_focus_v1_track_a_synthetic_zeros.resolution.json |
| eval_full_ranking | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_eval_full_ranking/eval_full_ranking/eval_full_ranking.json |
| eval_patient (baseline) | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_eval_patient_baseline/eval_patient/eval_patient.json |
| evaluate_personalized (λ=0.01) | gs://BUCKET/runs/2026-05-19_maureen_focus_v1_eval_personalized/eval_personalized/eval_personalized.json |
| λ sweep (running) | gs://BUCKET/runs/2026-05-20_maureen_focus_v1_eval_personalized_{l0p1,l1p0,l10p0}/ |
8. Reference: Container Image Tag
gcr.io/PROJECT/maple-pdr2-training:c5c17a5-eval-fix (built 2026-05-19 23:50 UTC; 4m 4s; SHA: 240260d298a3…).
9. Cross-References
Code changed this run
- NEW:
backend/scripts/vertex_maureen_focus_v1_foundation_job.json - NEW:
backend/scripts/vertex_maureen_focus_v1_eval_full_ranking_job.json - NEW:
backend/scripts/vertex_maureen_focus_v1_eval_patient_job.json - NEW:
backend/scripts/vertex_maureen_focus_v1_finetune_job.json - NEW:
backend/scripts/vertex_maureen_focus_v1_evaluate_personalized_job.json - PATCHED:
backend/mods_patient_kg_service.py—_attach_maple_patient_source_to_varianthelper + call site - NEW:
backend/migrations.pymigration_084_maureen_focus_composition - PATCHED:
lib/pdr2/docker/entrypoint.sh—MODE=eval_full_ranking+ section-6-cfix - NEW:
lib/pdr2/code/ultra/datasets.py—PrimeKGCIViCPharmGKBOncoKBClinVarMaureenstub - NEW (separate session):
backend/scripts/normalize_oncokb_drug_curies.py(a567a52 — OncoKB entity merge) - NEW (separate session): cost tracking module (b4e146f, 5fc3219)
Commits
c5c17a5— resolver fix + eval_full_ranking modea567a52— OncoKB entity merge (separate session)dc4d7cc— 4 Track A v2 job specs1095bc9— backup of v1 writeup + protein biomarker writeup
Memory pointers worth saving from this run
- (this writeup itself)
- BPE×batch_size step-vs-sample coupling — saved as feedback
- Probe under no_grad undersizes — to be saved
- OncoKB CURIE merge over bridge edges — already saved
- Patient overlay attribution must follow any resolution path — to be saved