Ensemble members with noise injection should be able to generate predictable and reproducible noise. #191

Activity

1. s-sasaki-earthsea-wizard commented on Jul 19, 2026

   

   Syota-Sasaki (s-sasaki-earthsea-wizard)

   on Jul 19, 2026

   Contributor

   More actions

   I had the same need for reproducible ensemble members, so I prototyped this proposal on a branch — and it works as hoped: with an optional per-member noise_seed, a given ensemble member reproduces the exact same noise sequence across runs and across fresh model instances, and calling rng_reset() restores the sequence from the start. noise_seed=None keeps the current behaviour unchanged. All 43 tests under tests/v1p5/ pass.

   The branch: feature/per-member-noise-seed (pinned at 1ff8fe5). Details below.

   Implementation

   • An optional noise_seed: int | None = None on Swin3DTransformerBackbone, plumbed through Aurora (and therefore AuroraV1p5Ensemble). When set, noise is drawn from a dedicated per-instance torch.Generator instead of the global RNG. None preserves the current noise-generation semantics.
   • rng_reset() restores the initial RNG state, so processing the same sequence of batches in the same order reproduces the same noise sequence. Exposed on the model via Aurora.rng_reset().
   • The generator is created lazily on the first noise sample because it must live on the same device as the noise tensor (a CPU generator cannot be used for CUDA tensors).

   A few design decisions on the points the issue left open — feedback welcome:

   • With noise_seed=None, rng_reset() is a warning-emitting no-op rather than an exception, so it is safe to call unconditionally while still surfacing misuse.
   • When a seed is configured, rng_reset() also flushes the noise cache used for noise accumulation: leftover cached noise would otherwise contaminate the reproduced sequence.
   • The existing reset_noise() (cache only) and the new rng_reset() (RNG state + cache) stay separate methods, with the distinction documented.
   • Moving the model to another device recreates the generator from the initial seed (with a warning), which restarts the noise sequence from the beginning; this is documented in the docstring.

   Verification

   tests/v1p5/test_noise_seed.py covers:

   • same seed + rng_reset() reproduces the exact noise sequence, also with noise accumulation enabled;
   • the same seed on a fresh instance reproduces the same sequence (the core requirement here);
   • without rng_reset() the stream keeps advancing, and different seeds produce different noise;
   • a seeded stream is independent of the global RNG;
   • noise_seed=None preserves the current global-RNG semantics, and the no-op rng_reset() leaves both the global RNG state and the noise cache untouched.

   To reproduce:

   git clone -b feature/per-member-noise-seed https://github.com/s-sasaki-earthsea-wizard/aurora.git
   cd aurora
   git checkout 1ff8fe5
   pip install -e ".[dev]"
   pytest tests/v1p5/ -v

   Expected result: all 43 tests pass.

   One note for testing: with randomly initialised weights the adaptive-LN modulation is zero-initialised (init_respostnorm), so the noise context does not affect the output at initialization. The tests therefore record the sampled noise directly instead of comparing model outputs.

   Happy to open a PR if this direction looks right to the maintainers.
2. wesselb commented on Aug 19, 2026

   

   Wessel (wesselb)

   on Aug 19, 2026

   Collaborator

   More actions

   Hey Agnieszka Słowik (@Slowika) and Syota-Sasaki (@s-sasaki-earthsea-wizard)! Thanks for opening and issue and apologies for the slow reply from our side.

   Instead of adding an argument noise_seed to the constructor, how would it sound to add a keyword generator: torch.Generator | tuple[torch.Generator | None, ...] | None to Aurora.forward, which is then passed to Swin3DTransformerBackbone.forward and used in the torch.randn calls?

   If the user generates multiple ensemble members simultaneously by using a batch size, one could pass a tuple of generators to control the randomness for each ensemble member simultaneously.
3. wesselb commented on Aug 19, 2026

   

   Wessel (wesselb)

   on Aug 19, 2026

   Collaborator

   More actions

   CC Jonathan Weyn (@jweyn)
4. Slowika commented on Aug 20, 2026

   

   Agnieszka Słowik (Slowika)

   on Aug 20, 2026

   ContributorAuthor

   More actions

   Hi Wessel (@wesselb) ! Thank you for the reply.

   I like this approach! It gives the user a higher level of control over using fixed random seeds.

   Syota-Sasaki (@s-sasaki-earthsea-wizard) Are you happy to make this change on your branch and open a PR? I don't think this is a big change.
5. s-sasaki-earthsea-wizard commented on Aug 24, 2026

   

   Syota-Sasaki (s-sasaki-earthsea-wizard)

   on Aug 24, 2026

   Contributor

   More actions

   Wessel (@wesselb) Agnieszka Słowik (@Slowika) Thank you for the feedback! I have opened #202 implementing the generator keyword on Aurora.forward as suggested, with a tuple of per-member generators supported for batched ensemble members. The previous prototype branch stays untouched for reference. Details and a few design decisions are described in the PR.
6. Wessel (wesselb)

   added a commit that references this issue on Sep 25, 2026

   Add `generator` argument to `forward` for reproducible noise in stoch…

   ...

   b628d7c
7. Wessel (wesselb)

   closed this as completedin #202on Sep 25, 2026