README.md

---
license: cc-by-nc-4.0
language: [pt, en]
tags:
  - world-model
  - patient-digital-twin
  - rare-disease
  - diffusion-forcing
  - knowledge-graph
  - brazilian-sus
  - datasus
  - primekg
library_name: pytorch
pipeline_tag: time-series-forecasting
extra_gated_prompt: >-
  Research only. Not a medical device. No clinical use without physician
  oversight and applicable regulatory clearance.
extra_gated_fields:
  Name: text
  Affiliation: text
  Intended use: text
  I agree to non-clinical research use only: checkbox
---

# GEMEO/SUS — Recurrence-Aware Patient World Model

> The flagship instance of [**GEMEO Architecture v2.0**](https://huggingface.co/Raras-AI/gemeo-arch) —
> a patient world model for rare disease, trained on real Brazilian SUS
> (DATASUS) data. Predicts novel clinical events and long-context outcomes,
> grounded in a biomedical knowledge graph. In the lineage of Dreamer
> (Hafner 2025), Diffusion Forcing (Chen, NeurIPS 2024), Sora and Genie.

**Family:** [`gemeo-arch`](https://huggingface.co/Raras-AI/gemeo-arch) (architecture) · **`gemeo-sus`** (this, flagship) · [`gemeo-twin-stack`](https://huggingface.co/Raras-AI/gemeo-twin-stack) (6-mode app layer) · [`rarebench-br-trajectory`](https://huggingface.co/datasets/Raras-AI/rarebench-br-trajectory) (benchmark)

## State-of-the-art results

Evaluated on the public [RareBench-BR Trajectory v2](https://huggingface.co/datasets/Raras-AI/rarebench-br-trajectory) benchmark + a new-onset task, with mandatory baselines on the same candidate space and 95% bootstrap CIs. GEMEO leads on every novelty and long-context task:

| Task | GEMEO | Strong baseline | Margin |
|---|---:|---:|---:|
| **New-onset prediction** (Top-1) | **53.7%** | 38.2% (frequency) | **+15.5 pp** |
| **Will-change** (AUROC) | **0.906** | 0.889 (count-based) | +0.017 |
| **Transition-within-12mo** (AUROC) | **0.827** | 0.790 (count-based) | +0.037 |
| **Treatment discontinuation** (AUROC) | **0.838** | 0.696 (count-based) | **+0.142** |

Long-context outcomes — especially **treatment discontinuation** (dropout drives bad outcomes in rare disease) — are where the world model's learned representation pulls clearly ahead of count-based methods, exactly as the EHR literature predicts for context-rich tasks (arXiv 2511.00782). The recurrence-aware objective makes the model predict *novel* events, not repeats.

## Architecture

```
GEMEO/SUS (19.97M params)
├── Token embedding (tied with LM head)
├── PositionalFeatureEmbed(age, calendar_year, position)
├── 8 × Transformer blocks (SwiGLU + RMSNorm + RoPE + AdaLN-Zero)
│   └── Gated PrimeKG cross-attention (tanh(α), α init=0)
└── Tied LM head
```

- **Backbone:** Causal Diffusion Forcing transformer (per-token σ; Chen et al., NeurIPS 2024).
- **Training objective:** recurrence-weighted loss (RAVEN, arXiv 2603.24562) — first occurrences carry full weight, so the model learns genuinely new events.
- **Conditioning:** gated cross-attention to a real PrimeKG ego-subgraph (disease–gene, disease–phenotype edges).
- **Recipe:** warm-start, WSD LR schedule, bf16, single H100, ~5 min, ≈ $0.40.
- MEDS v0.4.1 substrate · 42,265 DATASUS rare-disease trajectories.

Subgroup fairness is clean across pediatric / adult / elderly bands. The bundled `cdf_v7_10digit_rbt.pt` is the 10-digit variant used for the RBT-v2 transition evaluation; `cdf_v6_raven.pt` is the recurrence-aware flagship.

## Usage

```python
import torch, sys; sys.path.append("src")
import torch.nn as nn
from diffusion_forcing_v13 import CDFv13Transformer, CDFv13Config

class PositionalFeatureEmbed(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.age_proj=nn.Linear(1,d//4); self.year_proj=nn.Linear(1,d//4)
        self.pos_proj=nn.Linear(1,d//4); self.combine=nn.Linear(3*(d//4),d); self.norm=nn.LayerNorm(d)
    def forward(self, ages, years, positions):
        a=ages.clamp(0,100)/100; y=(years-2010).clamp(0,20)/20; p=(positions/512).clamp(0,1)
        e=torch.cat([self.age_proj(a.unsqueeze(-1)), self.year_proj(y.unsqueeze(-1)),
                     self.pos_proj(p.unsqueeze(-1))], -1)
        return self.norm(self.combine(e))

ck = torch.load("cdf_v6_raven.pt", map_location="cpu", weights_only=False)
cfg = CDFv13Config(**{k:v for k,v in ck["config"].items() if k in CDFv13Config.__dataclass_fields__})
model = CDFv13Transformer(cfg); model.load_state_dict(ck["model_state"])
pfe = PositionalFeatureEmbed(cfg.d_model); pfe.load_state_dict(ck["pos_feat_state"])
```

## Scope

GEMEO leads on novelty (new-onset) and long-context outcomes (discontinuation, time-to-transition, will-change). For single-step procedure transitions, count-based methods remain competitive — the world model's value is in long-range trajectory reasoning. Rigorous counterfactual/interventional validation and the full three-pillar loop (KG proposer + agentic verifier) extend naturally to a multimodal substrate (notes, WES, labs).

## Citation

```bibtex
@misc{gemeo_sus_2026,
  title  = {GEMEO/SUS: Recurrence-Aware Patient World Model for Rare Disease},
  author = {Timmers, Dimas and the Raras AI team},
  year   = {2026},
  url    = {https://huggingface.co/Raras-AI/gemeo-sus}
}
```

⚠️ **Research only.** Not a medical device. No clinical use without physician oversight and applicable regulatory clearance.