qazwsxal / diffusion-extensions Goto Github PK

Thanks for sharing your marvelous project!
I have some questions about the use of so3_lerp and so3_scale.
According to the paper, it seems that in the following code, so3_scale should be used instead of so3_lerp towards identity rotation.

Is there anything I missed or misunderstood? Or does that mean so3_lerp could approximate so3_scale? Your response would help a lot.

Hi, I have a question about the conversion from the axis-angle representation vector to a skew-symmetric matrix, as implemented here:

In the above implementation, the skew-symmetric matrix $S(v)$ for $v = \left[ x, y, z \right]$ is defined as:

$$S(v) = \begin{bmatrix} 0 & -z & y \\\ z & 0 & -x \\\ -y & x & 0 \end{bmatrix}$$

which is different from the definition in your ICLR paper (Denoising Diffusion Probabilistic Models on SO(3) for Rotational Alignment, https://openreview.net/forum?id=BY88eBbkpe5, page 3, six lines above Eq. 6).

Did I misunderstand something? Which definition should be used here for such conversion?

Hello, thanks for providing such a wonderful method and code that helped me improve our project.

However, while testing the code you open sourced, I discovered that there may be a problem with the code that samples rotations in the form of axis angles from an isotropic Gaussian distribution:

In the above code, trap_start and trap_end are the indexes before and after the CDF is equal to unif, but they all obtain the index in the CDF corresponding to the first variance, rather than the corresponding variance.

I trained and tested the network's performance on the task of generating $\pm90$-degree rotations along the z-axis and got the following result plot:

.

The predicted rotation distribution differs greatly from the true value.

I optimized this part of the code so that trap_start and trap_end are the indexes before and after the CDF corresponding to their respective variances is equal to unif. The optimized core code is as follows:

        B = self.eps.shape[0]
        bs = torch.arange(0, B)
        trap_start = self.trap[idx_0, bs]
        trap_end = self.trap[idx_1, bs]

In addition to this, I made several other changes to make the code runnable:

(1)
Original code:

modified code:

R = process.p_sample(R, torch.full((BATCH,), i, device=device, dtype=torch.long))

(2)
Original code:

modified code:

sample = IsotropicGaussianSO3(model_stdev).sample()

The result after modifying the code is shown below:

The predicted rotation distribution is almost the same as the true value.

If I understand it wrong, please point out my problem so that I can better understand your method, thank you~