DreamFold
Inquiry about the vector field calculation of the SO3FM function
Hey,
First of all, thank you very much for the excellent work by the authors. I am currently reading through the paper and the code, and I have encountered a small issue that has left me a bit puzzled. Specifically, the problem is as follows. Due to the differences in the setup of Flow Matching in FoldFlow (where $t = 0$ is the target distribution and $t = 1$ is the random distribution), I noticed a discrepancy when calculating the SO(3) vector field. In the paper, it is written:
However, after reviewing the code example, the implementation in the code looks like this:
def vectorfield(self, rot_0, rot_t, t):
...
rot_t_minus_0 = rot_0.transpose(-1, -2) @ rot_t
if self.inference_scaling < 0:
u_t = rot_t @ (
log(rot_t_minus_0)
/ torch.clamp(t[:, None, None], min=-self.inference_scaling)
)
else:
u_t = rot_t @ (log(rot_t_minus_0) * self.inference_scaling)
...
return None, u_t
The issue I don’t fully understand is that the paper describes the computation as $r_t^T r_0$, while the code implements it as $r_0^T r_t$. This seems to differ from what is described in the paper.
Another point I’m confused about is that the paper mentions this computation can be simplified to $\frac{\log_{r_t} (r_0)}{t}$. Would it be possible to directly compute it as log_not_from_identity(r_t, r_0})/ t?
Lastly, considering the difference in the notation of $t$ between the paper (where $t = 0$ is the random distribution and $t = 1$ is the target distribution) and the usual flow matching notation, if we were to change the notation to the usual one, could the interpolation on SO(3) be represented as:
$$ r_t = \exp_{r_1} \left( t \log_{r_1} (r_0) \right)? $$
Thank you for your time!
Hey, First of all, thank you very much for the excellent work by the authors. I am currently reading through the paper and the code, and I have encountered a small issue that has left me a bit puzzled. Specifically, the problem is as follows. Due to the differences in the setup of Flow Matching in FoldFlow (where t = 0 is the target distribution and t = 1 is the random distribution), I noticed a discrepancy when calculating the SO(3) vector field. In the paper, it is written:
However, after reviewing the code example, the implementation in the code looks like this:
def vectorfield(self, rot_0, rot_t, t): ... rot_t_minus_0 = rot_0.transpose(-1, -2) @ rot_t if self.inference_scaling < 0: u_t = rot_t @ ( log(rot_t_minus_0) / torch.clamp(t[:, None, None], min=-self.inference_scaling) ) else: u_t = rot_t @ (log(rot_t_minus_0) * self.inference_scaling) ... return None, u_t The issue I don’t fully understand is that the paper describes the computation as r t T r 0 , while the code implements it as r 0 T r t . This seems to differ from what is described in the paper.
Another point I’m confused about is that the paper mentions this computation can be simplified to log r t ( r 0 ) t . Would it be possible to directly compute it as
log_not_from_identity(r_t, r_0})/ t?Lastly, considering the difference in the notation of t between the paper (where t = 0 is the random distribution and t = 1 is the target distribution) and the usual flow matching notation, if we were to change the notation to the usual one, could the interpolation on SO(3) be represented as:
r t = exp r 1 ( t log r 1 ( r 0 ) ) ?
Thank you for your time!
Hello, have you solved this issue? I’ve encountered the same problem. It seems that the description in the article doesn’t match the implementation in the code, and some parts of the code are also inconsistent. This has left me confused.
For example,infoldflow.model.so3_fm.py, we have :
class SO3FM:
def vectorfield(self, rot_0, rot_t, t):
"""uses rot_0 and rot_t and t to calculate ut"""
batch_size = t.shape[0]
t = (
torch.clamp(t, min=1e-4, max=1 - 1e-4)
.repeat_interleave(rot_0.shape[1])
.double()
)
rot_0 = rearrange(rot_0, "t n c d -> (t n) c d", c=3, d=3).double()
rot_t = rearrange(rot_t, "t n c d -> (t n) c d", c=3, d=3).double()
rot_t_minus_0 = rot_0.transpose(-1, -2) @ rot_t
if self.inference_scaling < 0:
u_t = rot_t @ (
log(rot_t_minus_0)
/ torch.clamp(t[:, None, None], min=-self.inference_scaling)
)
else:
u_t = rot_t @ (log(rot_t_minus_0) * self.inference_scaling)
rot_t = rearrange(rot_t, "(t n) c d -> t n c d", t=batch_size, c=3, d=3)
u_t = rearrange(u_t, "(t n) c d -> t n c d", t=batch_size, c=3, d=3)
return None, u_t
But, in foldflow.utils.so3_condflowmatcher.py,the code is :
class SO3ConditionalFlowMatcher:
def compute_conditional_flow_simple(self, t, xt):
xt = rearrange(xt, "b c d -> b (c d)", c=3, d=3)
def index_time_der(i):
return torch.autograd.grad(xt, t, i, create_graph=True, retain_graph=True)[
0
]
xt_dot = vmap(index_time_der, in_dims=1)(
torch.eye(9).to(xt.device).repeat(xt.shape[0], 1, 1)
)
return rearrange(xt_dot, "(c d) b -> b c d", c=3, d=3)
I don't konw which one should be trusted.
Let me try to bring some clarity:
@mrzzmrzz, in my opinion the answer to both of your questions is yes cause:
- $log_{r_t}(r_0) / t$ is exactly
log_not_from_identity(point=r0, base_point=rt). - If you switch the indices of the data and the prior, you'll need to switch the direction of the flow, so basically swapping the indices would give exactly the formula you wrote.
@a-green-hand-jack, both methods are very closely related and should give the same result when self.inference_scaling factor is negative.
The latter method computes the vector field as $\dot{x}$, since $\frac{dx}{dt} = u_t$. This computation is done via torch's autograd and gives the exact theoretical vector field that represents the standard ODE that we're solving for.
However, in practice, as the authors mention on page 8 of their paper, they do not follow the theoretical ODE for generating rotations but instead multiply the vector field by a time dependent factor $i\cdot t$, which reduces the growth of the norm of the vector field during the end of inference. That's why you see 2 different variants of the computation of the vector field in that method: one without it, which follows the standard formula, and one with it, which essentially kills time $t$ in the denominator.
On top of that, the multiplication by rot_t, transports the "velocity" back to the rot_t point where we actually need to have it for proper flow computation. This velocity is a skew-symmetric matrix that is computed as the logmap of rot_t_minus_0 . The result of that logmap lies on the tangent space of the identity element of SO(3), on its Lie algebra $\mathfrak{s}o(3)$. That's why we need to "bring" it back to rot_t.
Hope this helps and correct me if I'm wrong in my thinking :)
