Skip to content

Guard pretrained transfer, normalise content scale, and fit the standardiser on the full split - #48

Open
Rian354 wants to merge 10 commits into
mainfrom
feat/ssl-pretraining-and-encoders
Open

Guard pretrained transfer, normalise content scale, and fit the standardiser on the full split#48
Rian354 wants to merge 10 commits into
mainfrom
feat/ssl-pretraining-and-encoders

Conversation

@Rian354

@Rian354 Rian354 commented Aug 11, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

Six defects in the encoder and tokenisation path produced results that looked correct and were
not. A pretrained checkpoint trained a random backbone and reported no missing keys. A text
channel was 94% constant across patients. A standardiser fitted on 1/N of the train split under
torchrun. A note was cut to one fifth of its length.

Batch padding was the worst of them. No processor emits an event-level mask, so the unified
embedding marked every slot valid, and padding sorted ahead of every real event. All three backbone
families then read it.

This PR repairs them, and it includes the SSL pretraining package so that the frozen-encoder
correction has a file to change. It supersedes draft #45.

The pretraining section is deliberately short. Pretraining is not yet useful in practice here, and
this PR makes no claim that a pretrained model is better than a model trained from the start.

Depends on #46. scripts/pretrain_ssl.py imports LabsOnlyMIMIC4, and tests/test_pretrain.py
imports sample_oversample. Both arrive in #46. Merge #46 first.

Implemented

1. Strict transfer of a pretrained checkpoint

  • File: pyhealth/models/base_model.py
  • Downstream models give the unified backbone a different name for each architecture:
    _unified_backbone for transformer, _unified_jamba for jamba, _unified_blocks for ehrmamba.
  • The previous loader built the full model.state_dict(), overwrote the keys that matched, then
    called strict=False. Every key was present, so PyTorch reported no missing keys.
  • A jamba checkpoint that matched approximately 6 of 30 backbone tensors therefore trained a
    backbone that was mostly random, and the run looked correct.
  • load_pretrained_state_dict maps the keys for each architecture, refuses a shape mismatch,
    requires full backbone coverage, and returns the tensor counts.
  • Bookkeeping buffers (_checkpoint_config, _dummy_param) are not counted. To count them makes
    a complete load appear partial. A TransformerLayer carries _checkpoint_config, which alone
    gives 24/25 = 0.96 and fails min_backbone_match=1.0.

2. Scale-safe unified embedding

  • File: pyhealth/models/embedding/unified.py
  • Each event is final = (content + time_emb + type_emb) * validity. The measured norms at
    embedding_dim=128:
Term Source Norm
time_emb SinusoidalTimeEmbedding(t) 8.0
type_emb nn.Embedding(n_modalities, D) 11.8
content, text BERT [CLS] through a default Linear(768, D) 3.2
content, laboratory raw values (sodium 140, osmolality 290) through Linear(10, D) 761.4
  • This is a range of 238x through the same code path. The contribution of a modality came from
    the output scale of its encoder and not from a modelling decision.
  • Measured on real MIMIC-IV notes through Bio_ClinicalBERT: cosine similarity across patients was
    0.9215 for the content alone and 0.9953 for the content with the constant terms. The signal
    was 5.9% of the squared norm of the final vector.
  • normalize_content applies F.layer_norm to the content term before the additive terms.
Check Before After
Text separation across patients 0.0485 0.3836 (7.9x)
Laboratory information kept (linear recovery of glucose) 1.0000 0.9999
Parameters added - 0
  • The correction is parameter free, so an existing checkpoint continues to load.

3. Standardiser that reads the full train split

  • File: pyhealth/processors/lab_standardizer.py (new), wired into the numeric path of
    unified.py
  • SampleDataset is a subclass of litdata.StreamingDataset. _DistributedEnv divides both
    __len__ and __iter__ by WORLD_SIZE. torchrun sets WORLD_SIZE but does not set
    GLOBAL_RANK, and torch.distributed is not initialised when the dataset is built. Every rank
    therefore reported global_rank=0 and fitted the same first part.

Measured against real litdata with 20 samples:

Condition len() Iteration Indexing
Single process 20 20 20
WORLD_SIZE=4 5 5 20
  • Production SSL starts with torchrun. On 4 GPUs the statistics came from 25% of the split, and
    the digest recorded train/4 samples. The single-process downstream runner could never match
    that digest.
  • The fit now reads region_of_interest, which WORLD_SIZE does not divide. An earlier attempt
    used patient_to_index. SampleDataset.subset copies that map from the parent, so after
    split_by_patient it still holds parent indices while __getitem__ is restricted to the
    subset. Fitting on a real training split raised ValueError: index 237 didn't find a match within the chunk intervals. The fit refuses to iterate a StreamingDataset that has no
    region of interest.
  • The standardiser applies before the projection. The projection mixes the features, so a
    transform after it cannot correct a feature whose physical unit gives it 300 times the magnitude
    of another.
  • The statistics are buffers, so they travel in state_dict. A checkpoint applies at inference the
    same transform that it trained under.

4. Frozen text encoder during pretraining

  • Files: scripts/pretrain_ssl.py, configs/pretrain/base.yaml
  • The configuration had the default freeze_encoder: false. Every downstream run uses
    --freeze-encoder.
  • SSL pretraining therefore updated all 110 million BERT parameters and trained a text path that
    inference then discards. The [CLS] cache serves only a frozen field, so the cache could not
    operate.
  • The measured cost was 2.4 hours for each epoch for notes, against approximately 5 minutes for
    laboratory values. A 12-hour limit stopped the run at 5 of 10 epochs.
  • The default is now true, and --no-freeze-encoder selects the previous behaviour.

5. Token budget of 512

  • File: pyhealth/processors/tuple_time_text_processor.py
  • max_length had the default 128. This cuts 95% of the extracted discharge notes to
    approximately one fifth of their length. Measured on 4,017 samples.
  • Changed to 512, which is the position-embedding limit of Bio_ClinicalBERT. A larger budget
    without a long-context tokenizer now gives an error. Before this change, Hugging Face reduced the
    value in silence.
  • Padding uses longest for each batch and not a fixed width. A note of 4 tokens no longer costs
    512 slots.

6. Cache for the frozen text encoder

  • File: pyhealth/models/embedding/unified.py
  • A frozen encoder is deterministic. A run of 50 epochs calculated the same forward pass of 110
    million parameters 50 times.
  • The cache has three conditions. It serves only a field in _frozen_text_fields, so a trainable
    encoder can never read it. The key covers only the tokens the attention mask marks real, so
    batch padding cannot give the same note a different key each epoch. A key over the padded row
    never hit under shuffling: epoch times on the full-scale notes run were 3458s, 3936s, 4048s,
    3835s. The cache has a maximum size, and it calculates a row again when the cache is full.
  • A cold batch encodes each distinct row one time. A missing note is a constant placeholder, so one
    batch holds many identical rows.
  • train() sets a frozen encoder to .eval(). Without this, dropout changes the output between
    passes, and the cache becomes an approximation instead of an identity.

7. Batch padding is masked in the unified sequence

  • Files: pyhealth/datasets/collate.py, pyhealth/models/embedding/unified.py,
    pyhealth/models/rnn.py
  • The collator pads short samples to the longest in the batch and fills value and time with 0.0.
    Nothing recorded that padding. No processor emits an event-level mask either, so the unified
    embedding took the mask is None branch and marked every slot valid. A padded slot then
    looked exactly like a real measurement taken at admission time.
  • Ordering made it worse. Padding carries time 0.0, so the ascending sort placed it before
    every real event, while all three backbone families assume the opposite: RNNLayer packs the
    first mask.sum() steps, get_last_visit indexes mask.sum() - 1, and TransformerLayer
    reads position 0 as its CLS vector.
  • The first attempt added pad_mask to collate_temporal, which has zero callers. The
    dataloader uses collate_fn_dict_with_padding. That function now reports {field}__pad_mask.
    _build_unified_inputs keys the field dict off schema(), so validity has to arrive on a
    separate key or it never reaches the model. All four backbones that build unified inputs thread
    it through. The unified embedding uses it for event validity, sorts invalid slots past every
    real one, and zeroes their embeddings.
  • The sort is now stable. The key is heavily tied, since all padding shares time 0.0 and
    events from one admission share offsets. An unstable sort changes RNN and Mamba outputs between
    torch builds and between CPU and CUDA.
  • pad_mask is deliberately not called mask. A field may carry its own {field}_mask meaning
    "was this value observed", which is a different question from "is this slot real".

8. The standardiser reads the observation field

  • File: pyhealth/models/embedding/unified.py
  • Contribution 3 above read feat_dict["mask"] for observation flags. No processor has ever
    populated that key    , so the standardiser raised on every run that enabled it:
    ValueError: The standardiser for 'labs' needs a paired labs_mask field.
  • Together with the end-to-end PR, this made the two PRs unrunnable as a pair. It now reads the
    sibling {field}_mask field, which is where the observation flags actually live.

9. Zero-length guard in RNNLayer

  • File: pyhealth/models/rnn.py
  • Required by contribution 7. While every mask was all ones, a zero length was unreachable. A
    correct mask makes a sample with no valid event reachable, and pack_padded_sequence rejects
    it. lengths is clamped to 1, so the caller sees a finite value instead of a crash.

10. Pad every ragged dimension, not only dimension 0

  • File: pyhealth/datasets/utils.py
  • Contribution 5 pads each sample to its longest note rather than to a fixed max_length. That
    makes both the event dimension and the token dimension vary across samples. pad_sequence pads
    dimension 0 only and requires every trailing dimension to match, so a batch of notes raised
    RuntimeError: The size of tensor a (7) must match the size of tensor b (14).
  • _pad_stack right-pads every dimension to the batch maximum.

11. Token mask and padding mask are not the same variable

  • File: pyhealth/models/embedding/unified.py
  • Preferring pad_mask for event validity reused that (B, N) tensor as the attention mask
    handed to the text encoder, which expects token level: RuntimeError: shape '[96, 512]' is invalid for input of size 96.
  • Three masks meet in this loop and are now named apart: mask is token level from the
    processor; pad_mask is event level from the collator; {field}_mask is "was this value
    observed" and belongs to the standardiser.

Validation

The cache does not change the result

Full scale, 144,586 train samples, the same configuration with and without --no-text-cache:

cache enabled  : train_loss [1.0934, 1.0600, 1.0612, 1.0636]
--no-text-cache: train_loss [1.0934, 1.0601, 1.0612, 1.0637]

The two runs agree to four decimal places across four epochs.

SSL encoders complete on corrected data

Transformer backbone, 128 dimensions, 2 layers, 4 heads, 10 epochs, campus cluster.

Encoder val_loss Seconds for each epoch
labs_only MAE 0.5262 -> 0.4623 240 -> 229
labs_only SimMIM 0.5320 -> 0.4460 206 -> 197
notes_labs MAE 0.6297 -> 0.6137 1618 -> 1176
notes_labs SimMIM 0.6196 -> 0.5829 1533 -> 1077

All four completed. The decrease from 1618 to 1176 seconds was measured before the padded-key
fix above. A later full-scale supervised run did not show that decrease, because shuffle changed
the pad width every epoch. Against the previous attempt with a trainable encoder: 5 of 10 epochs
at approximately 8,600 seconds each, stopped by the 12-hour limit.

One result is not clean: notes_labs MAE increases at epoch 2, from 0.6297 to 0.6443, and then
decreases.

Unit tests

File: tests/test_unified_padding.py. Later commits added tests for the dataloader collate, an
unpadded batch, and the backbone threading pad_mask into unified inputs. The last of those is
the test that would have caught the inert collate_temporal fix.

File: tests/test_encoders_and_tokenisation.py. Later commits added a fit-on-a-split test and a
test that the frozen-text cache key ignores batch padding.

Test group Purpose
Transfer guard (4) A complete checkpoint loads; a partial one, a shape mismatch, and a model with no registered backbone are refused
Standardiser (4) The fit ignores unobserved slots, an unobserved slot maps to zero, WORLD_SIZE does not change the statistics, and the statistics are in state_dict
Content scale (2) The correction is parameter free, and it restores separation between patients
Token budget (2) The default is 512, and padding is not a fixed width
Frozen cache (4) A repeated note encodes one time, a trainable encoder never reads the cache, the mask is part of the key, and a full cache recalculates

15 of the 16 fail on main. The one that passes,
test_content_normalisation_restores_separation_between_patients, demonstrates the mechanism with
tensors and calls no library code, so it is a description and not a guard.

tests/test_pretrain.py::test_load_pretrained_into_downstream_transformer now runs a real SimMIM
checkpoint into a real Transformer through the guarded loader and asserts full backbone coverage.
It previously called a private helper in the end-to-end runner.

No regression

The full suite was run two times: on main alone, and on main with #46 and this PR applied.

Tree Passed Failed
main 873 40
main + #46 + this PR 942 34

The 34 remaining failures are also present on main. They come from the environment and not from
the code: test_tfm_tokenizer and test_tuple_time_text_tokenizer cannot build a fast tokenizer
without sentencepiece, and test_audio_processor cannot load torchaudio against this build of
torch.

One LIME test differed between the two runs in each direction. Both pass 4 times of 4 in isolation
on both trees, so the difference comes from the order of the tests and from the global random
state, and not from this PR.

What this PR does not claim

No comparison between a pretrained model and a model trained from the start. The 26 encoders from
earlier work are superseded. They were trained on raw laboratory values, on notes that were
approximately 90% placeholder, and any encoder trained under torchrun holds statistics from 1/N
of the split. The transfer guard refuses them, and that is correct.

Questions for the reviewer

  1. Is the parameter-free layer_norm the correct correction? It also raises a low-signal modality
    to equal scale.
  2. Is an identical fit_scope digest too strict? A nested tuning subsplit can never reproduce the
    SSL split digest.
  3. freeze_encoder now has the default true for pretraining. Is any planned run dependent on the
    previous default?

@Rian354
Guard pretrained transfer, normalise content scale, fit the standardi…
c871cdb 
…ser on the full split

Four defects gave results that looked correct and were not.

A checkpoint that matched 6 of 30 backbone tensors trained a mostly random
backbone. The loader built the full target state dict, overwrote the keys that
matched, then called strict=False, so PyTorch reported no missing keys.
load_pretrained_state_dict now maps the backbone name for each architecture and
requires full coverage.

The content term reached a norm of 761 for raw laboratory values against 3.2 for
a BERT [CLS] vector, through the same code path. The constant time and type
terms then dominated the text channel: cosine similarity across patients was
0.9953. normalize_content applies a parameter-free layer_norm to the content
term before the additive terms, so an existing checkpoint still loads.

SampleDataset subclasses litdata.StreamingDataset, which divides __len__ and
__iter__ by WORLD_SIZE. torchrun sets WORLD_SIZE but not GLOBAL_RANK, so every
rank fitted statistics on the same first 1/N of the train split. The fit now
reads patient_to_index, which WORLD_SIZE does not divide.

max_length had the default 128, which cut 95% of extracted discharge notes to
about one fifth of their length. It is now 512, the position-embedding limit of
Bio_ClinicalBERT, and padding is per batch rather than a fixed width.

Also: SSL pretraining now freezes the text encoder by default, to match the
downstream path and to let the [CLS] cache operate; the cache is keyed on the
token identifiers and the attention mask, is bounded, and only serves a frozen
field. This supersedes draft #45, whose package is included here.
@Rian354 Rian354 mentioned this pull request Aug 11, 2026
Open
Rian354 and others added 7 commits August 12, 2026 13:50
@Rian354
Mask batch padding in the unified sequence, and read observation flag…
6e9a073 
…s from the right field

The collator pads short samples to the longest in the batch and fills value and
time with 0.0. Nothing recorded that padding. No processor emits an event-level
mask either, so the unified embedding took the `mask is None` branch and marked
every slot valid. A padded slot then looked exactly like a real measurement
taken at admission time.

Ordering made it worse. Padding carries time 0.0, so the ascending sort placed
it BEFORE every real event, while all three backbone families assume the
opposite: RNNLayer packs the first mask.sum() steps, get_last_visit indexes
mask.sum() - 1, and TransformerLayer reads position 0 as its CLS vector.

The collator now emits pad_mask for each temporal field, because the collator
creates the padding and nothing else can know about it. The unified embedding
uses it for event validity, sorts invalid slots past every real one, and zeroes
their embeddings. The sort is now stable: the key is heavily tied, since all
padding shares time 0.0 and events from one admission share offsets, and an
unstable sort changes RNN and Mamba outputs between torch builds and between
CPU and CUDA.

pad_mask is deliberately not called mask. A field may carry its own
{field}_mask meaning "was this value observed", which is a different question
from "is this slot real". The standardiser needs the former, and it was reading
feat_dict["mask"], which no processor has ever populated, so it raised on every
run that enabled it. It now reads the sibling {field}_mask field.

Restoring the zero-length clamp in RNNLayer is required by the above. While
every mask was all ones, a zero length was unreachable; a correct mask makes a
sample with no valid event reachable, and pack_padded_sequence rejects it.

Found while reviewing the same defects in the upstream consolidation branch.
@Rian354
Drive the standardiser fit from the region of interest, not the paren…
12add17 
…t index map

The fit was driven by patient_to_index. SampleDataset.subset copies that map
unchanged, so after split_by_patient it still holds indices into the PARENT
dataset while __getitem__ is restricted to the subset's own region. Fitting on a
real training split therefore raised:

  ValueError: The provided index 237 didn't find a match within the chunk
  intervals [Interval(chunk_start=0, roi_start_idx=135, roi_end_idx=136,
  chunk_end=297), ...]

region_of_interest is the only unsharded description of what a dataset actually
holds, and it is correct for a subset as well as for a full dataset. Measured
against real litdata with 300 samples: the full fit sees 300, an ROI-restricted
split sees its own 40, and under WORLD_SIZE=4 the fit still sees 300 with
statistics identical to the single-process fit.

This preserves the property the original change was for. len() and __iter__ are
still sharded by WORLD_SIZE; indexing and the region of interest are not.

Caught by a smoke run on real MIMIC-IV. The unit tests passed because their
fake dataset modelled provenance through patient_to_index, which is the
mechanism that turned out to be wrong; the fake now models the litdata contract
that was verified against real litdata.
@Rian354
Record padding in the collate the dataloader actually uses
d770167 
The previous commit added pad_mask to collate_temporal. Nothing calls it.
get_dataloader uses collate_fn_dict_with_padding in pyhealth/datasets/utils.py,
so the padding fix was inert in production while its unit tests passed, because
those tests called collate_temporal directly.

collate_fn_dict_with_padding now reports which event slots are real, under the
parallel key {field}__pad_mask. _build_unified_inputs keys the field dict off
schema(), and no processor's schema contains "mask", so the validity has to
arrive on a separate key or it never reaches the model. All four backbones that
build unified inputs now thread it through.

A batch whose rows are already equal length is not padded, and no mask is
invented for it.

Three new tests cover the real path: the dataloader collate records padding, an
unpadded batch reports none, and the backbone threads it into the unified
inputs. The last of these is the test that would have caught the inert fix.

Found by a smoke run on real MIMIC-IV.
@Rian354
Pad every ragged dimension in the collate, not only dimension 0
bbfbcf3 
Contribution 5 changed the text processor to pad to the longest note in a
sample rather than to a fixed max_length, so a note of 4 tokens no longer costs
512 slots. That makes the token dimension vary across samples as well as the
event dimension, and pad_sequence pads dimension 0 only and requires every
trailing dimension to match already:

  RuntimeError: The size of tensor a (7) must match the size of tensor b (14)
  at non-singleton dimension 1

_pad_stack right-pads every dimension to the batch maximum and reduces to
pad_sequence when only dimension 0 differs.

This is the other half of the token-budget change and should have shipped with
it. Caught by a smoke run on real MIMIC-IV notes; no unit test tokenised two
samples of different length into one batch.
@Rian354
Separate the token mask from the padding mask in the unified embedding
17aaeb0 
The previous commit made the event-validity variable prefer pad_mask. The TEXT
branch reuses that same variable as the attention mask it hands to the encoder,
so an event-level (B, N) mask reached a view that expects token level:

  RuntimeError: shape '[96, 512]' is invalid for input of size 96

Three masks meet in this loop and are now named apart: mask is token level from
the processor schema and is what a text encoder needs; pad_mask is event level
from the collator and says which slots are real; {field}_mask is a separate
field meaning "was this value observed" and belongs to the standardiser.

Event validity now prefers the collator's report and falls back to reducing the
token mask, with the nested-CODE flattening handled.
@Rian354
Key the frozen text cache on real tokens, not on batch padding
fb16cc1 
The collator pads each row to the widest note in its batch, and batch
composition changes every epoch because the loader shuffles. A key over the
padded row therefore gave one note a different key each epoch, so the cache
never hit.

Measured on the full-scale notes run: epoch time did not fall after the first
epoch, and rose slightly as the miss path also paid for hashing:

  3458s, 3936s, 4048s, 3835s

The key now covers only the positions the attention mask marks real, which is
invariant to how wide the batch happened to be.
@Rian354 @cursoragent
Name the standardiser provenance as region_of_interest, not the paren…
d08b8d7 
…t map

The fit has read region_of_interest since the subset-index fix. The error
messages and one test docstring still said patient_to_index, which is the
mechanism that raised on a real training split.

Co-authored-by: Cursor <cursoragent@cursor.com>
This was referenced Aug 13, 2026
Open
Closed
Rian354 and others added 2 commits August 17, 2026 21:20
@Rian354 @cursoragent
Keep padding_idx=0 on nested code embeddings.
bc00852 
NestedSequenceProcessor used padding_idx=None so empty visits could have a non-zero vector. Index 0 then received gradients, which moved a real code that shared the pad slot. Empty visits are now zero events, so the pad row stays frozen zeros. The test checks both the zeros and a zero gradient on that row.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Thread collate pad_mask through the unified MLP path.
996fea7 
The other unified heads already forwarded {field}__pad_mask. MLP did not, so a padded lab slot still looked like a real measurement at time 0. The new test is the same collate-to-inputs check that would have caught the inert Transformer fix.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354

Rian354 commented Aug 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Nested code embeddings keep padding_idx=0, so the pad row stays zeros and gets no gradient. Unified MLP now forwards the collate pad mask, same as the other heads.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Rian354