GEMEO Architecture v1.0 — spec + reference impl + figure

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figure1_gemeo_architecture.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,66 @@
+---
+license: cc-by-nc-4.0
+language: [en, pt]
+tags:
+  - world-model
+  - patient-digital-twin
+  - reference-architecture
+  - diffusion-forcing
+  - meds
+  - primekg
+  - rare-disease
+library_name: pytorch
+pipeline_tag: time-series-forecasting
+---
+
+# GEMEO Architecture v1.0
+
+> **A reference architecture for patient world models.** Six principles,
+> pluggable substrate, three open instances. *Not a model* — a recipe for
+> building a model.
+
+![GEMEO Architecture v1.0](./figure1_gemeo_architecture.png)
+
+This repo contains the **architecture specification** and a **reference
+implementation** (Apache-2.0 source, no weights). To use:
+
+1. Read [`gemeo_architecture_spec_v1.md`](./gemeo_architecture_spec_v1.md)
+   for the 6-principle conformance definition.
+2. Copy `reference_impl/` into your repo, adapt to your substrate (any
+   MEDS v0.4.1-compliant EHR), train.
+3. Name your instance `gemeo-<substrate>-v<n>` and submit a conformance
+   report.
+
+## Open instances (May 2026)
+
+| Instance | Substrate | Params | Status |
+|---|---|---|---|
+| [`Raras-AI/gemeo-sus-v2`](https://huggingface.co/Raras-AI/gemeo-sus-v2) | DATASUS (Brazil, 42K patients) | 19.86M | ✅ released |
+| [`Raras-AI/gemeo-twin-stack`](https://huggingface.co/Raras-AI/gemeo-twin-stack) | application layer | NeuralSurv + heads | ✅ released |
+| `Raras-AI/gemeo-mayo-v3` | Mayo Clinic Platform (planned) | 300M | in proposal |
+| `Raras-AI/gemeo-mimic-demo` | MIMIC-IV-DEMO | reference impl | in progress |
+
+## The six architectural principles
+
+1. **Diffusion Forcing backbone** with per-token σ ∼ 𝒰(0, 1)
+2. **Gated KG cross-attention** with tanh(α), α init = 0; real PrimeKG edges
+3. **MEDS v0.4.1 substrate** — `(subject_id, time, code, value)`
+4. **Bootstrap-then-learn** pattern per inference mode
+5. **Bidirectional health-system grounding** (formulary re-rank)
+6. **Audit-driven training** (Chinchilla scaling + SOTA component validation)
+
+Full definitions and conformance tests in [`gemeo_architecture_spec_v1.md`](./gemeo_architecture_spec_v1.md).
+
+## Citation
+
+```bibtex
+@misc{gemeo_arch_v1_2026,
+  title  = {GEMEO Architecture Specification v1.0:
+            Reference architecture for patient world models},
+  author = {Verdial, Dimas Quintas and Kawassaki, Alexandre and the Raras AI team},
+  year   = {2026},
+  url    = {https://huggingface.co/Raras-AI/gemeo-arch},
+}
+```
+
+⚠️ Research only. Not a medical device. No clinical use.

gemeo_architecture_spec_v1.md

@@ -0,0 +1,254 @@
+---
+title: "GEMEO Architecture Specification v1.0"
+subtitle: "Reference architecture for patient world models and healthcare digital twins"
+author: "Raras AI"
+date: "May 2026"
+geometry: margin=1in
+fontsize: 11pt
+mainfont: "Inter"
+---
+
+# GEMEO Architecture Specification v1.0
+
+> **Status:** Stable. This document defines GEMEO Architecture v1.0 — a reference design for **patient world models** in the lineage of Dreamer (Hafner 2019–2025), Diffusion Forcing (Chen NeurIPS 2024), Sora (OpenAI 2024), and Genie (DeepMind 2024–2025), applied to clinical event streams.
+>
+> **License:** CC-BY-NC 4.0 (specification text). Reference implementation in `Raras-AI/gemeo-arch` is Apache-2.0.
+>
+> **Authors:** Raras AI team. Correspondence: dimas@raras.ai
+
+---
+
+## 0. Scope and motivation
+
+A *patient world model* is a learned generative model of patient-trajectory dynamics. It should support: (a) trajectory rollout conditional on actions, (b) counterfactual reasoning under interventions, (c) risk and outcome inference, and (d) drug repurposing. GEMEO Architecture v1.0 is the smallest concrete design we have found that satisfies all four under realistic constraints — sparse clinical events, missing labels, asymmetric data substrates, and regulatory grounding.
+
+This document defines the architecture so that any institution holding longitudinal patient data can instantiate **its own GEMEO model** without sharing data with us. Implementations are expected to ship distinct *instances* (e.g. `gemeo-sus-v2`, `gemeo-mayo-v3`, `gemeo-mimic-demo`) while remaining GEMEO-conformant per §6 of this spec.
+
+## 1. The six architectural principles
+
+A model is **GEMEO-conformant** if and only if it implements all six principles below. Optional extensions are allowed (§5) but do not affect conformance.
+
+### Principle 1 — Diffusion Forcing backbone with per-token σ
+
+The trajectory backbone MUST be a transformer with **per-token noise level** σᵢ ∼ 𝒰(0, 1) sampled independently at each position during training (Chen et al., NeurIPS 2024 [DF]). Score-SDE backbones, pure autoregressive backbones, and uniform-σ backbones are **not** GEMEO-conformant.
+
+Rationale: per-token σ enables variable-horizon rollout and gap-filling under one trained model, which is necessary for both the trajectory mode and the bidirectional masked-fill mode used in clinical-evidence reconstruction.
+
+### Principle 2 — Gated cross-attention to a biomedical knowledge graph
+
+The model MUST attend to a *real* biomedical knowledge-graph ego-subgraph (PrimeKG [Chandak 2023] or equivalent), via cross-attention with a tanh-gated residual:
+
+```
+h' = h + tanh(α) · CrossAttn(h, KG_ego)
+```
+
+with α initialised to zero. The KG cross-attention MUST be inserted at ≥ 1 layer of the backbone. The KG ego-subgraph MUST contain real disease–gene and disease–phenotype edges; *simulated* or *RAG-only* KG conditioning is not GEMEO-conformant.
+
+Rationale: zero-initialised gating (Flamingo [Alayrac 2022], Genie [Bruce 2024]) lets the backbone train as an unconditional model first and adopt KG signal only when it improves the loss, avoiding early-training instability. Real edges anchor the latent state to known biology and prevent hallucinated trajectories under counterfactual rollout.
+
+### Principle 3 — MEDS v0.4.1 substrate
+
+Event input MUST be in the **Medical Event Data Standard v0.4.1** [McDermott 2024]. Each event is a tuple `(subject_id, time, code, numeric_value, text_value)`. Code-prefix patterns MUST follow the canonical MEDS conventions: `ICD10//`, `SIH//`, `APAC//`, `SIGTAP//`, `ORPHA//`, with `MEDS_BIRTH` and `MEDS_DEATH` reserved. The MEDS schema MUST be validated at export time via the official `meds.DataSchema`.
+
+Rationale: MEDS is the lingua franca for EHR foundation models (EHRSHOT, CoMET, CLMBR, MOTOR all use MEDS-derivatives). Compliance ensures any GEMEO instance can be benchmarked against any MEDS-based competitor and that EHR substrates can be swapped without re-engineering the model.
+
+### Principle 4 — Bootstrap-then-learn pattern per inference mode
+
+Each inference mode (trajectory, diagnosis, risk, counterfactual, repurposing, cohort) MUST expose the SAME function signature regardless of whether a learned head exists. When no learned head exists, the function MUST return a deterministic, rule-based or LLM-bootstrap output flagged `model="bootstrap"`. When a learned head exists, the function MUST return the learned output flagged with the head identifier (e.g. `model="neuralsurv"`).
+
+This is not a software-engineering nicety. It is a deployment guarantee: a clinic can deploy a GEMEO instance before any GPU-trained head exists. Capabilities turn on monotonically as checkpoints land in the runtime.
+
+### Principle 5 — Bidirectional health-system grounding
+
+Therapy recommendations MUST be re-ranked by the patient's actual health-system formulary. The re-ranking score MUST include at minimum:
+
+```
+score_grounding(r, p) = π_formulary(r, dx) · ρ_region(r, p_region) · (1 − dist(p, c_ref) / D_max)
+```
+
+where π_formulary ∈ {0, 1} marks formulary membership (PCDT in Brazil, NICE in UK, CMS-approved in US, etc.), ρ_region is the per-region empirical dispensation rate, and c_ref is the closest specialised referral centre. GEMEO instances that recommend therapies the local health system does not deliver are not conformant.
+
+Rationale: clinical AI deployed without grounding to local formularies is not actionable. This principle distinguishes GEMEO from US-only or geography-blind foundation models.
+
+### Principle 6 — Audit-driven training
+
+Every architectural decision MUST be auditable against a contemporaneous SOTA reference. Specifically:
+
+- **Compute scaling**: parameter count MUST respect Chinchilla scaling for the available training tokens. (Our v1 made the inverse error — 125M params on 20M tokens, 30–60× too large — and was downsized to 19.86M in v2.)
+- **Component validation**: each module (DF objective, KG cross-attention, WSD schedule, etc.) MUST cite the paper that established it as SOTA at training time and MUST be ablatable.
+- **Vocabulary integrity**: token codes MUST match the canonical MEDS prefixes (e.g. `ICD10//`, not `CID10//`). Synthetic or hallucinated KG edges are forbidden.
+
+A training run that cannot answer "why this layer, why this size, why this schedule" with a citation does not produce a GEMEO-conformant model.
+
+---
+
+## 2. Required components and module structure
+
+A reference implementation MUST expose the following Python module structure:
+
+```
+gemeo/
+├── __init__.py
+├── core.py                      # orchestrator
+├── api.py                       # FastAPI surface (/api/gemeo/*)
+├── types.py                     # typed dataclasses
+│
+├── cdf/                         # Principle 1 + 2: the world model itself
+│   ├── diffusion_forcing.py     # CDF transformer
+│   ├── adaln_zero.py            # DiT-style σ conditioning
+│   ├── primekg_attention.py     # gated KG cross-attention
+│   ├── wsd_scheduler.py         # WSD LR schedule
+│   ├── meds_export.py           # MEDS v0.4.1 schema validation
+│   └── sample.py                # inference / rollout
+│
+├── encoder.py                   # static patient embedding (HGT or bootstrap)
+├── cohort.py                    # patients-like-mine (Principle 4 bootstrap)
+├── subgraph.py                  # KG sparsification
+├── trajectory.py                # trajectory mode (delegates to cdf/)
+├── risk.py                      # NeuralSurv head (Principle 4)
+├── repurpose.py                 # drug repurposing (TxGNN slot)
+├── whatif.py                    # counterfactual engine
+├── ask.py                       # active learning
+├── ground_sus.py                # Principle 5: health-system grounding
+└── feedback.py                  # closed-loop label ingestion
+```
+
+A subset is allowed — any instance MAY omit modes it does not need (e.g. a research-only instance might skip `ground_sus.py`). But if a module IS present, it MUST conform to the signatures in `types.py`.
+
+---
+
+## 3. Training recipe
+
+### 3.1 Required training-time properties
+
+- **Per-token σ ∼ 𝒰(0.01, 0.99)** sampled independently. Absorbing-state corruption: position *i* becomes MASK with probability σᵢ.
+- **Loss**: masked cross-entropy over corrupted positions, with **Min-SNR weighting** (Hang 2023) on per-token loss.
+- **Conditional dropout 10%**: replace `cond` with `<NULL>` token for classifier-free guidance support at inference.
+- **WSD LR schedule**: 5% warmup / 80% stable / 15% linear decay (Hu et al. MiniCPM 2024). Cosine schedules are not GEMEO-conformant.
+- **bf16 mixed precision**.
+- **Embedding tying** between input and output projections.
+- **Architecture**: SwiGLU + RMSNorm + RoPE (Llama-style). Standard transformer with LayerNorm + ReLU is not conformant.
+
+### 3.2 Required validation gates
+
+Before publication, an instance MUST report:
+
+- Validation cross-entropy on a held-out random-σ task.
+- Integrated Calibration Index (ICI) per Austin & Steyerberg.
+- Gap-fill Top-K (K = 1, 3, 5, 10) on a contiguous mid-trajectory mask of size 24.
+- Multi-horizon Top-K decay curve at prefix lengths *k* ∈ {32, 64, 128}.
+- Per-event-class macro-AUROC with explicit honest reporting of horizon limitations (see §5 of the GEMEO/SUS-v2 model paper for the canonical honest format).
+- Per-subgroup fairness check (sex, age band, region).
+
+Instances that do not report all six are non-conformant for publication-grade claims.
+
+### 3.3 Compute budget guidance
+
+- **Tier-S** (≤ 10M tokens): backbone ≤ 20M params, single H100 < 1 hour, single instance.
+- **Tier-M** (10M–500M tokens): backbone 20–80M params, 8–24 H100-hours.
+- **Tier-L** (500M–10B tokens): backbone 80M–300M params, 100+ H100-hours.
+- **Tier-XL** (10B+ tokens): backbone 300M+ params, 1000+ H100-hours. Mayo / EHRSHOT / multi-modal class.
+
+A tier-S instance with claimed tier-XL performance is suspicious by construction.
+
+---
+
+## 4. Conformance tests
+
+A model is GEMEO-conformant if it passes the following automated tests, runnable via `gemeo-bench`:
+
+```bash
+gemeo-bench check Raras-AI/your-gemeo-instance
+```
+
+**Required tests:**
+
+1. `test_schema_meds` — model exports MEDS-validatable event streams.
+2. `test_per_token_sigma` — model accepts per-position σ vector.
+3. `test_kg_gating_init` — α = 0 at initialisation; sanity-check the gate.
+4. `test_gap_fill_recovery` — gap [24, 48) Top-10 ≥ 0.50 on test split (a real GEMEO instance, even tiny, recovers gaps far above random).
+5. `test_bootstrap_paths` — every inference mode returns a value with `model={bootstrap, learned}` identifier.
+6. `test_health_system_grounding` — at least one therapy recommendation differs by patient region.
+7. `test_audit_citations` — model card contains citations for: DF, AdaLN-Zero, WSD, MEDS, KG source, gating pattern.
+
+The reference test suite is bundled with `Raras-AI/gemeo-arch` and runs in < 60 seconds on CPU for a tiny instance.
+
+---
+
+## 5. Optional extensions (not required for conformance)
+
+- **Self Forcing training** (NeurIPS 2025 Spotlight [Self Forcing 2025]) — addresses tail exposure bias. Recommended for v2+ instances.
+- **Positional features** beyond RoPE — explicit age / calendar-year / region embeddings concatenated with the token embedding. Useful when temporal sparsity is high.
+- **Multimodal substrate** — clinical notes via Gemini / Llama-Med phenotype extraction, WES variant tokens, imaging features. These extend the MEDS event vocabulary but should retain MEDS-prefix conventions.
+- **CoMET-style multi-sample inference** — Monte-Carlo aggregation of *n* trajectories at inference. Note: not required, and per the GEMEO/SUS-v2 ablation in §7.6 of the model paper, it did not outperform the one-shot trained protocol at tier-S.
+
+---
+
+## 6. Versioning and naming
+
+- The architecture spec is versioned (`v1.0` here). Backward-incompatible changes increment the major version.
+- Model instances are named `gemeo/<substrate>-v<version>` (e.g. `gemeo/sus-v2`, `gemeo/mayo-v3`, `gemeo/mimic-demo-v1`).
+- The reference implementation is `Raras-AI/gemeo-arch` (architecture, no weights). Instances live at `Raras-AI/gemeo-<substrate>-v<n>`.
+- Bundled application layer: `Raras-AI/gemeo-twin-stack` (the six-mode wrapper around any conformant instance).
+
+---
+
+## 7. Validated instances (May 2026)
+
+| Instance | Substrate | Params | Status | Reference |
+|---|---|---|---|---|
+| `gemeo-sus-v2` | DATASUS (Brazil, 42 K patients) | 19.86 M | ✅ open | Verdial 2026 |
+| `gemeo-mayo-v3` | Mayo Clinic Platform (planned, 3 M) | 300 M (planned) | in proposal | Mayo Accelerate proposal |
+| `gemeo-mimic-demo` | MIMIC-IV-DEMO (Open, ~100 pts) | reference impl | in progress | — |
+
+External instances are welcome. Submit a pull request to `Raras-AI/gemeo-arch` with the conformance-test output to be listed.
+
+---
+
+## 8. Related architectures (not GEMEO-conformant, by design)
+
+We name these explicitly so that reviewers and adopters can distinguish:
+
+- **Sora** [OpenAI 2024]: world model for video, not patient events. DF lineage but no clinical substrate.
+- **Dreamer 4** [Hafner 2025]: world model for general environments. DF lineage. Not specialised to MEDS, no KG cross-attention, no health-system grounding.
+- **EHRWorld** [arXiv 2602.03569]: autoregressive Qwen-fine-tune on EHR. Not Diffusion Forcing. Not conformant to Principle 1.
+- **CoMET** [Epic 2508.12104]: large-scale generative medical event model. Not Diffusion Forcing — autoregressive only. Not conformant to Principle 1.
+- **MOTOR** [Steinberg 2023]: time-to-event foundation model. Not generative-dynamics. Conformant to substrate (MEDS) but not to Principles 1, 2.
+- **CLMBR-T** [Stanford]: contrastive foundation model. Not generative. Not conformant to Principle 1.
+- **Delphi-2M** [Nature 2025]: autoregressive transformer over UK Biobank + Danish registry. Not Diffusion Forcing. Not conformant to Principle 1.
+- **RareGraph-Synth** [arXiv 2510.06267]: score-SDE diffusion on *synthetic* rare-disease graphs. Not GEMEO-conformant: synthetic substrate (violates Principle 6's audit requirements) and score-SDE not Diffusion Forcing.
+
+GEMEO is the architecture that combines (DF + KG cross-attn + MEDS + health-system grounding + bootstrap-then-learn + audit) into a single deployable design. To our knowledge, no prior architecture satisfies all six.
+
+---
+
+## 9. Citation
+
+```bibtex
+@misc{gemeo_arch_v1_2026,
+  title  = {GEMEO Architecture Specification v1.0:
+            Reference architecture for patient world models},
+  author = {Verdial, Dimas Quintas and Kawassaki, Alexandre and the Raras AI team},
+  year   = {2026},
+  url    = {https://huggingface.co/Raras-AI/gemeo-arch},
+  note   = {Specification document. Reference implementation Apache-2.0.}
+}
+```
+
+## 10. References
+
+- [DF] Chen, B. et al. *Diffusion Forcing: Next-Token Prediction Meets Full-Sequence Diffusion.* NeurIPS 2024 (arXiv:2407.01392).
+- [Dreamer 4] Hafner, D. et al. *Training Agents Inside of Scalable World Models.* arXiv:2509.24527, Sept 2025.
+- [Sora] OpenAI. *Video generation models as world simulators.* Feb 2024.
+- [Genie] Bruce, J. et al. *Genie: Generative Interactive Environments.* ICML 2024.
+- [PrimeKG] Chandak, P., Huang, K., Zitnik, M. *Building a knowledge graph to enable precision medicine.* Nature Sci Data 2023.
+- [MEDS] McDermott, M. et al. *MEDS: Medical Event Data Standard v0.4.1.* GitHub, 2024.
+- [DiT] Peebles, W. & Xie, S. *Scalable Diffusion Models with Transformers.* ICCV 2023.
+- [Flamingo] Alayrac, J.-B. et al. *Flamingo: a Visual Language Model for Few-Shot Learning.* NeurIPS 2022.
+- [MiniCPM] Hu, S. et al. *MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies.* 2024.
+- [Min-SNR] Hang, T. et al. *Efficient Diffusion Training via Min-SNR Weighting Strategy.* ICCV 2023.
+- [Self Forcing 2025] *Self Forcing: Training Diffusion-Forcing Agents Without Exposure Bias.* arXiv:2506.08009, NeurIPS 2025 Spotlight.
+
+---
+
+*End of GEMEO Architecture Specification v1.0.*

reference_impl/__init__.py

@@ -0,0 +1,33 @@
+"""GEMEO-CDF: Causal Diffusion Forcing for clinical trajectories.
+
+Three "first in medicine" hooks:
+  1. DIFFUSION FORCING (Chen MIT NeurIPS 2024 → Dreamer 4 Hafner 2025 backbone)
+     — independent per-token noise levels unify AR + diffusion + counterfactual
+       in ONE loss. Zero clinical port as of May 2026.
+
+  2. LATENT ACTION MODEL (Genie / DeepMind 2024)
+     — VQ-VAE codebook over (state_t, state_{t+1}) deltas discovers a
+       treatment vocabulary without RxNorm/ATC labels. Solves the APAC
+       miscoding / sparsity / off-label labelling pain in DATASUS.
+
+  3. PROCESS REWARD VERIFIER (o3 / MAI-DxO 2025 pattern)
+     — small PRM scores top-K rollouts at inference, returns top-1 +
+       uncertainty band. Deliberative trajectory generation, novel in EHR.
+
+Modules:
+  diffusion_forcing.py  — core architecture (per-token noise + block-causal)
+  lam.py                — Latent Action Model (VQ-VAE codebook)
+  train_cdf.py          — training loop with diffusion forcing objective
+  sample.py             — sampling: AR mode / denoise mode / counterfactual
+  distill.py            — Shortcut Forcing distillation (Dreamer 4)
+  prm.py                — Process Reward Verifier
+"""
+from .diffusion_forcing import CDFTransformer, CDFConfig
+from .lam import LatentActionVQVAE, LAMConfig
+from .train_cdf import train_cdf
+
+__all__ = [
+    "CDFTransformer", "CDFConfig",
+    "LatentActionVQVAE", "LAMConfig",
+    "train_cdf",
+]

reference_impl/adaln_zero.py

@@ -0,0 +1,148 @@
+"""AdaLN-Zero conditioning module (DiT-style, Peebles 2023).
+
+Used in: DiT (ICCV 2023), Stable Diffusion 3 (Esser 2024), Sora, Lumina-Next,
+PixArt-Sigma. Standard for diffusion conditioning in 2025-2026.
+
+Why for Diffusion Forcing on EHR:
+  - Per-token sigma + global cond/action → per-token (scale, shift, gate)
+  - Gates init to zero ⇒ block starts as identity ⇒ no catastrophic init
+  - Much better CFG (dropped condition path goes through zero gates,
+    not corrupting residual stream)
+  - DFoT (Diffusion Forcing Transformer 2, ICLR 2026) confirms +3-8% win
+
+We fuse THREE conditioning signals:
+  - sigma (B, T) per-token noise level → time_emb (B, T, D)
+  - cond  (B,)   cohort-level treatment id → cond_emb (B, D) → broadcast
+  - action(B, T) per-token latent action id → action_emb (B, T, D)
+
+Combined into c_t (B, T, D) → ConditioningMLP → 6 modulation tensors
+per block. Each block uses them as:
+
+    h = x + gate_msa * Attn(scale_msa * Norm(x) + shift_msa)
+    h = h + gate_mlp * MLP(scale_mlp * Norm(h) + shift_mlp)
+"""
+from __future__ import annotations
+import torch
+import torch.nn as nn
+
+
+class AdaLNZeroModulator(nn.Module):
+    """Generates per-token (scale, shift, gate) for AdaLN-Zero block.
+
+    Input: fused conditioning vector c (B, T, d_model).
+    Output: 6 tensors of shape (B, T, d_model) each:
+      (scale_msa, shift_msa, gate_msa, scale_mlp, shift_mlp, gate_mlp)
+    """
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.modulator = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(d_model, 6 * d_model, bias=True),
+        )
+        # Zero-init for the gate-producing rows (AdaLN-Zero trick)
+        # We zero-init ALL outputs initially; gate stays zero so block is identity
+        nn.init.zeros_(self.modulator[-1].weight)
+        nn.init.zeros_(self.modulator[-1].bias)
+
+    def forward(self, c: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        # c: (B, T, d_model)
+        out = self.modulator(c)  # (B, T, 6*d_model)
+        return out.chunk(6, dim=-1)
+
+
+class AdaLNZeroBlock(nn.Module):
+    """Transformer block with AdaLN-Zero modulation.
+
+    Drop-in replacement for the standard pre-norm block. Reads
+    pre-computed modulation tensors and applies them around Attn + MLP.
+    """
+    def __init__(self, d_model: int, n_heads: int, ffn: int, dropout: float,
+                 rope=None, kg_xattn=None):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.rope = rope
+        self.kg_xattn = kg_xattn
+
+        self.norm1 = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.proj = nn.Linear(d_model, d_model, bias=False)
+        self.norm2 = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model, ffn, bias=False),
+            nn.GELU(),
+            nn.Linear(ffn, d_model, bias=False),
+        )
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor,
+                scale_msa, shift_msa, gate_msa,
+                scale_mlp, shift_mlp, gate_mlp,
+                kg_ctx: torch.Tensor | None = None) -> torch.Tensor:
+        import torch.nn.functional as F
+        B, T, D = x.shape
+        # MSA branch
+        h = self.norm1(x) * (1 + scale_msa) + shift_msa
+        qkv = self.qkv(h).reshape(B, T, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
+        if self.rope is not None:
+            q, k = self.rope(q, k, T)
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=(~attn_mask).float().masked_fill(attn_mask, float("-inf"))[None, None],
+            dropout_p=self.dropout.p if self.training else 0.0,
+        )
+        out = out.transpose(1, 2).reshape(B, T, D)
+        x = x + gate_msa * self.dropout(self.proj(out))
+        # KG cross-attention (between MSA and MLP)
+        if self.kg_xattn is not None and kg_ctx is not None:
+            x = self.kg_xattn(x, kg_ctx)
+        # MLP branch
+        h = self.norm2(x) * (1 + scale_mlp) + shift_mlp
+        x = x + gate_mlp * self.dropout(self.mlp(h))
+        return x
+
+
+class FusedConditioner(nn.Module):
+    """Fuse (sigma, cond, action) into one per-token conditioning vector.
+
+    Output (B, T, d_model) consumed by AdaLNZeroModulator per layer.
+    """
+    def __init__(self, d_model: int, n_conditions: int, n_actions: int,
+                 use_action: bool = True):
+        super().__init__()
+        self.d_model = d_model
+        self.use_action = use_action
+        # Sigma → sinusoidal embedding
+        self.sigma_proj = nn.Sequential(
+            nn.Linear(d_model, d_model), nn.SiLU(), nn.Linear(d_model, d_model),
+        )
+        self.cond_emb = nn.Embedding(n_conditions, d_model)
+        if use_action:
+            self.action_emb = nn.Embedding(n_actions + 1, d_model)
+        self.fuse = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(d_model, d_model),
+        )
+
+    def sinusoidal(self, sigma: torch.Tensor) -> torch.Tensor:
+        import math
+        half = self.d_model // 2
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, device=sigma.device) / half
+        )
+        ang = sigma.float().unsqueeze(-1) * freqs
+        emb = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)
+        return self.sigma_proj(emb)
+
+    def forward(self, sigma: torch.Tensor, cond: torch.Tensor,
+                action: torch.Tensor | None = None) -> torch.Tensor:
+        # sigma (B, T) → time_emb (B, T, D)
+        time_emb = self.sinusoidal(sigma)
+        # cond (B,) → (B, D) → broadcast to (B, T, D)
+        cond_emb = self.cond_emb(cond).unsqueeze(1).expand_as(time_emb)
+        fused = time_emb + cond_emb
+        if self.use_action and action is not None:
+            fused = fused + self.action_emb(action)
+        return self.fuse(fused)

reference_impl/diffusion_forcing_v13.py

@@ -0,0 +1,314 @@
+"""GEMEO-CDF v13 — audit-driven Chinchilla-correct architecture.
+
+Per the SOTA audit (May 2026):
+  - Path B (CLMBR fine-tune) BLOCKED: CLMBR-T-base is HF-gated (manual approval)
+  - Path A adopted: small from-scratch model + KG adapters + MEDS interop
+
+Architecture:
+  - 12M backbone params (Chinchilla-respecting for ~20M token corpus)
+  - d_model=384, n_layers=8, n_heads=6, ffn=1024, ctx=512
+  - SwiGLU MLP (ffn:d_model = 2.67)
+  - Tied embeddings (saves ~12M at vocab=32k)
+  - Dropout 0.1 everywhere (small-data critical)
+  - Block-causal attention (Diffusion Forcing)
+  - Per-token sigma noise (independent)
+  - GATED KG cross-attention (tanh(α)·xattn, α init=0)
+    - Layers 4, 6, 7 (3 of 8)
+    - Lets model learn to use KG progressively, doesn't disrupt early loss
+  - DF objective + LM-aux loss (joint training, paper-grade)
+
+Sources audited:
+  - CoMET (Aug 2025): tokens-per-param ratio
+  - CLMBR (Stanford): adapter pattern for cross-site transfer
+  - MDLM (Sahoo 2024): masked diffusion, matches AR at equal FLOPs
+  - Genie (DeepMind 2024): gated cross-attention pattern
+  - SD3 (Esser 2024): AdaLN-Zero zero-init gates
+"""
+from __future__ import annotations
+import math
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class CDFv13Config:
+    # Vocab + sequence
+    vocab_size: int = 32768       # MEDS-derived (will be much smaller in practice)
+    mask_token: int = 32767
+    max_seq_len: int = 512
+    block_size: int = 16
+    # Architecture (Chinchilla-correct for ~20M tokens)
+    d_model: int = 384
+    n_heads: int = 6
+    n_layers: int = 8
+    ffn: int = 1024                # SwiGLU effective; flag below uses 2 projections
+    dropout: float = 0.1
+    emb_dropout: float = 0.1
+    use_swiglu: bool = True
+    use_rmsnorm: bool = True
+    tie_embeddings: bool = True
+    # Diffusion forcing
+    cond_dropout: float = 0.10
+    # KG conditioning (GATED adapters)
+    use_kg: bool = True
+    kg_dim: int = 3072
+    kg_attn_layers: list = field(default_factory=lambda: [4, 6, 7])
+    # Latent action
+    use_latent_action: bool = False  # Dropped per audit (concept shaky)
+    n_latent_actions: int = 512
+    # Conditioning
+    n_conditions: int = 64
+
+
+class RMSNorm(nn.Module):
+    """Root-mean-square LayerNorm (LLaMA/Mistral style)."""
+    def __init__(self, d: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (norm * self.weight.float()).to(x.dtype)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU MLP (used in LLaMA/Gemma/Mistral)."""
+    def __init__(self, d_in: int, d_hidden: int, dropout: float = 0.1):
+        super().__init__()
+        self.w_gate = nn.Linear(d_in, d_hidden, bias=False)
+        self.w_up = nn.Linear(d_in, d_hidden, bias=False)
+        self.w_down = nn.Linear(d_hidden, d_in, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dropout(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))
+
+
+class RotaryEmbedding(nn.Module):
+    """RoPE (Su et al. 2021)."""
+    def __init__(self, dim: int, max_seq: int = 8192, base: float = 10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        t = torch.arange(max_seq).float()
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("cos", emb.cos(), persistent=False)
+        self.register_buffer("sin", emb.sin(), persistent=False)
+
+    def forward(self, q, k, seq_len):
+        cos = self.cos[:seq_len].to(q.dtype).to(q.device)
+        sin = self.sin[:seq_len].to(q.dtype).to(q.device)
+        def rot_half(x):
+            x1, x2 = x.chunk(2, dim=-1)
+            return torch.cat([-x2, x1], dim=-1)
+        return (q * cos) + (rot_half(q) * sin), (k * cos) + (rot_half(k) * sin)
+
+
+class PerTokenSigmaEmbed(nn.Module):
+    """Sinusoidal embedding of per-position diffusion noise sigma in [0,1]."""
+    def __init__(self, d: int):
+        super().__init__()
+        self.d = d
+        self.proj = nn.Sequential(
+            nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d),
+        )
+
+    def forward(self, sigma: torch.Tensor) -> torch.Tensor:
+        half = self.d // 2
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, device=sigma.device) / half
+        )
+        ang = sigma.float().unsqueeze(-1) * freqs
+        emb = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)
+        return self.proj(emb)
+
+
+class GatedKGCrossAttention(nn.Module):
+    """Cross-attention to KG ego-subgraph, with GATED output.
+
+    `tanh(alpha) * cross_attn(x_seq, x_kg)` where alpha is a learnable scalar
+    initialized to 0. This means at init the cross-attention contributes
+    NOTHING to the residual stream, so the model trains identically to
+    no-KG until it discovers KG is useful. Prevents catastrophic loss
+    spikes on small data.
+
+    Pattern from: Genie (DeepMind 2024), Flamingo (DeepMind 2022).
+    """
+    def __init__(self, d_model: int, kg_dim: int, n_heads: int = 8, dropout: float = 0.1):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        # Project KG to d_model (run inline so we don't need separate KGProjector module)
+        self.kg_in_proj = nn.Linear(kg_dim, d_model, bias=False)
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.norm_q = RMSNorm(d_model)
+        self.norm_kv = RMSNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        # Gate (scalar per block, init=0)
+        self.alpha = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x_seq: torch.Tensor, kg_raw: torch.Tensor) -> torch.Tensor:
+        """
+        x_seq: (B, T, d_model)
+        kg_raw: (B, N_kg, kg_dim)   -- raw KG embeddings (e.g. 3072)
+        """
+        B, T, D = x_seq.shape
+        kg_proj = self.kg_in_proj(kg_raw)  # (B, N_kg, D)
+        N_kg = kg_proj.size(1)
+        q = self.q_proj(self.norm_q(x_seq))
+        kv = self.kv_proj(self.norm_kv(kg_proj))
+        k, v = kv.chunk(2, dim=-1)
+        q = q.reshape(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.reshape(B, N_kg, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.reshape(B, N_kg, self.n_heads, self.head_dim).transpose(1, 2)
+        out = F.scaled_dot_product_attention(
+            q, k, v, dropout_p=self.dropout.p if self.training else 0.0)
+        out = out.transpose(1, 2).reshape(B, T, D)
+        gate = torch.tanh(self.alpha)
+        return x_seq + gate * self.dropout(self.out_proj(out))
+
+
+class CDFv13Block(nn.Module):
+    """Pre-norm transformer block + optional gated KG cross-attn."""
+    def __init__(self, cfg: CDFv13Config, rope: RotaryEmbedding,
+                 layer_idx: int):
+        super().__init__()
+        self.cfg = cfg
+        self.rope = rope
+        self.layer_idx = layer_idx
+        norm_cls = RMSNorm if cfg.use_rmsnorm else nn.LayerNorm
+        self.norm1 = norm_cls(cfg.d_model)
+        self.norm2 = norm_cls(cfg.d_model)
+        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        if cfg.use_swiglu:
+            self.mlp = SwiGLU(cfg.d_model, cfg.ffn, cfg.dropout)
+        else:
+            self.mlp = nn.Sequential(
+                nn.Linear(cfg.d_model, cfg.ffn, bias=False),
+                nn.GELU(),
+                nn.Linear(cfg.ffn, cfg.d_model, bias=False),
+                nn.Dropout(cfg.dropout),
+            )
+        self.dropout = nn.Dropout(cfg.dropout)
+        self.head_dim = cfg.d_model // cfg.n_heads
+        # Gated KG cross-attention (only in specified layers)
+        self.use_kg_in_layer = cfg.use_kg and layer_idx in cfg.kg_attn_layers
+        if self.use_kg_in_layer:
+            self.kg_xattn = GatedKGCrossAttention(
+                cfg.d_model, cfg.kg_dim, cfg.n_heads, cfg.dropout)
+
+    def forward(self, x, attn_mask, kg_raw=None):
+        B, T, D = x.shape
+        # MSA
+        h = self.norm1(x)
+        qkv = self.qkv(h).reshape(B, T, 3, self.cfg.n_heads, self.head_dim)
+        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
+        q, k = self.rope(q, k, T)
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=(~attn_mask).float().masked_fill(attn_mask, float("-inf"))[None, None],
+            dropout_p=self.cfg.dropout if self.training else 0.0,
+        )
+        out = out.transpose(1, 2).reshape(B, T, D)
+        x = x + self.dropout(self.proj(out))
+        # Gated KG cross-attn (if enabled at this layer)
+        if self.use_kg_in_layer and kg_raw is not None:
+            x = self.kg_xattn(x, kg_raw)
+        # MLP
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class CDFv13Transformer(nn.Module):
+    """Audit-compliant CDF v13: 12M backbone + KG adapters + DF objective."""
+
+    def __init__(self, cfg: CDFv13Config | None = None):
+        super().__init__()
+        self.cfg = cfg or CDFv13Config()
+        c = self.cfg
+        norm_cls = RMSNorm if c.use_rmsnorm else nn.LayerNorm
+
+        self.tok_emb = nn.Embedding(c.vocab_size, c.d_model)
+        self.emb_dropout = nn.Dropout(c.emb_dropout)
+
+        # Per-token sigma embedding (additive)
+        self.sigma_emb = PerTokenSigmaEmbed(c.d_model)
+        # Global condition embedding (additive, broadcast)
+        self.cond_emb = nn.Embedding(c.n_conditions, c.d_model)
+
+        # RoPE
+        self.rope = RotaryEmbedding(c.d_model // c.n_heads, max_seq=c.max_seq_len * 2)
+
+        # Blocks
+        self.blocks = nn.ModuleList([
+            CDFv13Block(c, self.rope, layer_idx=i) for i in range(c.n_layers)
+        ])
+        self.final_norm = norm_cls(c.d_model)
+        self.head = nn.Linear(c.d_model, c.vocab_size, bias=False)
+        if c.tie_embeddings:
+            self.head.weight = self.tok_emb.weight
+
+        # Block-causal mask buffer
+        T = c.max_seq_len
+        block_id = torch.arange(T) // c.block_size
+        mask = block_id.unsqueeze(0) < block_id.unsqueeze(1)
+        self.register_buffer("block_mask", mask, persistent=False)
+
+        # Init
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+            if m.bias is not None: nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    def forward(self, x, sigma, cond, kg_raw=None):
+        B, T = x.shape
+        h = self.tok_emb(x) + self.sigma_emb(sigma) + self.cond_emb(cond).unsqueeze(1)
+        h = self.emb_dropout(h)
+        mask = self.block_mask[:T, :T]
+        for blk in self.blocks:
+            h = blk(h, mask, kg_raw=kg_raw)
+        h = self.final_norm(h)
+        return self.head(h)
+
+    def diffusion_forcing_loss(self, x_clean, cond, kg_raw=None,
+                                mode: str = "uniform") -> torch.Tensor:
+        """Standard absorbing-state DF loss with per-token sigma.
+
+        mode: 'uniform' (default — safer for discrete than logit-normal per audit)
+              'logit_normal' (SD3-style — keep as ablation only)
+        """
+        B, T = x_clean.shape
+        device = x_clean.device
+        # CFG cond dropout
+        drop = torch.rand(B, device=device) < self.cfg.cond_dropout
+        cond = torch.where(drop, torch.zeros_like(cond), cond)
+        if kg_raw is not None:
+            drop_kg = (torch.rand(B, device=device) < self.cfg.cond_dropout).float()
+            kg_raw = kg_raw * (1 - drop_kg).reshape(B, 1, 1)
+        # Sample per-token sigma
+        if mode == "logit_normal":
+            sigma = torch.sigmoid(torch.randn(B, T, device=device)).clamp(0.01, 0.99)
+        else:
+            sigma = torch.rand(B, T, device=device).clamp(0.01, 0.99)
+        # Absorbing-state corruption
+        corrupt = torch.rand(B, T, device=device) < sigma
+        x_noisy = torch.where(corrupt, self.cfg.mask_token, x_clean)
+        logits = self.forward(x_noisy, sigma, cond, kg_raw=kg_raw)
+        ce = F.cross_entropy(
+            logits.reshape(-1, self.cfg.vocab_size),
+            x_clean.reshape(-1),
+            reduction="none",
+        ).reshape(B, T)
+        n = corrupt.float().sum().clamp(min=1.0)
+        return (ce * corrupt.float()).sum() / n

reference_impl/eval_sota.py

@@ -0,0 +1,290 @@
+"""SOTA evaluation suite for CDFv13 — audit-proof.
+
+Per the May 2026 SOTA audit, replaces "Top-1 mid-position" (not recognized)
+with the canonical EHR foundation model metric stack:
+
+  Classification (next-event, downstream tasks):
+    - AUROC + AUPRC + Brier
+    - Calibration: ICI (Austin & Steyerberg 2019)
+    - Decision-curve analysis (Vickers)
+    - Bootstrap 95% CI (≥2000 resamples) — required for rare disease
+
+  Survival (DATASUS SIM mortality):
+    - Uno's C (concordance_index_ipcw) — preferred over Harrell at high censoring
+    - Integrated Brier Score (1/3/5y)
+    - Time-dependent AUC
+
+  Counterfactual / causal:
+    - ATE with bootstrap CI
+    - E-value (VanderWeele)
+    - Negative-control outcome + exposure
+    - Tipping-point analysis
+
+  Generation fidelity (CoMET / SynthEHRella):
+    - Dim-wise probability match
+    - MMD (Maximum Mean Discrepancy) with RBF kernel
+    - TSTR (Train-on-Synthetic-Test-on-Real)
+
+  Subgroup fairness (npj DM requirement):
+    - Stratified metrics: sex, age band, UF region
+
+  Split strategy (DATASUS rare disease):
+    - Temporal: train ≤2022, val 2023, test 2024-2025
+    - Geographic: train SE+S, test N+NE (UF cross-region = "external")
+    - Patient-level 5-fold CV (variance estimation)
+"""
+from __future__ import annotations
+import math
+import numpy as np
+import torch
+from typing import Callable
+
+
+# ---------- Classification ----------
+
+def auroc(y: np.ndarray, p: np.ndarray) -> float:
+    from sklearn.metrics import roc_auc_score
+    if len(np.unique(y)) < 2: return float("nan")
+    return roc_auc_score(y, p)
+
+
+def auprc(y: np.ndarray, p: np.ndarray) -> float:
+    from sklearn.metrics import average_precision_score
+    if len(np.unique(y)) < 2: return float("nan")
+    return average_precision_score(y, p)
+
+
+def brier(y: np.ndarray, p: np.ndarray) -> float:
+    from sklearn.metrics import brier_score_loss
+    return brier_score_loss(y, p)
+
+
+def ici(y: np.ndarray, p: np.ndarray, frac: float = 0.75) -> float:
+    """Integrated Calibration Index (Austin & Steyerberg 2019).
+    Lowess-smoothed deviation from perfect calibration.
+    """
+    from statsmodels.nonparametric.smoothers_lowess import lowess
+    sm = lowess(y, p, frac=frac, return_sorted=True)
+    return float(np.mean(np.abs(sm[:, 1] - sm[:, 0])))
+
+
+def net_benefit(y: np.ndarray, p: np.ndarray, threshold: float) -> float:
+    """Net benefit at a given decision threshold (Vickers DCA)."""
+    tp = ((p >= threshold) & (y == 1)).sum()
+    fp = ((p >= threshold) & (y == 0)).sum()
+    n = len(y)
+    if threshold >= 1.0: return 0.0
+    return tp / n - (fp / n) * (threshold / (1 - threshold))
+
+
+def decision_curve(y: np.ndarray, p: np.ndarray,
+                   thresholds: list[float] = None) -> dict:
+    """Decision-curve analysis: net benefit across thresholds vs treat-all/treat-none."""
+    if thresholds is None:
+        thresholds = [0.01, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5]
+    model_nb = [net_benefit(y, p, t) for t in thresholds]
+    treat_all_nb = [(y.mean()) - (1 - y.mean()) * (t / (1 - t)) if t < 1 else 0
+                    for t in thresholds]
+    treat_none_nb = [0.0] * len(thresholds)
+    return {
+        "thresholds": thresholds,
+        "model": model_nb,
+        "treat_all": treat_all_nb,
+        "treat_none": treat_none_nb,
+    }
+
+
+def bootstrap_ci(y: np.ndarray, p: np.ndarray, metric_fn: Callable,
+                 n_boot: int = 2000, seed: int = 0,
+                 ci: tuple[float, float] = (2.5, 97.5)) -> tuple[float, float, float]:
+    """Bootstrap 95% CI for any (y, p) -> scalar metric."""
+    rng = np.random.default_rng(seed)
+    n = len(y)
+    stats = []
+    for _ in range(n_boot):
+        idx = rng.integers(0, n, n)
+        if len(np.unique(y[idx])) < 2: continue
+        try:
+            stats.append(metric_fn(y[idx], p[idx]))
+        except Exception:
+            continue
+    if not stats: return (float("nan"),) * 3
+    return (
+        float(np.percentile(stats, ci[0])),
+        float(np.median(stats)),
+        float(np.percentile(stats, ci[1])),
+    )
+
+
+# ---------- Survival ----------
+
+def uno_c_index(y_train_event, y_train_time, y_test_event, y_test_time,
+                risk_score, tau: float = None) -> float:
+    """Uno's C-index (IPCW concordance), preferred at high censoring.
+    Requires scikit-survival.
+    """
+    try:
+        from sksurv.metrics import concordance_index_ipcw
+    except ImportError:
+        return float("nan")
+    # Build structured arrays
+    y_train = np.array(
+        list(zip(y_train_event.astype(bool), y_train_time.astype(float))),
+        dtype=[("event", "?"), ("time", "<f8")],
+    )
+    y_test = np.array(
+        list(zip(y_test_event.astype(bool), y_test_time.astype(float))),
+        dtype=[("event", "?"), ("time", "<f8")],
+    )
+    if tau is None:
+        tau = float(y_test_time.max()) * 0.95
+    c, *_ = concordance_index_ipcw(y_train, y_test, risk_score, tau=tau)
+    return float(c)
+
+
+def integrated_brier_score(y_train_event, y_train_time, y_test_event, y_test_time,
+                            surv_pred: np.ndarray, times: np.ndarray) -> float:
+    """Integrated Brier Score (lower is better)."""
+    try:
+        from sksurv.metrics import integrated_brier_score as ibs_fn
+    except ImportError:
+        return float("nan")
+    y_train = np.array(
+        list(zip(y_train_event.astype(bool), y_train_time.astype(float))),
+        dtype=[("event", "?"), ("time", "<f8")],
+    )
+    y_test = np.array(
+        list(zip(y_test_event.astype(bool), y_test_time.astype(float))),
+        dtype=[("event", "?"), ("time", "<f8")],
+    )
+    return float(ibs_fn(y_train, y_test, surv_pred, times))
+
+
+# ---------- Causal / Counterfactual ----------
+
+def e_value(rr: float) -> float:
+    """E-value (VanderWeele & Ding 2017): min strength of unmeasured
+    confounder needed to explain away an observed RR.
+    """
+    rr = max(rr, 1e-9)
+    if rr >= 1.0:
+        return rr + math.sqrt(rr * (rr - 1))
+    rr_inv = 1.0 / rr
+    return rr_inv + math.sqrt(rr_inv * (rr_inv - 1))
+
+
+def negative_control_check(nc_ate: float, threshold: float = 0.02) -> bool:
+    """Negative-control outcome: ATE on a control outcome should be ~0."""
+    return abs(nc_ate) < threshold
+
+
+def tipping_point(observed_effect: float, ci_half_width: float) -> float:
+    """How much would unmeasured confounding need to shift effect to nullify?"""
+    if abs(observed_effect) <= ci_half_width:
+        return 0.0
+    return float(abs(observed_effect) - ci_half_width)
+
+
+# ---------- Generation fidelity (SynthEHRella triad) ----------
+
+def dim_wise_probability(real_seq: torch.Tensor, synth_seq: torch.Tensor,
+                          vocab_size: int) -> float:
+    """Compare per-token Bernoulli rates between real and synthetic batches.
+
+    Returns mean abs difference (lower = closer match).
+    """
+    real_one_hot = F.one_hot(real_seq, vocab_size).float().mean(dim=(0, 1))
+    synth_one_hot = F.one_hot(synth_seq, vocab_size).float().mean(dim=(0, 1))
+    return float((real_one_hot - synth_one_hot).abs().mean())
+
+
+def mmd_rbf(x: torch.Tensor, y: torch.Tensor, sigma: float = 1.0) -> float:
+    """Maximum Mean Discrepancy with RBF kernel.
+
+    x, y: (B, D) flattened embeddings. Returns MMD^2 (lower = closer).
+    """
+    def rbf(a, b):
+        d = (a.unsqueeze(1) - b.unsqueeze(0)).pow(2).sum(-1)
+        return torch.exp(-d / (2 * sigma ** 2))
+    return float(rbf(x, x).mean() + rbf(y, y).mean() - 2 * rbf(x, y).mean())
+
+
+# ---------- Subgroup fairness ----------
+
+def stratified_metrics(y: np.ndarray, p: np.ndarray,
+                       groups: np.ndarray,
+                       metric_fn: Callable = auroc) -> dict[str, float]:
+    """Compute metric per subgroup (sex, age band, UF region)."""
+    out = {}
+    for g in np.unique(groups):
+        mask = groups == g
+        if mask.sum() > 10:
+            try:
+                out[str(g)] = metric_fn(y[mask], p[mask])
+            except Exception:
+                out[str(g)] = float("nan")
+    return out
+
+
+# ---------- DATASUS split strategies ----------
+
+def temporal_split(events: list[dict], train_until: int = 2022,
+                   val_year: int = 2023):
+    """Temporal split for DATASUS: train ≤2022, val 2023, test 2024+."""
+    train, val, test = [], [], []
+    for e in events:
+        y = e.get("year") or 2020
+        if y <= train_until: train.append(e)
+        elif y == val_year: val.append(e)
+        else: test.append(e)
+    return train, val, test
+
+
+def geographic_split(patients: list[dict], external_ufs: set = None):
+    """Geographic split: train on SE+S, test on N+NE.
+    For DATASUS this is the closest analog to "external validation."
+    """
+    if external_ufs is None:
+        external_ufs = {"AC", "AL", "AP", "AM", "BA", "CE", "MA", "PA",
+                       "PB", "PE", "PI", "RN", "SE", "TO", "RR", "RO"}
+    train, test = [], []
+    for p in patients:
+        uf = next((e.get("uf_code") for e in p.get("events", []) if e.get("uf_code")),
+                  None)
+        (test if uf in external_ufs else train).append(p)
+    return train, test
+
+
+# ---------- Combined eval report ----------
+
+def full_eval_report(y: np.ndarray, p: np.ndarray,
+                     groups_sex: np.ndarray = None,
+                     groups_age: np.ndarray = None,
+                     groups_uf: np.ndarray = None,
+                     n_boot: int = 2000) -> dict:
+    """Generate a full audit-proof report for a binary classification task.
+
+    Returns a dict with point estimates + bootstrap CIs + DCA + fairness.
+    """
+    import torch.nn.functional as F  # local import to keep top clean
+
+    auroc_lo, auroc_med, auroc_hi = bootstrap_ci(y, p, auroc, n_boot)
+    auprc_lo, auprc_med, auprc_hi = bootstrap_ci(y, p, auprc, n_boot)
+    brier_lo, brier_med, brier_hi = bootstrap_ci(y, p, brier, n_boot)
+
+    report = {
+        "n_eval": len(y),
+        "prevalence": float(y.mean()),
+        "auroc": {"point": auroc(y, p), "ci95": [auroc_lo, auroc_hi], "median": auroc_med},
+        "auprc": {"point": auprc(y, p), "ci95": [auprc_lo, auprc_hi], "median": auprc_med},
+        "brier": {"point": brier(y, p), "ci95": [brier_lo, brier_hi], "median": brier_med},
+        "ici": ici(y, p),
+        "decision_curve": decision_curve(y, p),
+    }
+    if groups_sex is not None:
+        report["fairness_sex"] = stratified_metrics(y, p, groups_sex, auroc)
+    if groups_age is not None:
+        report["fairness_age"] = stratified_metrics(y, p, groups_age, auroc)
+    if groups_uf is not None:
+        report["fairness_uf"] = stratified_metrics(y, p, groups_uf, auroc)
+    return report

reference_impl/meds_export.py

@@ -0,0 +1,336 @@
+"""MEDS v0.4.1 exporter for DATASUS — audit-proof, interop-ready.
+
+Verified against:
+  - meds 0.4.1 schemas (DataSchema, CodeMetadataSchema)
+  - https://github.com/Medical-Event-Data-Standard/meds
+  - CLMBR/MOTOR/EHRSHOT/CoMET tokenization conventions
+
+Code conventions (interop-compatible):
+  - Static (time=None): GENDER//, RACE//, UF//, MUN//, ORPHA//
+  - Birth/Death: MEDS_BIRTH, MEDS_DEATH (reserved)
+  - Diagnoses: ICD10//<cid> (NOT CID10// — interop with OHDSI/Athena)
+  - Hospitalization: SIH//ADM, SIH//DIS (numeric_value=LOS_days on DIS)
+  - Procedures: SIGTAP//<10-digit> (Brazil-local namespace)
+  - Drugs (APAC): APAC//<sigtap> (numeric_value=monthly_cost_brl)
+  - Outpatient (BPA-I): BPAI//<sigtap>
+  - Visits: Visit//{IP, OP, ER} (matches CLMBR convention)
+
+Outputs canonical MEDS dataset:
+  /out/
+  ├── data/                       # parquet shards by subject
+  │   ├── shard_0.parquet
+  │   └── ...
+  ├── metadata/
+  │   └── codes.parquet           # REQUIRED: every unique code with description + parent_codes
+  └── dataset_metadata.json       # MEDS dataset metadata
+"""
+from __future__ import annotations
+import os
+import json
+import logging
+from collections import defaultdict, Counter
+from datetime import datetime
+from typing import Iterator
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import meds
+
+log = logging.getLogger("gemeo.cdf.meds_export")
+
+
+def _parse_date(s) -> datetime | None:
+    """Parse date string from various DATASUS formats."""
+    if s is None: return None
+    s = str(s).strip()
+    if not s or s in ("0", "None", "nan"): return None
+    try:
+        if "-" in s:
+            return datetime.strptime(s[:10], "%Y-%m-%d")
+        if len(s) == 8:
+            return datetime.strptime(s, "%Y%m%d")
+    except ValueError:
+        return None
+    return None
+
+
+def _ym(year, month) -> datetime | None:
+    if year is None: return None
+    try:
+        return datetime(int(year), int(month) if month else 1, 1)
+    except (ValueError, TypeError):
+        return None
+
+
+def datasus_patient_to_meds_rows(p: dict, subject_id: int) -> list[tuple]:
+    """Convert one DATASUS patient trajectory to a list of MEDS rows.
+
+    Each row is (subject_id, time, code, numeric_value, text_value).
+    Returns rows ready to write to a parquet shard.
+    """
+    rows = []
+
+    # ---- Static (time=None) ----
+    if p.get("sex"):
+        rows.append((subject_id, None, f"GENDER//{p['sex']}", None, None))
+    # ORPHA is rare-disease specific (parallel to ICD10)
+    for orpha in p.get("orphas", []):
+        rows.append((subject_id, None, f"ORPHA//{orpha}", None, None))
+
+    # ---- Birth (use birth_year as Jan 1) ----
+    birth_year = p.get("birth_year")
+    birth_dt = datetime(int(birth_year), 1, 1) if birth_year else None
+    if birth_dt:
+        rows.append((subject_id, birth_dt, "MEDS_BIRTH", None, None))
+
+    # ---- Events ----
+    for e in p.get("events", []):
+        et = e.get("type")
+
+        if et == "admission":  # SIH-RD
+            t = _ym(e.get("year"), e.get("month")) or _parse_date(e.get("admission_date"))
+            if not t: continue
+            rows.append((subject_id, t, "SIH//ADM", None, None))
+            rows.append((subject_id, t, "Visit//IP", None, None))
+            cid = e.get("cid_princ", "")
+            if cid: rows.append((subject_id, t, f"ICD10//{cid}", None, None))
+            proc = e.get("primary_procedure")
+            if proc: rows.append((subject_id, t, f"SIGTAP//{proc[:10]}", None, None))
+            los = e.get("los_days")
+            disch_dt = _parse_date(e.get("discharge_date")) or t
+            if e.get("death_during_stay"):
+                rows.append((subject_id, disch_dt, "MEDS_DEATH", None, None))
+            else:
+                rows.append((subject_id, disch_dt, "SIH//DIS",
+                            float(los) if los is not None else None, None))
+
+        elif et == "treatment":  # APAC-SIA orphan drug
+            t = _ym(e.get("year"), e.get("month"))
+            if not t: continue
+            cid = e.get("cid", "")
+            if cid: rows.append((subject_id, t, f"ICD10//{cid}", None, None))
+            proc = e.get("procedure_code", "")[:10]
+            if proc:
+                cost = e.get("monthly_cost_brl")
+                rows.append((subject_id, t, f"APAC//{proc}",
+                            float(cost) if cost is not None else None, None))
+
+        elif et == "outpatient_proc":  # BPA-I
+            t = _parse_date(e.get("auth_date")) or _ym(e.get("year"), e.get("month"))
+            if not t: continue
+            cid = e.get("cid", "")
+            if cid: rows.append((subject_id, t, f"ICD10//{cid}", None, None))
+            proc = e.get("procedure_code", "")[:10]
+            if proc:
+                rows.append((subject_id, t, f"BPAI//{proc}", None, None))
+
+        elif et == "death":  # SIM
+            t = _parse_date(e.get("date_of_death")) or _ym(e.get("year"), e.get("month"))
+            if not t: continue
+            rows.append((subject_id, t, "MEDS_DEATH", None, None))
+            cid = (e.get("cause_cid") or e.get("cid_princ") or e.get("cid", ""))
+            if cid: rows.append((subject_id, t, f"ICD10//{cid}", None, None))
+
+    # Sort: nulls first (static), then by time
+    rows.sort(key=lambda r: (r[1] is not None, r[1] or datetime(1900, 1, 1)))
+    return rows
+
+
+def export_to_meds(patients: list[dict], out_dir: str,
+                   shard_size: int = 5000,
+                   dataset_name: str = "GEMEO-DATASUS",
+                   version: str = "v13"):
+    """Export a list of DATASUS patient trajectories to MEDS v0.4.1 format.
+
+    Parameters
+    ----------
+    patients : list of dict
+        Each dict must have: patient_id, sex, birth_year, orphas (list),
+        events (list of dicts with 'type', 'year', 'month', etc.)
+    out_dir : str
+        Output directory (will create data/ and metadata/ subdirs)
+    shard_size : int
+        Number of subjects per parquet shard
+    """
+    os.makedirs(f"{out_dir}/data", exist_ok=True)
+    os.makedirs(f"{out_dir}/metadata", exist_ok=True)
+
+    log.info(f"Exporting {len(patients)} patients to MEDS at {out_dir}")
+
+    # Map patient_id (string hash) → int64 subject_id (MEDS requires int64)
+    pid_to_sid = {p["patient_id"]: i for i, p in enumerate(patients)}
+
+    # ---- Stream rows ----
+    all_codes = Counter()
+    shard_idx = 0
+    shard_rows = []
+    n_events = 0
+    n_subjects = 0
+
+    for p in patients:
+        sid = pid_to_sid[p["patient_id"]]
+        rows = datasus_patient_to_meds_rows(p, sid)
+        shard_rows.extend(rows)
+        n_events += len(rows)
+        n_subjects += 1
+        for r in rows:
+            all_codes[r[2]] += 1
+        # Write shard when full
+        if n_subjects % shard_size == 0 and shard_rows:
+            _write_shard(shard_rows, f"{out_dir}/data/shard_{shard_idx}.parquet")
+            shard_idx += 1
+            shard_rows = []
+
+    # Write remaining
+    if shard_rows:
+        _write_shard(shard_rows, f"{out_dir}/data/shard_{shard_idx}.parquet")
+
+    log.info(f"  wrote {shard_idx + 1} data shards, {n_events} rows, {n_subjects} subjects")
+
+    # ---- codes.parquet (REQUIRED in MEDS v0.4) ----
+    code_rows = []
+    for code, count in all_codes.most_common():
+        # parent_codes: empty for Brazil-local namespaces; populated for ICD10 -> SNOMED if mapped
+        parent_codes = _get_parent_codes(code)
+        code_rows.append({
+            "code": code,
+            "description": _get_description(code, count),
+            "parent_codes": parent_codes,
+        })
+    code_table = pa.Table.from_pylist(code_rows, schema=meds.CodeMetadataSchema.schema())
+    pq.write_table(code_table, f"{out_dir}/metadata/codes.parquet")
+    log.info(f"  wrote metadata/codes.parquet ({len(code_rows)} unique codes)")
+
+    # ---- dataset_metadata.json ----
+    md = {
+        "dataset_name": dataset_name,
+        "dataset_version": version,
+        "etl_name": "gemeo.cdf.meds_export",
+        "etl_version": "1.0.0",
+        "meds_version": meds.__version__,
+        "n_subjects": n_subjects,
+        "n_events": n_events,
+        "n_unique_codes": len(all_codes),
+        "top_codes": dict(all_codes.most_common(30)),
+    }
+    with open(f"{out_dir}/dataset_metadata.json", "w") as f:
+        json.dump(md, f, indent=2, default=str)
+    log.info(f"  wrote dataset_metadata.json")
+
+    return md
+
+
+def _write_shard(rows: list[tuple], path: str):
+    """Write a list of (subject_id, time, code, numeric_value, text_value) to parquet."""
+    if not rows: return
+    # Build columnar arrays
+    subject_id = pa.array([r[0] for r in rows], type=pa.int64())
+    time = pa.array([r[1] for r in rows], type=pa.timestamp("us"))
+    code = pa.array([r[2] for r in rows], type=pa.string())
+    numeric_value = pa.array([r[3] for r in rows], type=pa.float32())
+    text_value = pa.array([r[4] for r in rows], type=pa.large_string())
+    table = pa.Table.from_arrays(
+        [subject_id, time, code, numeric_value, text_value],
+        names=["subject_id", "time", "code", "numeric_value", "text_value"],
+    )
+    # Validate against MEDS schema
+    expected_schema = meds.DataSchema.schema()
+    # Cast if needed
+    table = table.cast(expected_schema, safe=False)
+    pq.write_table(table, path, compression="zstd")
+
+
+# Brazilian-specific mapping tables (extend as needed)
+ICD10_CHAPTERS = {
+    "A": "Certain infectious and parasitic diseases",
+    "B": "Certain infectious and parasitic diseases",
+    "C": "Neoplasms",
+    "D": "Neoplasms / Diseases of the blood and immune",
+    "E": "Endocrine, nutritional and metabolic diseases",
+    "F": "Mental, Behavioral and Neurodevelopmental disorders",
+    "G": "Diseases of the nervous system",
+    "H": "Diseases of the eye / ear",
+    "I": "Diseases of the circulatory system",
+    "J": "Diseases of the respiratory system",
+    "K": "Diseases of the digestive system",
+    "L": "Diseases of the skin and subcutaneous tissue",
+    "M": "Diseases of the musculoskeletal system",
+    "N": "Diseases of the genitourinary system",
+    "O": "Pregnancy, childbirth and the puerperium",
+    "P": "Certain conditions originating in the perinatal period",
+    "Q": "Congenital malformations, deformations and chromosomal abnormalities",
+    "R": "Symptoms, signs and abnormal clinical and laboratory findings",
+    "S": "Injury, poisoning and certain other consequences of external causes",
+    "T": "Injury, poisoning and certain other consequences of external causes",
+    "V": "External causes of morbidity",
+    "W": "External causes of morbidity",
+    "X": "External causes of morbidity",
+    "Y": "External causes of morbidity",
+    "Z": "Factors influencing health status and contact with health services",
+}
+
+
+def _get_description(code: str, count: int) -> str:
+    """Generate a brief description for a code (used in codes.parquet)."""
+    if code in ("MEDS_BIRTH",): return "Birth event (reserved)"
+    if code in ("MEDS_DEATH",): return "Death event (reserved)"
+    parts = code.split("//")
+    if len(parts) < 2: return f"Unknown code (n={count})"
+    domain, val = parts[0], "//".join(parts[1:])
+    if domain == "GENDER": return f"Patient sex = {val}"
+    if domain == "ORPHA": return f"Orphanet rare disease {val}"
+    if domain == "ICD10":
+        ch = ICD10_CHAPTERS.get(val[0], "Unknown chapter")
+        return f"ICD-10 {val} ({ch})"
+    if domain == "SIH": return f"SIH hospitalization {val}"
+    if domain == "Visit": return f"Visit type {val}"
+    if domain == "SIGTAP": return f"SIGTAP procedure {val}"
+    if domain == "APAC": return f"APAC orphan-drug authorization {val}"
+    if domain == "BPAI": return f"BPA-I outpatient procedure {val}"
+    if domain == "UF": return f"Residence UF {val}"
+    return f"{domain} code {val}"
+
+
+def _get_parent_codes(code: str) -> list[str]:
+    """Return parent codes for ontology hierarchy (currently minimal)."""
+    parts = code.split("//")
+    if len(parts) < 2: return []
+    domain, val = parts[0], "//".join(parts[1:])
+    parents = []
+    if domain == "ICD10" and len(val) >= 3:
+        # ICD-10 chapter as parent
+        chapter = val[0]
+        if chapter in ICD10_CHAPTERS:
+            parents.append(f"ICD10//chapter_{chapter}")
+        # 3-char prefix as parent (e.g., E84.0 → E84)
+        if "." in val:
+            parents.append(f"ICD10//{val.split('.')[0]}")
+        elif len(val) > 3:
+            parents.append(f"ICD10//{val[:3]}")
+    if domain == "SIGTAP" and len(val) >= 4:
+        # 4-digit group as parent (SIGTAP 10-digit → 4-digit group)
+        parents.append(f"SIGTAP//group_{val[:4]}")
+    return parents
+
+
+def load_meds_dataset(meds_dir: str) -> dict:
+    """Load a MEDS dataset back from parquet for inspection or downstream processing."""
+    import glob
+    shards = sorted(glob.glob(f"{meds_dir}/data/*.parquet"))
+    tables = [pq.read_table(p) for p in shards]
+    data = pa.concat_tables(tables) if tables else None
+    codes = pq.read_table(f"{meds_dir}/metadata/codes.parquet")
+    md = json.load(open(f"{meds_dir}/dataset_metadata.json"))
+    return {"data": data, "codes": codes, "metadata": md}
+
+
+if __name__ == "__main__":
+    # Quick test on real patient data
+    logging.basicConfig(level=logging.INFO,
+                        format="%(asctime)s %(levelname)s %(message)s")
+    PATIENTS = "/tmp/datasus_patient_trajectories_v2.json"
+    if os.path.exists(PATIENTS):
+        patients = json.load(open(PATIENTS))[:50]   # 50 patients smoke test
+        md = export_to_meds(patients, "/tmp/meds_smoke_test")
+        print("\n=== smoke test result ===")
+        print(json.dumps(md, indent=2, default=str))

reference_impl/primekg_attention.py

@@ -0,0 +1,262 @@
+"""PrimeKG cross-attention — graph-RAG into the Diffusion Forcing denoiser.
+
+Now uses REAL EDGES from raras-app/data/graph-ml/hetero_graph.json:
+  - disease → has_phenotype → phenotype  (curated phenotype linkage)
+  - disease → associated_with → gene     (causal gene evidence)
+  - gene → interacts_with → gene         (PPI network)
+  - phenotype → is_a → phenotype         (HPO ontology)
+
+Ego-subgraph BFS:
+  1. Start from disease node (ORPHA → PrimeKG index)
+  2. 1-hop: pull connected phenotypes (top-K by edge weight or count)
+  3. 1-hop: pull connected genes
+  4. 2-hop: gene→gene neighbors (interacting partners)
+  5. Concatenate fused embeddings of all selected nodes → cross-attn context
+
+Falls back to cosine-similarity if graph not loaded.
+
+White-space architecture (May 2026):
+  - EHRWorld, CLARITY, Time-Aware G-Transformer all skip KG conditioning
+  - PhenoKG/RareNet use KG for RETRIEVAL (rare disease diagnosis)
+  - We use it for GENERATION (counterfactual trajectory completion)
+"""
+from __future__ import annotations
+import os
+import json
+import logging
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+log = logging.getLogger("gemeo.cdf.kg")
+
+# Try raras-app paths first (richer, including hetero_graph edges + node_texts)
+RARAS_KG_DIR = "/Users/dimas/raras-app/data/graph-ml"
+LOCAL_KG_DIR = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
+
+
+def _kg_path(name: str) -> str:
+    """Prefer raras-app path if available, fall back to local fp16."""
+    raras = os.path.join(RARAS_KG_DIR, name)
+    if os.path.exists(raras):
+        return raras
+    local = os.path.join(LOCAL_KG_DIR, name)
+    return local if os.path.exists(local) else None
+
+
+@lru_cache(maxsize=1)
+def load_kg(prefer_raras: bool = True) -> dict | None:
+    """Load PrimeKG: fused embeddings + node ids + edges + texts.
+
+    Returns dict:
+        emb        : {kind: torch.Tensor(N, 3072)}
+        idx2id     : {kind: {pos: id_str}}
+        id2idx     : {kind: {id_str: pos}}
+        edges      : {edge_type: {'src': [...], 'dst': [...]}}
+        adj        : {edge_type: {src_idx: [dst_idx, ...]}}  -- precomputed
+        texts      : {kind: [str, ...]}  -- aligned to position
+        num_nodes  : {kind: int}
+    """
+    # Try raras-app full file first, then local fp16
+    emb_path = (os.path.join(RARAS_KG_DIR, "fused_embeddings.npz")
+                if prefer_raras and os.path.exists(os.path.join(RARAS_KG_DIR, "fused_embeddings.npz"))
+                else _kg_path("fused_embeddings_fp16.npz"))
+    if not emb_path or not os.path.exists(emb_path):
+        log.warning("PrimeKG fused embeddings not found")
+        return None
+
+    nids_path = _kg_path("node_ids.json")
+    graph_path = _kg_path("hetero_graph.json")
+    texts_path = _kg_path("node_texts.json")
+
+    fz = np.load(emb_path)
+    nids = json.load(open(nids_path)) if nids_path else {}
+    graph = json.load(open(graph_path)) if graph_path else {"edges": {}, "num_nodes": {}}
+    texts = json.load(open(texts_path)) if texts_path else {}
+
+    out = {"emb": {}, "id2idx": {}, "idx2id": {}, "edges": {}, "adj": {},
+           "texts": texts, "num_nodes": graph.get("num_nodes", {})}
+
+    for kind in ("disease", "phenotype", "gene"):
+        if kind in fz.files:
+            out["emb"][kind] = torch.from_numpy(fz[kind].astype(np.float32))
+            if kind in nids:
+                out["idx2id"][kind] = {int(k): v for k, v in nids[kind].items()}
+                out["id2idx"][kind] = {v: int(k) for k, v in nids[kind].items()}
+
+    # Build adjacency from edges
+    for edge_type, edata in graph.get("edges", {}).items():
+        adj = {}
+        srcs = edata.get("src", []) if isinstance(edata, dict) else []
+        dsts = edata.get("dst", []) if isinstance(edata, dict) else []
+        for s, d in zip(srcs, dsts):
+            adj.setdefault(int(s), []).append(int(d))
+        out["adj"][edge_type] = adj
+        out["edges"][edge_type] = edata
+
+    log.info(f"  KG loaded from {emb_path}")
+    log.info(f"  disease={out['emb'].get('disease', torch.empty(0)).shape}, "
+             f"phenotype={out['emb'].get('phenotype', torch.empty(0)).shape}, "
+             f"gene={out['emb'].get('gene', torch.empty(0)).shape}")
+    log.info(f"  edges: {list(out['edges'].keys())}")
+    return out
+
+
+def ego_subgraph_real(orpha_code: str, k_pheno: int = 16, k_gene: int = 16,
+                      k_gene_2hop: int = 0, kg: dict | None = None) -> torch.Tensor:
+    """BFS ego-subgraph using REAL PrimeKG edges (not cosine similarity).
+
+    Returns concatenated embeddings (N, 3072) where:
+      - 1 disease node (the query)
+      - up to k_pheno phenotype nodes (direct edges)
+      - up to k_gene gene nodes (direct edges)
+      - up to k_gene_2hop gene-gene 2-hop neighbors
+
+    Falls back to cosine similarity if no edges available.
+    """
+    if kg is None:
+        kg = load_kg()
+    if kg is None or "disease" not in kg["emb"]:
+        return None
+
+    d_id = kg["id2idx"]["disease"].get(str(orpha_code))
+    if d_id is None:
+        return None
+
+    d_emb = kg["emb"]["disease"][d_id]
+    nodes = [d_emb.unsqueeze(0)]
+
+    # Phenotype neighbors (via disease__has_phenotype__phenotype)
+    adj = kg["adj"].get("disease__has_phenotype__phenotype", {})
+    pheno_neighbors = adj.get(d_id, [])
+    if pheno_neighbors and "phenotype" in kg["emb"]:
+        pheno_neighbors = pheno_neighbors[:k_pheno]
+        nodes.append(kg["emb"]["phenotype"][pheno_neighbors])
+    elif "phenotype" in kg["emb"]:
+        # Fallback: cosine similarity
+        pool = kg["emb"]["phenotype"]
+        sim = F.cosine_similarity(d_emb.unsqueeze(0), pool, dim=-1)
+        top = sim.topk(min(k_pheno, pool.size(0))).indices
+        nodes.append(pool[top])
+
+    # Gene neighbors (via disease__associated_with__gene)
+    g_adj = kg["adj"].get("disease__associated_with__gene", {})
+    gene_neighbors = g_adj.get(d_id, [])
+    if gene_neighbors and "gene" in kg["emb"]:
+        gene_neighbors = gene_neighbors[:k_gene]
+        nodes.append(kg["emb"]["gene"][gene_neighbors])
+
+        # 2-hop: gene-gene neighbors of the genes we just pulled
+        if k_gene_2hop > 0:
+            gg_adj = kg["adj"].get("gene__interacts_with__gene", {})
+            seen = set(gene_neighbors)
+            second_hop = []
+            for g in gene_neighbors:
+                for g2 in gg_adj.get(g, []):
+                    if g2 not in seen:
+                        second_hop.append(g2)
+                        seen.add(g2)
+                        if len(second_hop) >= k_gene_2hop: break
+                if len(second_hop) >= k_gene_2hop: break
+            if second_hop:
+                nodes.append(kg["emb"]["gene"][second_hop])
+    elif "gene" in kg["emb"]:
+        pool = kg["emb"]["gene"]
+        sim = F.cosine_similarity(d_emb.unsqueeze(0), pool, dim=-1)
+        top = sim.topk(min(k_gene, pool.size(0))).indices
+        nodes.append(pool[top])
+
+    return torch.cat(nodes, dim=0)
+
+
+# Keep old API name for backward compat
+ego_subgraph = ego_subgraph_real
+
+
+class KGCrossAttention(nn.Module):
+    """Cross-attention from sequence (B, T, d_model) to KG ego (B, N, d_model)."""
+    def __init__(self, d_model: int, n_heads: int = 8, dropout: float = 0.1):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.norm_q = nn.LayerNorm(d_model)
+        self.norm_kv = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x_seq: torch.Tensor, x_kg: torch.Tensor) -> torch.Tensor:
+        B, T, D = x_seq.shape
+        _, N, _ = x_kg.shape
+        q = self.q_proj(self.norm_q(x_seq))
+        kv = self.kv_proj(self.norm_kv(x_kg))
+        k, v = kv.chunk(2, dim=-1)
+        q = q.reshape(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.reshape(B, N, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.reshape(B, N, self.n_heads, self.head_dim).transpose(1, 2)
+        out = F.scaled_dot_product_attention(
+            q, k, v, dropout_p=self.dropout.p if self.training else 0.0)
+        out = out.transpose(1, 2).reshape(B, T, D)
+        return x_seq + self.dropout(self.out_proj(out))
+
+
+class KGProjector(nn.Module):
+    """Project 3072-d KG embeddings to d_model with LayerNorm."""
+    def __init__(self, kg_dim: int, d_model: int):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(kg_dim, d_model),
+            nn.GELU(),
+            nn.LayerNorm(d_model),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x)
+
+
+def build_kg_batch(orpha_strings: list[str], d_model: int,
+                   projector: KGProjector,
+                   k_pheno: int = 16, k_gene: int = 16,
+                   k_gene_2hop: int = 0) -> torch.Tensor:
+    """Build (B, N, d_model) batched KG context for a batch of patient ORPHAs.
+
+    Falls back to zero context for missing ORPHAs.
+    """
+    kg = load_kg()
+    if kg is None:
+        return torch.zeros(len(orpha_strings), 1, d_model,
+                          device=next(projector.parameters()).device)
+    N = 1 + k_pheno + k_gene + k_gene_2hop
+    egos = []
+    for orpha in orpha_strings:
+        e = ego_subgraph_real(orpha, k_pheno, k_gene, k_gene_2hop, kg)
+        if e is None:
+            e = torch.zeros(N, kg["emb"]["disease"].size(-1))
+        elif e.size(0) < N:
+            pad = torch.zeros(N - e.size(0), e.size(-1))
+            e = torch.cat([e, pad], dim=0)
+        egos.append(e[:N])
+    egos = torch.stack(egos, dim=0)
+    return projector(egos.to(next(projector.parameters()).device))
+
+
+def precompute_kg_for_dataset(orpha_codes: list[str], projector: KGProjector,
+                              k_pheno: int = 16, k_gene: int = 16,
+                              batch_size: int = 32) -> torch.Tensor:
+    """Pre-compute KG context for an entire dataset in batches.
+
+    Returns (N_patients, kg_nodes, d_model) tensor on projector device.
+    Saves to disk-cacheable format.
+    """
+    out = []
+    for i in range(0, len(orpha_codes), batch_size):
+        batch = orpha_codes[i:i + batch_size]
+        ctx = build_kg_batch(batch, projector.proj[0].out_features,
+                             projector, k_pheno, k_gene)
+        out.append(ctx.cpu())
+    return torch.cat(out, dim=0)

reference_impl/sample.py

@@ -0,0 +1,230 @@
+"""Sampling primitives for CDF: AR mode, denoise mode, counterfactual rollouts.
+
+Diffusion Forcing flexibility — the same model handles:
+
+  AR mode:
+    Sigma_future = 1, sigma_past = 0. Roll forward like an autoregressive
+    transformer but with per-token noise control.
+
+  Denoise mode (bidirectional):
+    Sigma low everywhere. Run k denoise steps, model fills the whole sequence.
+
+  Counterfactual mode (the TTE primitive):
+    Sigma=0 on observed tokens (clamp them clean), sigma=1 on tokens to
+    generate. Condition on (cohort, intervention_action_id). Sample N times,
+    compare distributions of outcome tokens.
+
+CFG (classifier-free guidance) wraps any mode:
+    logits_g = (1 + gamma) * logits(c) - gamma * logits(null_c)
+
+Shortcut Forcing (Dreamer 4) reduces denoise steps from 32-64 to 4 via
+distilled student model — implemented in distill.py.
+"""
+from __future__ import annotations
+import logging
+
+import torch
+import torch.nn.functional as F
+
+from .diffusion_forcing import CDFTransformer
+
+log = logging.getLogger("gemeo.cdf.sample")
+
+
+@torch.no_grad()
+def sample_denoise(
+    model: CDFTransformer,
+    cond: torch.Tensor,
+    *,
+    seed_prefix: torch.Tensor | None = None,
+    observed_mask: torch.Tensor | None = None,    # (B, T) True = clamped clean
+    action: torch.Tensor | None = None,
+    gamma: float = 2.0,
+    n_steps: int = 32,
+    null_cond: int = 0,
+    schedule: str = "cosine",
+) -> torch.Tensor:
+    """Denoise-mode sampling: fully-masked sequence + iterative refinement.
+
+    Supports:
+      - seed_prefix: clean tokens kept at sigma=0 for positions [0, L)
+      - observed_mask: arbitrary positions to clamp (counterfactual mode)
+      - CFG via (cond, null_cond) pair
+    """
+    cfg = model.cfg
+    device = cond.device
+    B = cond.size(0)
+    T = cfg.max_seq_len
+
+    # Init with MASK
+    x = torch.full((B, T), cfg.mask_token, device=device, dtype=torch.long)
+    fixed_mask = torch.zeros(B, T, dtype=torch.bool, device=device)
+    if seed_prefix is not None:
+        L = seed_prefix.size(1)
+        x[:, :L] = seed_prefix
+        fixed_mask[:, :L] = True
+    if observed_mask is not None:
+        fixed_mask |= observed_mask
+
+    # Build noise schedule
+    if schedule == "cosine":
+        # smooth cosine from 1 -> 0
+        ts = torch.cos(torch.linspace(0, torch.pi/2, n_steps+1, device=device))
+    else:
+        ts = torch.linspace(1.0, 0.0, n_steps+1, device=device)
+
+    null = torch.full_like(cond, null_cond)
+    null_action = (torch.full_like(action, cfg.n_latent_actions)
+                   if action is not None and cfg.use_latent_action else None)
+
+    for k in range(n_steps):
+        # Per-token sigma: fixed positions at 0, dynamic positions at ts[k]
+        sigma = torch.where(fixed_mask, torch.zeros_like(ts[k:k+1]).expand(B, T),
+                            torch.full((B, T), ts[k].item(), device=device))
+        logits_c = model(x, sigma, cond, action)
+        if gamma > 0:
+            logits_n = model(x, sigma, null, null_action)
+            logits = (1 + gamma) * logits_c - gamma * logits_n
+        else:
+            logits = logits_c
+        logits[:, :, cfg.mask_token] = -1e9
+
+        probs = F.softmax(logits, dim=-1)
+        confs, preds = probs.max(dim=-1)
+
+        # Confidence-based remasking: reveal top-(1 - ts[k+1]) fraction of free tokens
+        t_next = ts[k+1].item()
+        target_kept = int(round((1 - t_next) * T))
+        revealed = (x != cfg.mask_token) | fixed_mask
+        already = revealed.sum(dim=-1)
+        new_x = x.clone()
+        for b in range(B):
+            need = max(0, target_kept - int(already[b].item()))
+            if need == 0:
+                continue
+            confs_b = torch.where(revealed[b], torch.full_like(confs[b], -1e9), confs[b])
+            topi = confs_b.topk(need).indices
+            new_x[b, topi] = preds[b, topi]
+        x = new_x
+
+    # Final cleanup
+    mask_left = x == cfg.mask_token
+    if mask_left.any():
+        sigma_final = torch.zeros(B, T, device=device)
+        logits_c = model(x, sigma_final, cond, action)
+        if gamma > 0:
+            logits_n = model(x, sigma_final, null, null_action)
+            logits = (1 + gamma) * logits_c - gamma * logits_n
+        else:
+            logits = logits_c
+        logits[:, :, cfg.mask_token] = -1e9
+        preds = logits.argmax(-1)
+        x = torch.where(mask_left, preds, x)
+    return x
+
+
+@torch.no_grad()
+def sample_ar(
+    model: CDFTransformer,
+    cond: torch.Tensor,
+    prefix: torch.Tensor,
+    *,
+    action: torch.Tensor | None = None,
+    max_new: int = 50,
+    temperature: float = 1.0,
+    gamma: float = 0.0,
+    null_cond: int = 0,
+) -> torch.Tensor:
+    """AR-mode sampling: future tokens at sigma=1, past at sigma=0.
+
+    Faster than denoise mode when you only want to continue a prefix.
+    """
+    cfg = model.cfg
+    device = cond.device
+    B = cond.size(0)
+    x = prefix.clone().to(device)
+    if x.dim() == 1: x = x.unsqueeze(0)
+    null = torch.full_like(cond, null_cond)
+    null_action = (torch.full_like(action, cfg.n_latent_actions)
+                   if action is not None and cfg.use_latent_action else None)
+
+    for _ in range(max_new):
+        T_now = x.size(1)
+        if T_now >= cfg.max_seq_len:
+            break
+        # Pad with MASK
+        x_pad = torch.cat([x, torch.full((B, 1), cfg.mask_token,
+                                          device=device, dtype=torch.long)], dim=1)
+        sigma = torch.zeros(B, T_now + 1, device=device)
+        sigma[:, -1] = 1.0
+        a_pad = None
+        if action is not None and cfg.use_latent_action:
+            a_pad = torch.cat([action[:, :T_now],
+                               torch.full((B, 1), cfg.n_latent_actions,
+                                          device=device, dtype=torch.long)], dim=1)
+        logits = model(x_pad, sigma, cond, a_pad)
+        if gamma > 0:
+            logits_n = model(x_pad, sigma, null, null_action)
+            logits = (1 + gamma) * logits - gamma * logits_n
+        logits[:, :, cfg.mask_token] = -1e9
+        p = F.softmax(logits[:, -1] / max(temperature, 1e-3), dim=-1)
+        nxt = torch.multinomial(p, 1)
+        x = torch.cat([x, nxt], dim=1)
+    return x
+
+
+@torch.no_grad()
+def counterfactual_rollout(
+    model: CDFTransformer,
+    seed_prefix: torch.Tensor,
+    treatment_cond: int,
+    untreated_cond: int,
+    *,
+    treatment_action: int | None = None,
+    untreated_action: int | None = None,
+    n_samples: int = 100,
+    gamma: float = 2.0,
+    n_steps: int = 32,
+) -> dict:
+    """Sample paired counterfactual trajectories under treatment vs no-treatment.
+
+    Two ways to specify the intervention:
+      - via cond id (cohort-level): treatment_cond / untreated_cond
+      - via latent action id (per-token): treatment_action / untreated_action
+    """
+    cfg = model.cfg
+    device = next(model.parameters()).device
+    seed = seed_prefix.unsqueeze(0).expand(n_samples, -1).to(device)
+    T = cfg.max_seq_len
+
+    cond_tx = torch.full((n_samples,), treatment_cond, device=device, dtype=torch.long)
+    cond_null = torch.full((n_samples,), untreated_cond, device=device, dtype=torch.long)
+
+    action_tx = action_null = None
+    if cfg.use_latent_action:
+        action_tx = torch.full((n_samples, T),
+                               treatment_action if treatment_action is not None
+                               else cfg.n_latent_actions,
+                               device=device, dtype=torch.long)
+        action_null = torch.full((n_samples, T),
+                                 untreated_action if untreated_action is not None
+                                 else cfg.n_latent_actions,
+                                 device=device, dtype=torch.long)
+
+    traj_tx = sample_denoise(model, cond_tx, seed_prefix=seed,
+                              action=action_tx, gamma=gamma, n_steps=n_steps)
+    traj_null = sample_denoise(model, cond_null, seed_prefix=seed,
+                                action=action_null, gamma=gamma, n_steps=n_steps)
+    return {
+        "traj_treated": traj_tx, "traj_untreated": traj_null,
+        "n": n_samples, "treatment_cond": treatment_cond,
+        "untreated_cond": untreated_cond, "gamma": gamma,
+    }
+
+
+def outcome_rate(traj: torch.Tensor, target_ids: list[int]) -> float:
+    if not target_ids:
+        return 0.0
+    target = torch.tensor(target_ids, device=traj.device)
+    has = (traj.unsqueeze(-1) == target).any(dim=(-1, -2))
+    return has.float().mean().item()

reference_impl/wsd_scheduler.py

@@ -0,0 +1,72 @@
+"""WSD (Warmup-Stable-Decay) LR scheduler — manual implementation.
+
+Per MiniCPM (Hu et al. 2024) and the data-constrained scaling literature:
+  Phase 1 (warmup, 1-5% of total_steps): linear 0 → peak_lr
+  Phase 2 (stable, 60-80%):              constant peak_lr
+  Phase 3 (decay, 10-25%):               linear or 1/sqrt to peak_lr * 0.1
+
+Beats cosine for:
+  - data-limited regimes (we can extend stable phase if loss still falls)
+  - continue-pretrain (sharp decay enables clean fine-tune handoff)
+"""
+from __future__ import annotations
+import math
+import torch
+from torch.optim.lr_scheduler import LambdaLR
+
+
+def wsd_lr_schedule(step: int, total_steps: int,
+                    warmup_steps: int = 500,
+                    stable_frac: float = 0.80,
+                    decay_frac: float = 0.15,
+                    min_lr_ratio: float = 0.1,
+                    decay_type: str = "linear") -> float:
+    """Return LR multiplier in [min_lr_ratio, 1.0] for a given step."""
+    if step < warmup_steps:
+        return step / max(1, warmup_steps)
+    # remainder of steps after warmup
+    remaining = total_steps - warmup_steps
+    if remaining <= 0:
+        return 1.0
+    stable_steps = int(stable_frac * remaining)
+    decay_steps = int(decay_frac * remaining)
+    pos = step - warmup_steps
+    if pos < stable_steps:
+        return 1.0
+    decay_pos = pos - stable_steps
+    if decay_pos >= decay_steps:
+        return min_lr_ratio
+    progress = decay_pos / max(1, decay_steps)
+    if decay_type == "linear":
+        return 1.0 - (1.0 - min_lr_ratio) * progress
+    elif decay_type == "cosine":
+        return min_lr_ratio + 0.5 * (1 - min_lr_ratio) * (1 + math.cos(math.pi * progress))
+    elif decay_type == "inv_sqrt":
+        return max(min_lr_ratio, 1.0 / math.sqrt(1 + progress * 10))
+    else:
+        raise ValueError(f"unknown decay_type: {decay_type}")
+
+
+def get_wsd_scheduler(optimizer: torch.optim.Optimizer,
+                      total_steps: int,
+                      warmup_steps: int = 500,
+                      stable_frac: float = 0.80,
+                      decay_frac: float = 0.15,
+                      min_lr_ratio: float = 0.1,
+                      decay_type: str = "linear") -> LambdaLR:
+    """Build a LambdaLR scheduler with WSD schedule."""
+    def fn(step):
+        return wsd_lr_schedule(step, total_steps, warmup_steps,
+                              stable_frac, decay_frac, min_lr_ratio, decay_type)
+    return LambdaLR(optimizer, lr_lambda=fn)
+
+
+if __name__ == "__main__":
+    # Visualize the schedule
+    total = 10000
+    warmup = 500
+    print(f"WSD schedule preview: total={total}, warmup={warmup}, stable=80%, decay=15%")
+    print(f"  step    lr_mult")
+    for s in [0, 250, 500, 1000, 5000, 8000, 8500, 9000, 9500, 9800, 9999]:
+        m = wsd_lr_schedule(s, total, warmup, 0.80, 0.15, 0.1, "linear")
+        print(f"  {s:>5}   {m:.4f}")