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

Superposing generated structures #24

Closed amorehead closed 4 months ago

amorehead commented 1 year ago

Hello again.

I was curious if you had any thoughts regarding superpositioning of the model's generated structures relative to the corresponding input (i.e., native) PDB structures for structural loss scoring. For example, in the structure prediction (not necessarily generation) community, one's predicted structures are usually superposed onto the corresponding native structure for the sake of calculating structural losses such that the model only needs to optimize the predicted structure up to a 3D rotation and translation (which will be handled e.g., by the Kabsch structural alignment algorithm).

I understand that this diffusion process is formally defined on a translation-invariant subspace, so 3D translations are likely already addressed (w.r.t. the generated structures). Moreover, since the model uses SE(3)-equivariant network layers to learn its structural distribution, the distribution it learns should be SE(3)-invariant. However, 3D rotations of the generated structures do not appear (to me) to be aligned with the input PDB structures when it comes to the calculation of the backbone atom loss, so this loss function appears (to me) to break the assumption that the final distribution the model learns should be rotation-invariant. Unless, by chance, these backbone atom loss terms are being computed from a residue-local perspective and not at a structure-global level?

I'd be curious to get your thoughts on this.

Best, Alex

blt2114 commented 1 year ago

Hi Alex,

Thanks for the question. This is an interesting idea. We haven't seriously considered within the SE(3)^N diffusion framework for the primary, score-matching loss. But it's possible that additional auxiliary losses involving an alignment step of this sort could improve performance.

I'm not sure what you mean by "this diffusion process is formally defined on a translation-invariant subspace". Can you elaborate? This is not how I think about our procedure.

Best, Brian

Wangchentong commented 1 year ago

@amorehead good point ! Guess backbone atom loss calculated from residue level local frame would be more suitable to keep SE3 invariant(which is the definition of Fape loss hhh), at least from my point of view.

amorehead commented 1 year ago

@blt2114, I see. When I mentioned translation-invariant subspace, I was borrowing the language used by other equivariant diffusion models like EDM that subtract the center of mass from all their input point clouds prior to model training. This appears to be the case in this repository's dataloader as well (i.e., zero-centering each training example).

jasonkyuyim commented 4 months ago

Closing due to inactivity.