Fix GAMENet's Dynamic Memory - #1191
Open
lehendo wants to merge 2 commits into
Open
Conversation
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.
fbonc
approved these changes
Aug 20, 2026
fbonc
suggested changes
Aug 28, 2026
Contributor
There was a problem hiding this comment.
This problem predates your changes in this PR, but i think this implementation is incorrect for mixed-length batches now that dynamic memory is actually considered.
e.g., if patient A has 4 visits and patient B has 2, padding produces:
A: A1 A2 A3 A4
B: B1 B2 PAD PAD
queries[:, :-1] removes only the final column, producing:
A: A1 A2 A3
B: B1 B2 PAD
So this (correctly) excludes A’s current visit (A4), but it leaves B’s current visit (B2) AND padding in dynamic memory. This means softmax can assign attention to these invalid positions, taking weight away from B’s actual history (B1 here), which is bad.
I think to fix use the visit mask to exclude padding and each patient’s last valid visit before softmax. Also, patients without any prior visits should receive a zero dynamic-memory result.
…urrent visit A reviewer (fbonc) on PR sunlabuiuc#1191 flagged that the previous commit (using real drug history instead of hardcoded zeros) exposed a pre-existing bug: DM_keys/DM_values were sliced as queries[:, :-1, :] / prev_drugs[:, :-1, :], which only drops the batch's LAST padded column. For a patient shorter than the batch's longest sequence, this leaves that patient's own current-visit position -- and any padding beyond it -- inside the dynamic- memory attention pool, stealing softmax weight from their genuine previous visits. This didn't matter before the prior commit, since prev_drugs was always all-zero regardless of attention weights; now that it carries real history, misallocated attention actively degrades recommendation quality for any patient shorter than the batch max. Verified the reviewer's claim by reading the code directly: query = get_last_visit(queries, mask) correctly finds each patient's own last valid visit via the mask, but DM_keys/DM_values never used the mask at all. Fix: build a per-patient boolean mask over the DM_keys/DM_values positions (valid iff non-padded AND strictly before this patient's own current visit), mask the attention logits to -inf there before the softmax, and zero out (via nan_to_num) the resulting NaN row for patients with zero valid previous visits (e.g. a patient's very first visit), matching the reviewer's suggested fix exactly. Added regression tests: the main one perturbs prev_drugs at exactly the positions that should now be masked out and asserts the output is unchanged (confirmed this test fails against the pre-fix code with a 71% relative difference, and passes after the fix); plus a zero-prior-visit edge case (no NaN) and a full-length-patient sanity check (masking is a no-op when there's no padding). Also fixed a separate, pre-existing bug found during verification while in this file: GAMENetLayer's own docstring example used mismatched shapes (queries with 5 visits, prev_drugs with 4), causing an einsum shape mismatch -- confirmed this predates this PR entirely (same failure on master before any of these changes). Corrected the shapes and the Args docstring, which had documented prev_drugs as pre-truncated to visit-1 when the layer itself performs that truncation internally.
Contributor
|
The mixed-length mask looks good. One remaining issue: when a patient has no previous visits, the authors’ implementation reuses the graph-memory result, but this code uses zero and the new test locks that in. Could we match the reference behavior for first visits? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
GAMENet.forward() hardcoded prev_drugs = torch.zeros(...) instead of the patient's actual drug history
require drugs_hist in the dataset schema, precompute a remap from drugs_hist's own input vocabulary to the drugs label_vocab used by ehr_adj/ddi_adj/the Memory Bank, and build the real multi-hot prev_drugs tensor from that remapped history