search

jasonkyuyim / se3_diffusion

Implementation for SE(3) diffusion model with application to protein backbone generation
https://arxiv.org/abs/2302.02277
MIT License
305 stars 50 forks source link

Reproduction SE3 Diffusion on CATH40 #22

Closed pengzhangzhi closed 4 months ago

pengzhangzhi commented 1 year ago

Hi Jason and all,

I reproduce the se3 diffusion training on cath40. the training log can be found at https://wandb.ai/traffic_flow_prediction/se3-diffusion/runs/f5dut9k5/workspace?workspace=user-zhangzhipeng I have a few observations.

Just share my observations with those who are interested in SE3 Diffusion.

pengzhangzhi commented 1 year ago

Hi Jason, The SE3 Diffusion is trained by optimizing the score predictions. Have you tried optimizing the X0 prediction? As the model just predicts X0 and then calculates the score. Why not directly optimize on X0?

jasonkyuyim commented 1 year ago

The model does optimize the x0 prediction. It is the same thing as optimizing the score up to a scaling factor https://github.com/jasonkyuyim/se3_diffusion/blob/master/experiments/train_se3_diffusion.py#L562

I would expect the model to perform worse on CATH since it is a smaller dataset. Are you sampling with low noise? That fixes any bond breaks. CATH to my knowledge also has a lot of loopy proteins. We found removing the loopy proteins to give better samples.

pengzhangzhi commented 1 year ago

Thank you! I am sorry I did not state my question properly. I know you have a loss supervising the x0 prediction at a time close to 0. My question is supervise x0 prediction at all times. I guess it is still the same thing? I can submit the ckpt trained on cath so that everyone else can play with it : ) The checkpoint is available at https://github.com/pengzhangzhi/se3_diffusion/releases/download/ckpt/step_710000.pth

jasonkyuyim commented 1 year ago

The model supervises x0 at all times. Only the auxiliary losses come in close to 0. I'll take a look at the checkpoint.

amorehead commented 1 year ago

Great work, @pengzhangzhi. I wonder if these results would be improved after retraining your copy of the model and dataset with the patched rotation loss.

jasonkyuyim commented 1 year ago

Update on this. Unfortunately I haven't had time to look at the checkpoint. I don't expect to have time in the near future as well. I agree with @amorehead you should try training with the fixed loss and see if it helps.

jasonkyuyim commented 4 months ago

Closing due to inactivity.