Add NestedMultiHotProcessor and vectorise HALO's visit encoding - #1233
Open
janezdu wants to merge 1 commit into
Open
Add NestedMultiHotProcessor and vectorise HALO's visit encoding#1233janezdu wants to merge 1 commit into
janezdu wants to merge 1 commit into
Conversation
EHR generation feeds HALO a nested list of per-visit codes. The existing NestedSequenceProcessor emits code indices padded to the longest visit seen during fit, so one outlier sets the width for the whole dataset: on eICU a single visit holds 3,951 code entries (the same diagnosis re-charted through a stay) while a typical visit holds about five. HALO then unpacked that back into multi-hot vectors with a triple-nested Python loop. Add NestedMultiHotProcessor (registered as "nested_multihot"), which emits one multi-hot row per visit, sized by the vocabulary rather than by the worst-case visit -- 8.6x smaller per patient on a 921-code eICU vocabulary (4.5 KB vs 38.7 KB). Repeats within a visit collapse to a single 1, which is what set-membership models already did with the index form. Switch EHRGeneration to it and rewrite HALO._encode_visits as a vectorised placement into the context window. The old loop cost an .item() per patient -- a CUDA sync each -- and a single-element kernel launch per code, which measured at 84.5% of a training step at batch 128 on an A100 (13.887 s encode vs 2.553 s forward+backward), and ~99.8% in steady state once CUDA warmup is excluded. decode_dataset is updated to invert the new encoding: reading multi-hot rows as indices would see only 0s and 1s and decode every visit as empty. <pad> (0) and <unk> (1) keep their indices, so the vocabulary is interchangeable between the two processors. Tests: test_nested_multihot_processor.py covers the processor directly, and test_halo_encode_equivalence.py asserts HALO sees identical tensors from either processor, so the switch changes no results. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
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.
EHR generation feeds HALO a nested list of per-visit codes. The existing NestedSequenceProcessor emits code indices padded to the longest visit seen during `it, so one outlier sets the width for the whole dataset: on eICU a single visit holds 3,951 code entries (the same diagnosis re-charted through a stay) while a typical visit holds about five. HALO then unpacked that back into multi-hot vectors with a triple-nested Python loop.
Add NestedMultiHotProcessor (registered as
nested_multihot), which emits one multi-hot row per visit, sized by the vocabulary rather than by the worst-case visit -- 8.6x smaller per patient on a 921-code eICU vocabulary (4.5 KB vs 38.7 KB). Repeats within a visit collapse to a single 1, which is what set-membership models already did with the index form.Switch EHRGeneration to it and rewrite
HALO._encode_visitsas a vectorised placement into the context window. The old loop cost an.item()per patient -- a CUDA sync each -- and a single-element kernel launch per code, which measured at 84.5% of a training step at batch 128 on an A100 (13.887 s encode vs 2.553 s forward+backward), and ~99.8% in steady state once CUDA warmup is excluded.decode_dataset is updated to invert the new encoding: reading multi-hot rows as indices would see only 0s and 1s and decode every visit as empty.
(0) and (1) keep their indices, so the vocabulary is interchangeable between the two processors.
Tests: test_nested_multihot_processor.py covers the processor directly, and test_halo_encode_equivalence.py asserts HALO sees identical tensors from either processor, so the switch changes no results.
tl;dr NestedMultiHotProcessor + vectorised HALO._encode_visits. 8.6× less memory per patient, and encoding drops from 84.5% of a training step to negligible. Equivalence test proves results don't change. 10 files, +472/−33.