Clarification on auto_prob and batch_size Settings for Multi-System Training

[1292211370]

Dear Community Administrator,

I would like to use the se_e2_a operator to train an aqueous solution system containing six types of elements. My dataset includes over a dozen different systems. The largest system contains more than 4,000 data points, while the smallest has only about 300. Following the recommended parameters, I set "auto_prob": "prob_sys_size" and "batch_size=auto".

My question is: since the data volume varies greatly among systems, should I switch to "auto_prob": "prob_uniform" in order to ensure sufficient sampling of all systems? If so, should I also adjust the batch size accordingly?

Thank you very much!

[anyangml]

ob": "prob_uniform" in order to ensure sufficient sampling of all systems? If so, should I also adjust the batch size accordingly?

Suppose you have three systems: Sys1 (4000 frames), Sys2 (1000 frames), and Sys3 (200 frames). According to the Doc: if you set prob_uniform, each system is sampled with equal probability (1/3). If the systems contain a similar number of atoms, this means that— statistically — you train about 1 epoch on Sys1, 4 epochs on Sys2, and 20 epochs on Sys3.
If you want each system to be trained for a similar number of epochs (again assuming similar atom counts), use prob_sys_size instead.

If your systems differ significantly in atom counts, you may want to use prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;… to manually specify the sampling probability for each system. In this case, set the weight proportional to the number of frames in each system.

[1292211370]

Thank you for your reply. Since in my dataset, the amount of data from different systems varies by up to 10 times, I used "auto_prob": "prob_uniform". Compared with "auto_prob": "prob_sys_size", my model showed some improvement. Specifically, when the training step was set to 1,000,000, the model trained with "auto_prob": "prob_sys_size" caused the system to collapse rapidly during MD runs, while the model trained with "auto_prob": "prob_uniform" did not. However, when the training step reached 4,000,000, models trained with both parameters caused the system to collapse rapidly during MD runs. From the literature I reviewed, this issue may be related to the value of batch_size. In fact, I have seen some papers using batch sizes of 1, 4, or 8. Do you think batch_size is a parameter worth testing?

[anyangml]

Thank you for your reply. Since in my dataset, the amount of data from different systems varies by up to 10 times, I used "auto_prob": "prob_uniform". Compared with "auto_prob": "prob_sys_size", my model showed some improvement. Specifically, when the training step was set to 1,000,000, the model trained with "auto_prob": "prob_sys_size" caused the system to collapse rapidly during MD runs, while the model trained with "auto_prob": "prob_uniform" did not. However, when the training step reached 4,000,000, models trained with both parameters caused the system to collapse rapidly during MD runs. From the literature I reviewed, this issue may be related to the value of batch_size. In fact, I have seen some papers using batch sizes of 1, 4, or 8. Do you think batch_size is a parameter worth testing?

With the information provided, it is difficult for me to explain your observation. Regarding batch size, in general, we recommend using the largest possible value (typically auto:256 or auto:512, depending on the GPUs). A larger batch size usually helps stabilize the training process.