Hi! As part of a master’s thesis at Karolinska Institutet, we are reproducing AlphaGenome pretraining from scratch (1 Mb, fold 0) in JAX on GPUs, where only one 1 Mb interval fits per device. Following the Supplementary Information and the released code, our training becomes unstable around the learning-rate peak. We would be grateful for a few details that the SI does not cover:
- RMSBatchNorm statistics: were they computed over all 64 intervals of the global batch, or per device? And was the gradient stopped through the variance (as in the training branch of the released RMSBatchNorm), or did it flow through, as in standard batch norm?
- Distillation: with one sample per GPU, how were the RMSBatchNorm statistics handled (synchronized across devices, per device, or frozen)?
- Initialization: do the initialization defaults in the released code (e.g. residual_scale 0.1, and a 1e-6 truncated normal for linear_embedding and the MLP output layer) match what the published models were trained with? The SI pseudocode gives 0.9 for residual_scale.
- Gradient clipping: was it used in pretraining, and if so at what global norm? (The fine-tuning colab clips at 0.5.)
- Weight decay: which parameters was it applied to? (The released fold_0 weights suggest kernels but not learnt_scale.)
- RNA-seq gene-level loss: which genes were included (all GENCODE v46 genes, or e.g. only protein-coding and lncRNA), and were gene bodies or exons used?
- Loss spikes: did you see any during pretraining, e.g. on intervals with very high small-RNA coverage such as RN7SL1/2 on chr14? If so, how did you handle them (clipping, skipping batches, restarts)?
Short answers are very welcome, even “yes/no” or “not sure” helps. Thank you!