Hi! We are reproducing NEO visual pre-training with the official 2B recipe and consistently observe short but large spikes in text CE, independent of the tested optimizer hyperparameters.
The spikes affect all text datasets and all distributed ranks simultaneously, while the same batches are stable in other runs. Our current hypothesis is that the spatial H/W Q/K branch-particularly zero-initialized k_proj_hw followed by RMSNorm - may occasionally destabilize text attention.
Did you observe similar loss spikes during NEO pre-training? If so, did you use any additional stabilization such as a lower LR for spatial Q/K, gradual spatial-branch scaling, per-tensor update clipping, or masking the H/W contribution for text-text attention?
Thank you!
Hi! We are reproducing NEO visual pre-training with the official 2B recipe and consistently observe short but large spikes in text CE, independent of the tested optimizer hyperparameters.
The spikes affect all text datasets and all distributed ranks simultaneously, while the same batches are stable in other runs. Our current hypothesis is that the spatial H/W Q/K branch-particularly zero-initialized k_proj_hw followed by RMSNorm - may occasionally destabilize text attention.
Did you observe similar loss spikes during NEO pre-training? If so, did you use any additional stabilization such as a lower LR for spatial Q/K, gradual spatial-branch scaling, per-tensor update clipping, or masking the H/W contribution for text-text attention?
Thank you!