From 2f5f4f00c1e9e9b9aa67b22b1ae60571c73895dc Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 18 Aug 2026 17:59:50 -0700 Subject: [PATCH] Fix GAMENet's Dynamic Memory: use real drug history, not zeros GAMENet.forward() hardcoded prev_drugs = torch.zeros(...) instead of the patient's actual drug history, making the graph-augmented Dynamic Memory mechanism the model is named for permanently inert. Per the paper (Shang et al., "GAMENet: Graph Augmented MEmory Networks for Recommending Medication Combination," AAAI 2019, arXiv:1809.01852), the Dynamic Memory's values (Eq. 6) should be [c_m^1; ...; c_m^{t-1}] -- each previous visit's actual administered drugs, retrieved via temporal attention (Eq. 7). That attention/retrieval math was already correct; only the memory's content was wrong. pyhealth.tasks.drug_recommendation already produces exactly the needed field (drugs_hist: nested per-visit drug history, current/target visit zeroed out) -- GAMENet just never consumed it (an unused batch_to_multihot import was a leftover sign of the abandoned wiring). Fix: require drugs_hist in the dataset schema (a silent zeros-fallback would just reintroduce the same silently-wrong-results bug in a new form), precompute a remap from drugs_hist's own input vocabulary to the drugs label_vocab used by ehr_adj/ddi_adj/the Memory Bank (they are tokenized independently), and build the real multi-hot prev_drugs tensor from that remapped history. Verified end-to-end on real hardware: unit tests (including two new regression tests), a 30-patient stress test with variable visit counts, and the actual documented example script trained against real synthetic MIMIC-III data. --- docs/api/models/pyhealth.models.GAMENet.rst | 7 ++ .../drug_recommendation_mimic4_gamenet.py | 1 - pyhealth/models/gamenet.py | 110 +++++++++++++++++- tests/core/test_gamenet.py | 57 +++++++++ 4 files changed, 168 insertions(+), 7 deletions(-) diff --git a/docs/api/models/pyhealth.models.GAMENet.rst b/docs/api/models/pyhealth.models.GAMENet.rst index 55a7aadbc..84942b02d 100644 --- a/docs/api/models/pyhealth.models.GAMENet.rst +++ b/docs/api/models/pyhealth.models.GAMENet.rst @@ -3,6 +3,13 @@ The separate callable GAMENetLayer and the complete GAMENet model. +GAMENet requires ``drugs_hist`` (nested per-visit drug history, with the +current/target visit already zeroed out, e.g. as produced by +:mod:`pyhealth.tasks.drug_recommendation`) in the dataset's ``input_schema``. +This is used to populate the paper's Dynamic Memory (Eq. 6): each previous +visit's actual administered drugs, retrieved via the query-key temporal +attention in Eq. 7. + .. autoclass:: pyhealth.models.GAMENetLayer :members: :undoc-members: diff --git a/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py b/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py index bd5b33cb0..16680fe28 100644 --- a/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py +++ b/examples/drug_recommendation/drug_recommendation_mimic4_gamenet.py @@ -6,7 +6,6 @@ from pyhealth.tasks import drug_recommendation_mimic4_fn # import dataloader related functions -from pyhealth.datasets.splitter import split_by_patient from pyhealth.datasets import split_by_patient, get_dataloader # import gamenet model diff --git a/pyhealth/models/gamenet.py b/pyhealth/models/gamenet.py index 46afe057f..7f2383564 100644 --- a/pyhealth/models/gamenet.py +++ b/pyhealth/models/gamenet.py @@ -225,12 +225,22 @@ class GAMENet(BaseModel): Note: This model is only for medication prediction which takes conditions and procedures as feature_keys, and drugs as label_key. - It only operates on the visit level. Thus, we have disable the + It only operates on the visit level. Thus, we have disable the feature_keys, label_key, and mode arguments. Note: This model only accepts ATC level 3 as medication codes. + Note: + Requires ``drugs_hist`` (nested per-visit drug history, current + visit excluded, e.g. as produced by + :mod:`pyhealth.tasks.drug_recommendation`) in the dataset's + input_schema. This populates the paper's Dynamic Memory (Eq. 6): + each previous visit's key-value pair of (patient query, actual + administered drugs), retrieved via temporal attention at inference + time (Eq. 7) to condition the recommendation on the patient's own + medication history. + Args: dataset: the dataset to train the model. It is used to query certain information such as the set of all tokens. @@ -252,15 +262,19 @@ class GAMENet(BaseModel): ... samples=[ ... { ... "patient_id": "patient-0", - ... "visit_id": "visit-0", - ... "conditions": [["cond-33", "cond-86"], ["cond-80"]], - ... "procedures": [["proc-12", "proc-45"], ["proc-23"]], - ... "drugs": ["drug-1", "drug-2", "drug-3"], + ... "visit_id": "visit-2", + ... "conditions": [["cond-33", "cond-86"], ["cond-80"], ["cond-91"]], + ... "procedures": [["proc-12", "proc-45"], ["proc-23"], ["proc-67"]], + ... # drugs_hist: per-visit drugs administered so far, + ... # with the current (target) visit zeroed out. + ... "drugs_hist": [[], ["drug-1"], []], + ... "drugs": ["drug-2", "drug-3"], ... } ... ], ... input_schema={ ... "conditions": "nested_sequence", ... "procedures": "nested_sequence", + ... "drugs_hist": "nested_sequence", ... }, ... output_schema={"drugs": "multilabel"}, ... dataset_name="test", @@ -303,6 +317,11 @@ def __init__( assert "conditions" in self.dataset.input_schema, "conditions must be in input_schema" assert "procedures" in self.dataset.input_schema, "procedures must be in input_schema" + assert "drugs_hist" in self.dataset.input_schema, ( + "drugs_hist must be in input_schema (nested per-visit drug history, " + "current visit excluded) -- required to populate the paper's Dynamic " + "Memory (Eq. 6); see e.g. pyhealth.tasks.drug_recommendation." + ) assert "drugs" in self.dataset.output_schema, "drugs must be in output_schema" # feature_keys and label_key for GAMENet @@ -314,6 +333,19 @@ def __init__( self.embedding_model = EmbeddingModel(dataset, embedding_dim) self.label_size = len(self.dataset.output_processors[self.label_key].label_vocab) + # drugs_hist is tokenized against its own input vocabulary (built by + # NestedSequenceProcessor), which is generally NOT the same indexing + # as the drugs label_vocab used by ehr_adj/ddi_adj/the Memory Bank. + # Precompute a remap table (drugs_hist vocab index -> drugs label_vocab + # index) so historical drug codes can be converted into the same + # multi-hot space as the drug label, as required to build the Dynamic + # Memory's values in Eq. 6. Codes with no match (, , or a + # history code absent from the output label_vocab) map to + # self.label_size, a sentinel "trash" bin sliced away after scatter. + self.register_buffer( + "_drug_hist_to_label", self._build_drug_hist_vocab_map() + ) + # adj matrix ehr_adj = self.generate_ehr_adj() ddi_adj = self.generate_ddi_adj() @@ -392,6 +424,52 @@ def generate_ddi_adj(self) -> torch.tensor: ddi_adj[label_vocab[atc_j], label_vocab[atc_i]] = 1 return ddi_adj + def _build_drug_hist_vocab_map(self) -> torch.Tensor: + """Maps drugs_hist input-vocabulary indices to drugs label_vocab + indices, so historical drug codes align with the same multi-hot + space as the drug label / EHR & DDI graphs. + + Returns: + LongTensor of shape [drugs_hist_vocab_size]. Entry i is the + label_vocab index of the drugs_hist-vocab code at index i, or + self.label_size (a sentinel "trash" index, sliced away after + scatter) if that code has no corresponding drug label (this + covers the drugs_hist processor's own / tokens, plus + any history code absent from the output label_vocab). + """ + hist_vocab = self.dataset.input_processors["drugs_hist"].code_vocab + label_vocab = self.dataset.output_processors[self.label_key].label_vocab + + mapping = torch.full((len(hist_vocab),), self.label_size, dtype=torch.long) + for code, hist_idx in hist_vocab.items(): + if code in label_vocab: + mapping[hist_idx] = label_vocab[code] + return mapping + + def _build_prev_drugs(self, drugs_hist: torch.Tensor) -> torch.Tensor: + """Converts the raw drugs_hist tensor into the multi-hot Dynamic + Memory values Eq. 6 of the paper requires: [c_m^1; ...; c_m^{t-1}]. + + Args: + drugs_hist: LongTensor of shape [batch, visits, codes_per_visit], + indices into the drugs_hist input processor's vocabulary + (current visit already zeroed out by the task, per + pyhealth.tasks.drug_recommendation). + + Returns: + Multi-hot tensor of shape [batch, visits, label_size] aligned + with the drugs label_vocab / ehr_adj / ddi_adj / Memory Bank. + """ + batch_size, num_visits, _ = drugs_hist.shape + mapped = self._drug_hist_to_label[drugs_hist.clamp(min=0)] + # mapped values are in [0, label_size]; label_size is the sentinel + # trash bin absorbing //unmatched codes. + multihot = torch.zeros( + batch_size, num_visits, self.label_size + 1, device=drugs_hist.device + ) + multihot.scatter_(2, mapped, 1.0) + return multihot[:, :, : self.label_size] + def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """Forward propagation. @@ -401,6 +479,9 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: Expected keys: - conditions: tensor of shape [batch, visits, codes_per_visit] - procedures: tensor of shape [batch, visits, codes_per_visit] + - drugs_hist: tensor of shape [batch, visits, codes_per_visit], + nested per-visit drug history with the current visit + zeroed out (see pyhealth.tasks.drug_recommendation) - drugs: tensor of shape [batch, num_drugs] (multilabel) Returns: @@ -437,7 +518,24 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: batch_size = queries.size(0) num_visits = queries.size(1) - prev_drugs = torch.zeros(batch_size, num_visits, self.label_size, device=self.device) + # Dynamic Memory values (Eq. 6): each previous visit's actual + # administered drugs, multi-hot encoded in the drugs label_vocab + # space. drugs_hist already has the current visit zeroed out by the + # task (see pyhealth.tasks.drug_recommendation); GAMENetLayer drops + # the last (current) visit itself when slicing DM_keys/DM_values. + drugs_hist = kwargs["drugs_hist"].to(self.device) + if drugs_hist.size(1) != num_visits: + # conditions/procedures/drugs_hist are built in lockstep per + # visit by the task, but the default collate function pads + # each field independently -- align defensively just in case. + if drugs_hist.size(1) < num_visits: + pad = drugs_hist.new_zeros( + batch_size, num_visits - drugs_hist.size(1), drugs_hist.size(2) + ) + drugs_hist = torch.cat([drugs_hist, pad], dim=1) + else: + drugs_hist = drugs_hist[:, :num_visits, :] + prev_drugs = self._build_prev_drugs(drugs_hist) # [batch, visits] mask = (embedded["conditions"].sum(dim=-1) != 0).any(dim=-1) diff --git a/tests/core/test_gamenet.py b/tests/core/test_gamenet.py index 8b735b857..df4909153 100644 --- a/tests/core/test_gamenet.py +++ b/tests/core/test_gamenet.py @@ -14,6 +14,10 @@ def setUp(self): "visit_id": "visit-0", "conditions": [["cond-33", "cond-86"], ["cond-80", "cond-12"]], "procedures": [["proc-45", "proc-23"], ["proc-67"]], + # drugs_hist: per-visit drugs actually administered so far, + # with the current (target) visit already zeroed out, as + # produced by pyhealth.tasks.drug_recommendation. + "drugs_hist": [["drug-2"], []], "drugs": ["drug-1", "drug-2", "drug-3"], }, { @@ -21,6 +25,7 @@ def setUp(self): "visit_id": "visit-1", "conditions": [["cond-33"], ["cond-80"]], "procedures": [["proc-45"], ["proc-23", "proc-67"]], + "drugs_hist": [["drug-4"], []], "drugs": ["drug-2", "drug-4"], }, { @@ -28,6 +33,7 @@ def setUp(self): "visit_id": "visit-2", "conditions": [["cond-86", "cond-80"], ["cond-12"]], "procedures": [["proc-45", "proc-67"], ["proc-23"]], + "drugs_hist": [["drug-5", "drug-1"], []], "drugs": ["drug-1", "drug-4", "drug-5"], }, ] @@ -35,6 +41,7 @@ def setUp(self): self.input_schema = { "conditions": "nested_sequence", "procedures": "nested_sequence", + "drugs_hist": "nested_sequence", } self.output_schema = {"drugs": "multilabel"} @@ -63,10 +70,12 @@ def test_forward_input_format(self): self.assertIn("conditions", data_batch) self.assertIn("procedures", data_batch) + self.assertIn("drugs_hist", data_batch) self.assertIn("drugs", data_batch) self.assertEqual(len(data_batch["conditions"].shape), 3) self.assertEqual(len(data_batch["procedures"].shape), 3) + self.assertEqual(len(data_batch["drugs_hist"].shape), 3) self.assertEqual(len(data_batch["drugs"].shape), 2) def test_model_forward(self): @@ -128,6 +137,54 @@ def test_output_shapes(self): self.assertEqual(ret["loss"].shape, ()) + def test_missing_drugs_hist_raises(self): + """Regression test: constructing GAMENet without drugs_hist in the + input_schema must fail loudly, not silently fall back to zeroed + history (the original bug).""" + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": [["cond-33"], ["cond-80"]], + "procedures": [["proc-45"], ["proc-67"]], + "drugs": ["drug-1", "drug-2"], + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "nested_sequence", "procedures": "nested_sequence"}, + output_schema={"drugs": "multilabel"}, + dataset_name="test_missing_hist", + ) + with self.assertRaises(AssertionError): + GAMENet(dataset=dataset) + + def test_dynamic_memory_uses_real_drug_history(self): + """Regression test for the critical bug: the Dynamic Memory's + values (prev_drugs, Eq. 6 of the GAMENet paper) must be populated + from each patient's actual drugs_hist, not hardcoded zeros.""" + train_loader = get_dataloader(self.dataset, batch_size=3, shuffle=False) + data_batch = next(iter(train_loader)) + + drugs_hist = data_batch["drugs_hist"] + prev_drugs = self.model._build_prev_drugs(drugs_hist.to(self.model.device)) + + # Every sample in this test set has non-empty history at visit 0 + # (patient-0: drug-2, patient-1: drug-4, patient-2: drug-5/drug-1), + # so the resulting multi-hot tensor must NOT be all zeros. + self.assertGreater(prev_drugs.sum().item(), 0) + + # The visit-0 row for patient-0 should have exactly one drug + # ("drug-2") marked, at the index the label_vocab assigns it. + label_vocab = self.model.dataset.output_processors["drugs"].label_vocab + self.assertEqual(prev_drugs[0, 0].sum().item(), 1.0) + self.assertEqual(prev_drugs[0, 0, label_vocab["drug-2"]].item(), 1.0) + + # The current (target) visit's history was zeroed out by the task + # convention, so its row must be all zeros. + self.assertEqual(prev_drugs[0, 1].sum().item(), 0.0) + + if __name__ == "__main__": unittest.main()