Skip to content

Add NestedMultiHotProcessor and vectorise HALO's visit encoding - #1233

Open
janezdu wants to merge 1 commit into
sunlabuiuc:masterfrom
janezdu:multihot-ehrgen
Open

Add NestedMultiHotProcessor and vectorise HALO's visit encoding#1233
janezdu wants to merge 1 commit into
sunlabuiuc:masterfrom
janezdu:multihot-ehrgen

Conversation

@janezdu

@janezdu janezdu commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

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_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.

(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.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant