fix(determinism): make --seed actually reproducible - #6
Conversation
docs/prediction.md documents --seed as "Random seed for reproducible predictions", but two model inputs were drawn from RNGs that pl.seed_everything does not reach, so the same input could produce different atomistic features depending on --num_workers. 1. Ligand conformers. get_conformer embedded with ETKDGv3 and never set randomSeed (RDKit's default is -1, i.e. random). RDKit's global RNG starts from a fixed state per process, so a single-record run was already reproducible, but it advances with each embedding: in a batch the Nth ligand embedded inside a worker depends on how many preceded it, and therefore on how work is split across preprocess workers. Measured over 6 ligands, comparing --num_workers 1 vs 2 and 1 vs 4: 10 of 12 record/worker-count comparisons produced different geometry (up to 22.2 A); after the fix, 0 of 12. conformer_seed() now derives the seed from the base seed and the molecule's canonical SMILES, deliberately not from processing order. CCD and SDF ligands are unaffected: they already carry a conformer, so embedding never runs (verified: identical coordinates across seeds). 2. ref_pos augmentation. center_random_augmentation drew from the global torch RNG, which seed_everything(workers=True) seeds per DataLoader worker. With the conformer pinned so this was isolated, 0 of 6 records matched between --num_workers 0 and 1/2/4 (up to 14.5 A); after the fix, 6 of 6 match. The generator is now seeded from the existing per-record RandomState, so the augmentation distribution is unchanged. --seed is threaded to conformer generation so the flag stays meaningful, and base_seed=None keeps RDKit's previous non-deterministic behaviour. Also fixes `--num_workers 0`, which crashed with "max_workers must be greater than 0" because the value was reused for the preprocessing ProcessPoolExecutor. Note this is not output-preserving and cannot be: previous behaviour was random. The claim is identical output across runs and across --num_workers, not identical to any particular earlier run. Adds tests/test_determinism.py (6 tests). The two that rely only on existing APIs were confirmed to fail on main and pass here.
A note on the Ångström figuresThose are max absolute differences in a single atom coordinate, which mixes two things: a genuinely different conformation, and the arbitrary frame it happens to be generated in. A molecule that is merely rotated and translated shows a large raw coordinate difference while being chemically identical. To separate those, I compared rotation- and translation-invariant interatomic distance matrices across 4 unseeded ETKDG draws of the same ligand (imatinib, 37 heavy atoms, matching what nesso stores):
Radius of gyration across those draws spanned 5.55 to 6.36 Å, i.e. genuinely different conformations (extended versus folded), not one shape re-posed. As a sanity check on the method, applying a pure rotation plus translation to a single conformer gives a 17.6 Å raw coordinate difference and a 0.0000 Å interatomic distance difference. One case is different. In a separate run I pinned the conformer and varied only the |
Correcting something in the description above before this goes further. I ran a screening-scale check (40 diverse ligands against one target, run twice through the CLI) and found my "identical output across runs and worker counts" claim was too strong.
The correction: I chased the residual variance rather than guess at it. Comparing
So the input pipeline this PR targets really is deterministic. The remaining variance enters somewhere in the multi-worker DataLoader path. One thing worth separating out, since the first two rows of that table can be read as "this PR does nothing at
So even single-threaded, reordering or renaming input files changes every compound's conformer on What this PR does, stated accurately:
What it does not do: make multi-worker inference reproducible. That looks like a separate pre-existing bug, and it also cannot be fixed on its own; whoever addresses the DataLoader path still needs deterministic conformer generation underneath it or the result will not be reproducible either way. Happy to open a separate issue with the reproducer if useful. |
Seeding the ref_pos augmentation was not sufficient on its own. Two further
layers made a record's features depend on where it sat in the batch:
1. preprocess_yamls collected results with as_completed, which yields in
completion order, so the manifest order followed worker scheduling.
Measured over 40 records: order was stable at --num_workers 1 but differed
between runs at 2 and 4. Results are now collected by submission index.
2. InferenceDataset seeded RandomState from the dataset index, so that
scrambled order reached the augmentation. Even with order preserved, the
index still shifts whenever inputs are added, removed, renamed or
reordered. record_rng_seed() now derives the seed from the record id via
sha256, matching how conformer_seed keys on the molecule. hash() is unusable
here because it is salted per process.
Diagnosis: hashing every feature tensor reaching predict_step across two
--num_workers 4 CLI runs showed ref_pos differing in 39 of 40 records with all
26 other features identical, which pointed at the per-record RNG rather than
the transfer path. An earlier suspicion that pin_memory plus non_blocking
transfers were responsible was tested directly and disproved.
End to end over 40 ligands against one target, same library run twice:
comparison before after
nw=1 twice 0/40 0/40
nw=4 twice 17/40 0/40
nw=1 vs nw=4 39/40 0/40
nw=1 vs nw=2 n/a 0/40
All max deltas are exactly 0.0000 and all Spearman correlations exactly
1.000000, so this is bit-identical rather than close.
Adds two tests. test_ref_pos_independent_of_record_position featurizes the
same records in forward and reversed manifest order and was confirmed to fail
without the fix.
|
Found the root cause of the residual multi-worker nondeterminism I flagged above, and it is fixed in eaaf9c4. My earlier suspicion was wrong. It was not the DataLoader transfer path. I tested
Seeding the augmentation in the first commit was necessary but not sufficient, because the index it was keyed on was itself unstable. How I found it: hashing every feature tensor reaching The fix, two parts:
Result. Same 40 ligands against one target, run twice through the CLI:
Every max delta is exactly 0.0000 and every Spearman correlation exactly 1.000000, so this is bit-identical rather than merely close. That supersedes the "what it does not do" paragraph in my previous comment. End-to-end predictions are now reproducible across runs, worker counts, batch composition and input order, not just at Two new tests. |
|
@shenoynikhil this is settled now and ready for review whenever you have time. Apologies the thread ran long. I posted a claim, found while verifying at screening scale that it was too strong, and tracing that turned up two further causes rather than the one I expected. I left the corrections in place instead of force-pushing over them, so the comments read as the investigation rather than a clean summary. The description at the top is current and accurate; the comments below are the working record if you want the detail. Final state: 2 commits, three RNG causes fixed, CI green on all 7 checks. I am not planning to amend anything further. |
|
Thanks @mooreneural, I'll review it over the coming days |
2 participants
docs/prediction.mddocuments--seedas "Random seed for reproducible predictions", but three model inputs were drawn from RNGs thatpl.seed_everythingdoes not reach. The practical symptom: the same input produced different predictions depending on--num_workers, on the order of your input files, and on what else was in the batch.That last one matters most for screening. Because RDKit's unseeded RNG advances with every embedding, a compound's 3D conformer depended on how many molecules were embedded before it. Adding compounds to a screening library silently changed the geometry, and therefore the predicted affinity, of compounds already in it.
The three causes
1. Ligand conformers.
get_conformerembedded withETKDGv3and never setrandomSeed(RDKit's default is-1, random). RDKit's global RNG starts from a fixed state per process, so a single-record run was already reproducible; it advances with each embedding, so batches were not.conformer_seed()now derives the seed from the base seed and the molecule's canonical SMILES, deliberately not from processing order, sincepreprocess_yamlsparses in aProcessPoolExecutorwhere a counter would vary with worker scheduling. Keying on the molecule also means a given ligand embeds identically wherever it appears.2.
ref_posaugmentation.center_random_augmentationdrew from the global torch RNG, whichseed_everything(workers=True)seeds per DataLoader worker. It now takes an optionaltorch.Generator, seeded from the per-recordRandomState. The augmentation distribution is unchanged, only its RNG source.3. Batch position leaking into the per-record RNG.
preprocess_yamlscollected results withas_completed, which yields in completion order rather than submission order, so the manifest order followed worker scheduling.InferenceDatasetthen seeded itsRandomStatefrom the dataset index, feeding that scrambled order into the augmentation. Fixing 2 alone was not sufficient, because the index it keyed on was itself unstable.Results are now collected by submission index, and
record_rng_seed()derives the seed from the record id via sha256 rather than from position. Order-preservation alone is not enough: the index still shifts whenever inputs are added, removed, renamed or reordered.hash()is unusable here because it is salted per process.Measurements
Geometry, varying only the thing named:
--num_workers1 vs 2 vs 4 (6 ligands)--num_workers 4twiceref_posdiffered in 39 of 40 records, all others identicalNote on the first row: whether growing a batch perturbs a given compound depends on how the added molecules shift RDKit's RNG stream across workers, so it varies between repeats. I measured 2 of 5 in one run and 0 of 5 in another with a different set of added compounds. The order and worker-count rows reproduced identically on every repeat. That run-to-run variability is itself the defect being fixed, so it is worth stating rather than reporting a single figure as if it were stable.
End-to-end predictions, 40 diverse ligands against one target, run twice through the CLI:
nw=1twicenw=4twicenw=1vsnw=4nw=1vsnw=2(the default)Every max delta is exactly 0.0000 and every Spearman correlation exactly 1.000000, so this is bit-identical rather than merely close.
Does this reach the predictions? Measured on the released checkpoint (128 predictions: 4 arms × 16 replicates × 2 ligands), conformer choice moved
affinity_pred_valueacross a 0.232 log10 range for tyrosine, a ~1.7× spread in IC50, roughly a third of the model's own ensemble disagreement|value1 - value2|. The effect is ligand-dependent and not simply a function of flexibility: imatinib (37 heavy atoms) showed no spread above the augmentation baseline while tyrosine (13) did. I don't have an explanation for that and am not claiming one.Scope and honesty
--seedstill varies it, which is what a future conformer-ensembling option would use.base_seed=Nonerestores RDKit's previous non-deterministic behaviour for library callers.--num_workers > 1. I measured it, found the third cause above, and fixed it; the comments below have the full diagnosis.Why
base_seedis threaded rather than a module constantA module-level constant would make
--seedmeaningless for conformers, and the seed must be content-derived rather than process state to survive theProcessPoolExecutor. Every added parameter is keyword-with-default, so existing call signatures are unchanged (verified).Also included
--num_workers 0crashed withValueError: max_workers must be greater than 0, because the DataLoader value was reused for the preprocessing pool. Same flag, one-line fix (max(1, num_workers)).Tests
tests/test_determinism.py, 9 tests, ligand-only so they need no CCD asset and run on plain CI. Four of them use only pre-existing APIs and therefore run againstmainunchanged; all four fail there and pass here, so they genuinely catch the bug rather than merely describing the new code.test_ref_pos_independent_of_record_positionfeaturizes the same records in forward and reversed manifest order and was confirmed to fail without the position fix.Full suite: 29 passed, 7 skipped. The 7 are the pre-existing CCD-gated protein tests, which skip when the asset is absent (as on CI); with a CCD pickle present all 36 pass.
ruff 0.11.13clean.