Question about fixed 10k training steps across CV schemes (risk of overfitting for modulo/contiguous)

[yuzimu1]

Hi ProteinNPT authors/maintainers,

Thank you very much for your work on ProteinNPT.

While reading the paper and looking through the code, I noticed that during training the model seems to use a fixed budget of 10,000 gradient steps for each assay and for each CV splitting scheme.

I am a bit puzzled by this choice: under fold_random_5, the position distribution in train/test is relatively uniform, so the model may be less prone to overfitting in the same way. However, for fold_modulo_5 and especially fold_contiguous_5, the train/test position distributions are more shifted. In that case, I wonder whether training for 10k steps could lead to overfitting to the training-position distribution (even if training loss keeps improving), and whether a much smaller number of steps (e.g., ~1k) might be optimal for these more challenging split schemes.

Thanks again for your time and for releasing this work.

[pascalnotin]

Hi @yuzimu1 -- we briefly discuss that point in Appendix B.3 of the paper. In earlier iterations, we had experimented with using a separate validation set (eg., random subset of 10-20% of the training data) to automatically detect overfitting / do early stopping. But, in practice, always training for 10k steps (using the full extent of the training data) led to higher performance on the validation benchmark assays. Quite often, training had converged earlier than 10k steps (eg., no gain after 1k training steps), but we rarely saw performance degradation from training longer, potentially because the dual objective that ProteinNPT is trained with (target prediction + MLM reconstruction) mitigates the overfitting risk.
The lower performance on the contiguous / modulo schemes is more of an issue of domain adaptation than overfitting imo -- once trained, the models should in general perform relatively well for novel mutations at positions in domains that were part of the training data (no overfitting), but out-of-domain predictions are more challenging to handle in general.