Group-aware Geodesic Losses in SE(3)

• The paper introduces the Frame Aligned Frame Error (F2E), which replaces chordal metrics with geodesic distances to ensure constant gradients during optimization.
• It leverages the mathematical structure of SE(3) to combine rotational and translational components into a single, group-aware loss function applied to molecular complex modeling.
• Empirical results demonstrate significant improvements in docking accuracy, with relative performance gains of up to 182% compared to traditional chordal loss methods.

Group-aware geodesic losses represent a class of structural loss functions for learning over rigid-body transformations, formulated on the Lie group $\mathrm{SE}(3)$ and explicitly optimized to respect the geometry of group actions—especially when modeling interactions in multi-component molecular complexes. These losses generalize traditional structure comparison metrics to operate directly in the group-theoretic framework, enabling principled, gradient-stable optimization even under large frame misalignments. Notably, the Frame Aligned Frame Error (FAFE, or F2E) demonstrates that replacing traditional chordal (Frobenius norm–based) losses with geodesic metrics on $\mathrm{SE}(3)$ yields superior convergence and empirical performance in protein docking contexts, such as antibody–antigen complex modeling (Wu et al., 2024).

1. Mathematical Framework of Group Frames and $\mathrm{SE}(3)$

Rigid-body transformations in three-dimensional space form the matrix Lie group $\mathrm{SE}(3)=\mathrm{SO}(3)\ltimes \mathbb{R}^3$, where each frame or pose $g=(R, t)$ consists of a rotation matrix $R\in \mathrm{SO}(3)$ and a translation vector $t\in \mathbb{R}^3$. The homogeneous matrix representation is

$$G=\begin{pmatrix} R & t \ 0^\top & 1 \end{pmatrix}.$$

In this formalism, composition and inversion follow matrix group operations: $$g_1g_2 = (R_1R_2, R_1t_2 + t_1), \quad g^{-1} = (R^\top, -R^\top t).$$ Left- and right-invariant distances depend only on $g_1^{-1}g_2$ or $g_2g_1^{-1}$, respectively, and are natural candidates for geometric loss formulations over both individual residues and entire molecular subunits.

2. Frame Aligned Point Error (FAPE) and Chordal Distances

AlphaFold2's main structural loss, Frame Aligned Point Error (FAPE), is formulated as an $\mathrm{SE}(3)$-invariant metric. Given true frames $\mathbf{T} = [T_1,\ldots,T_N]$ and predicted frames $\hat{\mathbf{T}} = [\hat{T}_1,\ldots,\hat{T}_N]$, the per-pair FAPE between residues $i\neq j$ is

$$L_{\mathrm{FAPE}}(i,j) = \|\hat{R}_i^\top(\hat{t}_j-\hat{t}_i) - R_i^\top(t_j-t_i)\|_2,$$

aggregated over inter-chain pairs. This expression estimates the error in local coordinates, and (under near-rigidity assumptions) can be interpreted as a chordal distance between group frames $\Delta T = T_B^{-1}T_A$ and $\Delta \hat{T} = \hat{T}_B^{-1}\hat{T}_A$: $$L_{\rm G\text{-}FAPE}(\Delta) \propto d_c(\Delta,\widehat\Delta),$$ where $d_c(R_1, R_2) = \|R_1 - R_2\|_F$ denotes the chordal (Frobenius) metric on $\mathrm{SO}(3)$, with translation handled via appropriate scaling.

3. Gradient Pathologies in Chordal Losses

Chordal distance-based losses, such as FAPE, exhibit gradient-vanishing pathologies at large angular deviations. The functional relationship

$$d_c(\theta) = 2\sqrt{2}\,\sin\left(\frac{\theta}{2}\right)$$

implies that the corresponding gradient magnitude

$$\frac{\mathrm{d}}{\mathrm{d}\theta}d_c(\theta)\propto \cos(\theta/2)$$

vanishes as $\theta \rightarrow \pi$. Empirically, for antibody–antigen docking, the majority of predictions exhibit $\theta > \pi/2$, placing them in the region where the loss provides minimal effective learning signal. These gradient-degenerate regimes result in optimizers failing to reduce rotational misalignments for difficult docking targets (Wu et al., 2024).

4. Frame Aligned Frame Error (F2E): Formulation of the Group-aware Geodesic Loss

To address gradient vanishing, the Frame Aligned Frame Error (F2E) replaces the chordal metric with the true geodesic on $\mathrm{SE}(3)$. The per-pair loss is

$$L_{\mathrm{F2E}}(i,j) = d_G(\hat{T}_i^{-1}\hat{T}_j, T_i^{-1}T_j) = \|\log (\hat{T}_i^{-1}\hat{T}_j (T_i^{-1}T_j)^{-1})\|_F,$$

where $\log$ denotes the principal matrix logarithm. Decomposing into rotational and translational components yields the double-geodesic form

$$d_G((R_1,t_1),(R_2,t_2)) = \sqrt{d_\theta(R_1,R_2)^2 + \left\|\frac{t_1 - t_2}{\alpha}\right\|^2},$$

with $$d_\theta(R_1, R_2) = |\arccos((\mathrm{Tr}(R_1^\top R_2)-1)/2)|$$ and $\alpha$ a translation scaling factor. The complex-level (inter-chain) geodesic loss is then

$$L_{\mathrm{F2E}} = \frac{1}{N_A N_B} \sum_{i\in A,\,j\in B} L_{\mathrm{F2E}}(i,j).$$

This approach guarantees constant gradient magnitude for all rotation angles up to $\theta = \pi$, thereby avoiding stagnation in high-error regimes.

5. Group-wise Aggregation and Theoretical Equivalence

Under near-rigidity, aggregating pairwise F2E terms over residues reduces to a group-to-group geodesic between complex-level frames: $$L_{\mathrm{G-F2E}} \approx d_G(\Delta, \widehat{\Delta}),$$ where $\Delta = T_B^{-1}T_A$, $\widehat{\Delta} = \hat{T}_B^{-1}\hat{T}_A$. Neglecting intra-chain noise,

$$L_{\mathrm{G-F2E}} \approx \sqrt{n\,d_\theta^2 + \frac{k^2}{\alpha^2}\|\Delta t\|^2},$$

where $n$ is the chain length and $k^2 = \sum_i \|x_{i,\rm local}\|^2$. This formalism establishes the precise equivalence between aggregated residue-level losses and true geodesic distances between rigid bodies in $\mathrm{SE}(3)$.

6. Empirical Performance in Protein Docking

Fine-tuning AlphaFold2-Multimer using only the F2E (group-aware geodesic) loss—augmented by standard intra-chain auxiliary losses—was benchmarked on SABDab antibody–antigen complexes post–September 2021. The following docking results were obtained (Wu et al., 2024):

Metric	AF2.3 (chordal/FAPE)	AF2.3 + F2E (geodesic)	Relative Change
DockQ > 0.23 (full)	18.6%	52.3%	+182%
DockQ > 0.80 (full)	1.7%	3.9%
Avg DockQ (full)	0.080	0.195
DockQ > 0.23 (low-homology)	21.9%	43.8%	+100%

These results indicate substantial improvements in both standard and challenging low-homology settings, demonstrating the efficacy of group-aware geodesic losses when applied to difficult docking tasks.

7. Broader Implications and Methodological Significance

The formulation of F2E as a group-aware geodesic loss on $\mathrm{SE}(3)$ provides a general framework for structure prediction settings where rigid-body misalignments hinder learning under chordal metrics. The approach extends directly to any application reliant on pose or frame estimation, allowing for stable, rotationally robust optimization. Adoption of such losses represents a step towards greater geometric fidelity in deep structure learning, with demonstrated gains in complex molecular modeling (Wu et al., 2024). A plausible implication is the extension of group-aware geodesic objectives to other structured prediction tasks involving frames, transformations, or manifolds, especially where traditional Euclidean approximations are prone to gradient degeneracy.