Congrats on the new Aurora release!
I have been looking into the implementation of Aurora 1.5 Ensemble and I have some feedback.
In some applications we want to focus on individual ensemble members. To make this possible, it would be helpful to support deterministic, per-ensemble-member noise generation so that a given ensemble member produces the same noise sequence across model runs.
A possible implementation would be to assign each ensemble member its own RNG seed. These seeds could either be provided explicitly or generated automatically when the ensemble members are created.
We should not generate the same noise on every forward pass. The goal should be to reproduce the same noise sequence when processing the same sequence of batches in the same order. RNG state management could be exposed to the user so that they have control over the random number sequence predictability.
This could be done by adding a rng_reset() method to the Swin3DTransformerBackbone class, which restore the initial RNG state associated with the ensemble member. The Aurora ensemble model would also be extended with the ability to invoke rng_reset for every ensemble member. Users would invoke this method whenever deterministic semantics is required (for example, at the start of an inference run of the model).
To preserve current behaviour, this would be entirely optional. The per-ensemble-member seed would default to None (current semantics), with deterministic behaviour enabled only when an explicit seed is provided. With the current semantics, rng_reset() could either raise an exception or be a silent no-op.
Summary: support optional per-ensemble-member RNG seeds for the sake of making noise generation reproducible across runs for a given ensemble member.
Congrats on the new Aurora release!
I have been looking into the implementation of Aurora 1.5 Ensemble and I have some feedback.
In some applications we want to focus on individual ensemble members. To make this possible, it would be helpful to support deterministic, per-ensemble-member noise generation so that a given ensemble member produces the same noise sequence across model runs.
A possible implementation would be to assign each ensemble member its own RNG seed. These seeds could either be provided explicitly or generated automatically when the ensemble members are created.
We should not generate the same noise on every forward pass. The goal should be to reproduce the same noise sequence when processing the same sequence of batches in the same order. RNG state management could be exposed to the user so that they have control over the random number sequence predictability.
This could be done by adding a
rng_reset()method to theSwin3DTransformerBackboneclass, which restore the initial RNG state associated with the ensemble member. The Aurora ensemble model would also be extended with the ability to invokerng_resetfor every ensemble member. Users would invoke this method whenever deterministic semantics is required (for example, at the start of an inference run of the model).To preserve current behaviour, this would be entirely optional. The per-ensemble-member seed would default to
None(current semantics), with deterministic behaviour enabled only when an explicit seed is provided. With the current semantics,rng_reset()could either raise an exception or be a silent no-op.Summary: support optional per-ensemble-member RNG seeds for the sake of making noise generation reproducible across runs for a given ensemble member.