search

chaidiscovery / chai-lab

Chai-1, SOTA model for biomolecular structure prediction
https://www.chaidiscovery.com
Other
1.26k stars 158 forks source link

Guidance on changing num_diffn_timesteps #80

Open xwang112358 opened 1 month ago

xwang112358 commented 1 month ago

Thank you for the great work! In the ./example/predict_structure directory, I noticed that the diffusion step is set to 200 by default. I wonder if we could slightly decrease this value to reduce inference time while still maintaining decent results.

Allen

MattMcPartlon commented 1 month ago

Hi Allen,

We've seen pretty good results with less diffusion steps, though the other diffusion params, i.e. https://github.com/chaidiscovery/chai-lab/blob/306f53c7f45a6dd082aabd0d82a36d64fcf51f82/chai_lab/chai1.py#L500

might need to be changed. If you want to reduce to e.g. 80 steps then the current params are fine. If you want to go below that then I'd recommend trying different values of s_min, s_churn, and s_noise (ordered by importance). I found that 

S_churn \in (40.0, 80.0)
S_tmax=120.0
S_tmin=0.005
S_noise=1.005

works pretty well even in the 40-80 step range. To go below 40 steps you will need to use the scaling trick (\eta term in AF3 Algorithm 18 "Sample Diffusion"). This actually helps a lot with very few timesteps.

Also, if you're running on larger molecules just remember that the time spent in the trunk is usually much longer than the time spent in the diffusion module (O(L^3) vs O(SL^2) where S is the number of diffusion steps).

xwang112358 commented 1 month ago

Thank you for the insightful reply! @MattMcPartlon