InstaDeepAI

/

NTv3_pretrained

Like 1

Follow

InstaDeep Ltd 451

RMSBatchNorm normalises by the batch statistics in train mode but never updated its running variance var_ema, while the JAX layer does (EMA, momentum 0.999). A torch fine-tune therefore trained on batch statistics and then evaluated on the pre-fine-tuning ones: no error, just worse results.

This adds the update, following the Flax layer exactly (genomix_research/models/ntv3/layers.py): var_ema <- 0.999 * var_ema + 0.001 * E[x^2], skipped when the result is not finite.

• Eval mode, hence inference, is unchanged: only the self.training branch changes.
• All five models (NTv3_pretrained, NTv3_tracks, NTv3_annot_*) load this file through their auto_map, so the fix applies to all of them.
• The file is generated from genomix_research/models/ntv3/torch/layers.py by ntv3_huggingface/build_hf_modeling_files.py; the regenerated file differs from the published one by this change only.
• Tested against the Flax layer in tests_torch/test_ntv3_torch_parity.py (outputs and var_ema after two train-mode steps, then eval, plus the non-finite guard).

The Space notebook 02 hook that worked around this is skipped automatically once the layer has ema_momentum.