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..37c3d1d9e 100644 --- a/pyhealth/models/gamenet.py +++ b/pyhealth/models/gamenet.py @@ -129,7 +129,7 @@ class GAMENetLayer(nn.Module): Examples: >>> from pyhealth.models import GAMENetLayer >>> queries = torch.randn(3, 5, 32) # [patient, visit, hidden_size] - >>> prev_drugs = torch.randint(0, 2, (3, 4, 50)).float() + >>> prev_drugs = torch.randint(0, 2, (3, 5, 50)).float() >>> curr_drugs = torch.randint(0, 2, (3, 50)).float() >>> ehr_adj = torch.randint(0, 2, (50, 50)).float() >>> ddi_adj = torch.randint(0, 2, (50, 50)).float() @@ -172,12 +172,17 @@ def forward( Args: queries: query tensor of shape [patient, visit, hidden_size]. - prev_drugs: multihot tensor indicating drug usage in all previous - visits of shape [patient, visit - 1, num_drugs]. + prev_drugs: multihot tensor indicating drug usage in all + visits (including the current one, whose row should + already be zeroed out by the caller) of shape + [patient, visit, num_drugs]. This method itself drops the + current visit via [:, :-1, :] to derive DM_keys/DM_values. curr_drugs: multihot tensor indicating drug usage in the current visit of shape [patient, num_drugs]. mask: an optional mask tensor of shape [patient, visit] where 1 - indicates valid visits and 0 indicates invalid visits. + indicates valid visits and 0 indicates invalid visits. Also + used to exclude each patient's own current visit and any + padding beyond it from the dynamic-memory attention below. Returns: loss: a scalar tensor representing the loss. @@ -198,11 +203,30 @@ def forward( DM_keys = queries[:, :-1, :] DM_values = prev_drugs[:, :-1, :] + # For batches with variable-length patient histories, dropping only + # the batch's last (padded) column is not enough: a patient shorter + # than the batch's longest sequence still has their own current + # visit -- and pure padding beyond it -- sitting inside DM_keys/ + # DM_values. Build a per-patient mask over these positions so the + # attention below only ever sees genuine previous visits: valid iff + # the position is a real (non-padded) visit AND strictly before + # this patient's own current/last valid visit. + num_prev_positions = DM_keys.size(1) + position_idx = torch.arange(num_prev_positions, device=queries.device).unsqueeze(0) + last_visit_idx = (mask.long().sum(dim=1) - 1).clamp(min=0).unsqueeze(1) + prev_mask = (mask[:, :-1] > 0) & (position_idx < last_visit_idx) + """O: Output memory representation""" a_c = torch.softmax(torch.mm(query, MB.t()), dim=-1) o_b = torch.mm(a_c, MB) - a_s = torch.softmax(torch.einsum("bd,bvd->bv", query, DM_keys), dim=1) + attn_logits = torch.einsum("bd,bvd->bv", query, DM_keys) + attn_logits = attn_logits.masked_fill(~prev_mask, float("-inf")) + a_s = torch.softmax(attn_logits, dim=1) + # Patients with no valid previous visit (e.g. their very first + # visit) have an all -inf row, so softmax produces NaN; their + # dynamic-memory contribution should just be zero. + a_s = torch.nan_to_num(a_s, nan=0.0) a_m = torch.einsum("bv,bvz->bz", a_s, DM_values.float()) o_d = torch.mm(a_m, MB) @@ -225,12 +249,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 +286,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 +341,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 +357,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 +448,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 +503,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 +542,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..dfd79aeea 100644 --- a/tests/core/test_gamenet.py +++ b/tests/core/test_gamenet.py @@ -2,7 +2,7 @@ import torch from pyhealth.datasets import create_sample_dataset, get_dataloader -from pyhealth.models import GAMENet +from pyhealth.models import GAMENet, GAMENetLayer class TestGAMENet(unittest.TestCase): @@ -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,166 @@ 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) + + +class TestGAMENetLayerDynamicMemoryMasking(unittest.TestCase): + """Regression tests for GAMENetLayer.forward()'s dynamic-memory + attention over variable-length (padded) batches. + + Naively slicing DM_keys/DM_values as queries[:, :-1, :] and + prev_drugs[:, :-1, :] only drops the batch's last column. For a + patient shorter than the batch's longest sequence, that leaves the + patient's own current-visit position -- and any padding beyond it -- + inside the attention pool, stealing softmax weight from that + patient's genuine previous visits. Fixed by masking those positions + to -inf before the softmax. + """ + + def _make_layer(self, hidden_size, num_drugs, seed=0): + torch.manual_seed(seed) + ehr_adj = torch.randint(0, 2, (num_drugs, num_drugs)).float() + ddi_adj = torch.randint(0, 2, (num_drugs, num_drugs)).float() + layer = GAMENetLayer(hidden_size, ehr_adj, ddi_adj) + layer.eval() + return layer + + def test_dynamic_memory_ignores_own_current_visit_and_padding(self): + """Perturbing prev_drugs at a shorter patient's own current-visit + slot and at padding positions -- both of which fall inside the + naively-sliced DM_values range -- must not change that patient's + output at all: those positions must receive exactly zero dynamic- + memory attention weight. This would fail under the pre-fix + behavior, where nothing masks those positions out of the softmax. + """ + hidden_size, num_drugs, num_visits = 8, 5, 4 + layer = self._make_layer(hidden_size, num_drugs) + + torch.manual_seed(42) + queries = torch.randn(2, num_visits, hidden_size) + prev_drugs = torch.zeros(2, num_visits, num_drugs) + # Patient 0: 4 valid visits (mask all-ones), real history at 0,1,2. + prev_drugs[0, 0] = torch.tensor([1.0, 0, 0, 0, 0]) + prev_drugs[0, 1] = torch.tensor([0.0, 1, 0, 0, 0]) + prev_drugs[0, 2] = torch.tensor([0.0, 0, 1, 0, 0]) + # Patient 1: only 2 valid visits (current visit = index 1); one + # genuine prior visit at index 0. + prev_drugs[1, 0] = torch.tensor([0.0, 0, 0, 1, 0]) + + curr_drugs = torch.randint(0, 2, (2, num_drugs)).float() + mask = torch.tensor([[1.0, 1, 1, 1], [1.0, 1, 0, 0]]) + + with torch.no_grad(): + loss_base, y_prob_base = layer(queries, prev_drugs, curr_drugs, mask) + + prev_drugs_perturbed = prev_drugs.clone() + # Patient 1's own current-visit slot (index 1) and both padding + # slots (2, 3) -- all of which should be masked out. + prev_drugs_perturbed[1, 1] = torch.tensor([1.0, 1, 1, 1, 1]) + prev_drugs_perturbed[1, 2] = torch.tensor([1.0, 1, 1, 1, 1]) + prev_drugs_perturbed[1, 3] = torch.tensor([1.0, 1, 1, 1, 1]) + + with torch.no_grad(): + loss_pert, y_prob_pert = layer(queries, prev_drugs_perturbed, curr_drugs, mask) + + torch.testing.assert_close(y_prob_base[1], y_prob_pert[1]) + # Patient 0 (unaffected batch row, full-length sequence) must also + # be exactly unchanged. + torch.testing.assert_close(y_prob_base[0], y_prob_pert[0]) + torch.testing.assert_close(loss_base, loss_pert) + + def test_dynamic_memory_zero_for_first_visit_patient(self): + """A patient whose current visit is their very first visit has + zero valid previous visits -- the attention row is all -inf + pre-softmax. This must resolve to a zero dynamic-memory + contribution, not NaN propagating into the output.""" + hidden_size, num_drugs, num_visits = 8, 5, 3 + layer = self._make_layer(hidden_size, num_drugs, seed=1) + + torch.manual_seed(2) + queries = torch.randn(1, num_visits, hidden_size) + prev_drugs = torch.zeros(1, num_visits, num_drugs) + curr_drugs = torch.randint(0, 2, (1, num_drugs)).float() + # Single valid visit (the patient's first-ever visit); the rest is + # padding. + mask = torch.tensor([[1.0, 0, 0]]) + + with torch.no_grad(): + loss, y_prob = layer(queries, prev_drugs, curr_drugs, mask) + + self.assertFalse(torch.isnan(y_prob).any()) + self.assertFalse(torch.isnan(loss).any()) + + def test_full_length_patient_unaffected_by_masking_fix(self): + """Sanity check: for a patient whose sequence spans the batch's + full padded length (no padding, current visit is the batch's last + column), the mask covers exactly the same previous-visit positions + as the original unmasked slice -- output must be identical to a + run with mask=None (the pre-fix default).""" + hidden_size, num_drugs, num_visits = 8, 5, 4 + layer = self._make_layer(hidden_size, num_drugs, seed=3) + + torch.manual_seed(4) + queries = torch.randn(2, num_visits, hidden_size) + prev_drugs = torch.randint(0, 2, (2, num_visits, num_drugs)).float() + curr_drugs = torch.randint(0, 2, (2, num_drugs)).float() + mask_all_ones = torch.ones(2, num_visits) + + with torch.no_grad(): + loss_default, y_prob_default = layer(queries, prev_drugs, curr_drugs, mask=None) + loss_explicit, y_prob_explicit = layer( + queries, prev_drugs, curr_drugs, mask=mask_all_ones + ) + + torch.testing.assert_close(y_prob_default, y_prob_explicit) + torch.testing.assert_close(loss_default, loss_explicit) + + if __name__ == "__main__": unittest.main()