A theoretical problem about protein structure DDPM and SE(3) invariant model

[Guo-Stone]

Hi ! Dear authors, your derivation in this paper is very rigorous and amazing! Including RFdiffusion, there are two SE(3) equivariant models for protein structure generation based on diffusion model. My question is: is it possible for an SE(3) invariant model to do protein structure generation with diffusion model? To my way of thinking, an SE(3) invariant model plus an alignment algorithm (e.g. Kabsch algorithm) is equal to an SE(3) equivariant model. Looking forward to your answer!

[jasonkyuyim]

Hi! If I understand your question, you're asking if a SE(3) invariant score model can be used to perform SE(3) equivariant diffusion? Of course you can use any model to do diffusion but achieving equivariance is important. We utilize IPA which (by its name) is a SE(3) invariant model but with a change of frame it becomes equivariant (see the AF2 supp for more details). IPA is somewhat unique in this sense though by keeping track of these "local frames" at each residue. I'm not sure how you'd use an alignment algorithm to achieve equivariance with a invariant model. You could look into works like frame averaging which AIUI is a framework for converting neural networks to be equivariant/invariant.

[Guo-Stone]

Thanks for your patient answer! In my opinion: 1) When sampling, DDPM needs a network to predict x_pred (to approximate x0) in order to calculate p( x(t-1) | x_t, x_pred) . 2) In SE(3) equivariant diffusion, we have to ensure that x_pred, x_0 and xt are In the same coordinate system. Only that we can calculate p( x(t-1) | x_t, x_pred) directly and calculate the loss || x_pred - x_0 ||^2_2. 3) In your FrameDiff, you use an IPA module (SE3 invariant) to update invariant node features hl and update your frame T(l+1) = T_l * T_update. Because T_update derives from invariant node features hl, T(l+1) is equivariant as T_l.

I'm not sure whether my above comprehensions are right. And my questions are as following: 1) I want to use an SE(3) invariant model to perform SE(3) equivariant diffusion, which means that x_pred may not in the same coordinate system as x_t (or x0). So before calculating p( x(t-1) | x_t, x_pred), I have to align x_pred to xt, i.e. p( x(t-1) | x_t, Align(x_pred)). 2) When calculating loss, I want to replace || x_pred - x_0 ||^2_2 by || Align(x_pred) - x_0 ||^2_2 or FAPE(x_pred, x_0).

I hope I can express my questions clearly and look forward to your answer! orz

[dugu9sword]

Through computing FAPE-style loss, you gets an SE(3) invariant loss, instead of SE(3) equivariant prediction. They are not equivalant, and what we expect is "the score rotates with the input".

[jasonkyuyim]

Yes dugu's answer is correct. I think the equivariant loss fits the theory more closely but I have heard of many works where invariant losses lik FAPE also work. Hope you were able to solve your problem and sorry I dropped the ball on responding. I'm closing the issue for now.