← All writing

Technical note

A working intuition for diffusion on SO(3)

Why rotations make a useful case study for respecting geometry in generative models.

Many generative models begin in Euclidean space because Gaussian noise is easy to state and reverse. Rotations are a useful reminder that representation is not an implementation detail: a 3×33 \times 3 rotation matrix is constrained, and adding ordinary Gaussian noise walks off the manifold.

The space we are actually modeling

A three-dimensional rotation belongs to the special orthogonal group:

SO(3)={RR3×3RR=I,detR=1}.SO(3) = \{R \in \mathbb{R}^{3 \times 3} \mid R^\top R = I, \det R = 1\}.

The tangent space at the identity is the Lie algebra so(3)\mathfrak{so}(3), consisting of skew-symmetric matrices. A vector ωR3\omega \in \mathbb{R}^3 maps to this space with the hat operator:

ω^=[0ω3ω2ω30ω1ω2ω10].\widehat{\omega} = \begin{bmatrix} 0 & -\omega_3 & \omega_2 \\ \omega_3 & 0 & -\omega_1 \\ -\omega_2 & \omega_1 & 0 \end{bmatrix}.

The exponential map then takes a local direction back onto the group. That is the geometric move: perturb in a tangent space, then map the result to a valid rotation.

This is not a full derivation. It is the minimal picture I want in mind before choosing a parameterization or a noise process.

A practical lens

In a conventional diffusion process, the forward marginal is often written as

q(xtx0)=N(αˉtx0,(1αˉt)I).q(x_t \mid x_0) = \mathcal{N}\left(\sqrt{\bar{\alpha}_t}x_0, (1 - \bar{\alpha}_t)I\right).

For rotational variables, the central question is not how to reuse that expression verbatim. It is how to define a stochastic process whose samples remain meaningful rotations, and whose score or denoising target is well-conditioned in the representation being used.

The implementation questions follow directly:

  1. Choose a stable representation and conversion path.
  2. Define noise in the tangent space or with a distribution on SO(3)SO(3).
  3. Test equivariance and numerical behavior before trusting a training curve.
# Sketch: local perturbation, then return to the group.
omega_t = sigma_t * torch.randn(batch, 3)
R_t = R_0 @ so3_exp(omega_t)

For protein structure generation, these choices are not cosmetic. Frames, relative orientations, and coordinate updates all carry geometric constraints. Keeping them visible in the model design makes debugging more concrete.

Where this goes next

The next note will connect this local picture to frame-based protein representations and the design of a denoising objective. Until then, the important thing is simple: RSO(3)R \in SO(3) is a constraint worth honoring from the first line of code.

Notes by Zaid Yusuf