Table of Contents
Fetching ...

FAFE: Immune Complex Modeling with Geodesic Distance Loss on Noisy Group Frames

Ruidong Wu, Ruihan Guo, Rui Wang, Shitong Luo, Yue Xu, Jiahan Li, Jianzhu Ma, Qiang Liu, Yunan Luo, Jian Peng

TL;DR

This work tackles the challenge of antibody antigen complex modeling by identifying gradient vanishing in the Frame Aligned Point Error loss when rotation errors are large. It introduces Frame Aligned Frame Error, a geodesic distance based loss that jointly optimizes rotational and translational frame differences, and extends it to a group aware formulation GF2E to handle inter chain pose errors. Through LoRA based parameter efficient fine tuning of AF2 Multimer on antibody antigen data, the approach yields substantial improvements in DockQ metrics, especially for low homology targets. The results demonstrate that reframing inter chain errors as geodesic distances between group frames improves gradient stability and docking accuracy, with broad implications for immune complex modeling and potential extension to non protein components and pre training settings.

Abstract

Despite the striking success of general protein folding models such as AlphaFold2(AF2, Jumper et al. (2021)), the accurate computational modeling of antibody-antigen complexes remains a challenging task. In this paper, we first analyze AF2's primary loss function, known as the Frame Aligned Point Error (FAPE), and raise a previously overlooked issue that FAPE tends to face gradient vanishing problem on high-rotational-error targets. To address this fundamental limitation, we propose a novel geodesic loss called Frame Aligned Frame Error (FAFE, denoted as F2E to distinguish from FAPE), which enables the model to better optimize both the rotational and translational errors between two frames. We then prove that F2E can be reformulated as a group-aware geodesic loss, which translates the optimization of the residue-to-residue error to optimizing group-to-group geodesic frame distance. By fine-tuning AF2 with our proposed new loss function, we attain a correct rate of 52.3\% (DockQ $>$ 0.23) on an evaluation set and 43.8\% correct rate on a subset with low homology, with substantial improvement over AF2 by 182\% and 100\% respectively.

FAFE: Immune Complex Modeling with Geodesic Distance Loss on Noisy Group Frames

TL;DR

This work tackles the challenge of antibody antigen complex modeling by identifying gradient vanishing in the Frame Aligned Point Error loss when rotation errors are large. It introduces Frame Aligned Frame Error, a geodesic distance based loss that jointly optimizes rotational and translational frame differences, and extends it to a group aware formulation GF2E to handle inter chain pose errors. Through LoRA based parameter efficient fine tuning of AF2 Multimer on antibody antigen data, the approach yields substantial improvements in DockQ metrics, especially for low homology targets. The results demonstrate that reframing inter chain errors as geodesic distances between group frames improves gradient stability and docking accuracy, with broad implications for immune complex modeling and potential extension to non protein components and pre training settings.

Abstract

Despite the striking success of general protein folding models such as AlphaFold2(AF2, Jumper et al. (2021)), the accurate computational modeling of antibody-antigen complexes remains a challenging task. In this paper, we first analyze AF2's primary loss function, known as the Frame Aligned Point Error (FAPE), and raise a previously overlooked issue that FAPE tends to face gradient vanishing problem on high-rotational-error targets. To address this fundamental limitation, we propose a novel geodesic loss called Frame Aligned Frame Error (FAFE, denoted as F2E to distinguish from FAPE), which enables the model to better optimize both the rotational and translational errors between two frames. We then prove that F2E can be reformulated as a group-aware geodesic loss, which translates the optimization of the residue-to-residue error to optimizing group-to-group geodesic frame distance. By fine-tuning AF2 with our proposed new loss function, we attain a correct rate of 52.3\% (DockQ 0.23) on an evaluation set and 43.8\% correct rate on a subset with low homology, with substantial improvement over AF2 by 182\% and 100\% respectively.
Paper Structure (34 sections, 5 theorems, 38 equations, 8 figures, 1 table)

This paper contains 34 sections, 5 theorems, 38 equations, 8 figures, 1 table.

Key Result

Lemma 4.2

Given two rotations $R_i$ and $R_j \in \mathop{\mathrm{\mathrm{SO}(3)}}\nolimits3$, we have in which $d_c(\cdot)$ denotes the chordal distance defined in def:chordal_so3 and $d_{\theta}(\cdot)$ denotes the geodesic distance defined in def:geodesic_so3. The proof is provided in proof:geo2chordal.

Figures (8)

  • Figure 1: Problem and method overview. (a) An example from PDB 7XJF, the gray structure is predicted by AF2, and the green structure is the experimental ground truth. We can see there is a large rotation error. (b) Statistics of rotational errors in AF2.3 predictions on our evaluation set. Most of the predictions have a rotation error larger than $\frac{\pi}{2}$. (c) The comparison between FAPE (left) and F2E (right). FAPE calculates point-wise Euclidean distance error after alignment to local frames, while F2E calculates frame-wise geodesic distance error after the same alignment.
  • Figure 2: The loss value regarding rotation angle (difference between predicted and groundtruth) for chordal and geodesic distance. The gradient of chordal distance tends to vanish when the rotation angle approaches $\pi$, while the geodesic distance provides a stable gradient from $0$ to $\pi$.
  • Figure 3: Comparison of F2E with AF2.3 on two held evaluation sets. (a) The evaluation results on full evaluation set. For columns "DockQ$>$0.23" and "DockQ$>$0.8" the percentage is reported, and average DockQ is reported in the last column. (b) The evaluation results on low homology set with novel antigens.
  • Figure 4: (a) The visualization of rotational errors for AF2.3 and F2E. (b) The predicted structures of PDB 7Y1B against ground truth. The grey structure is the experimental ground truth, the green structure by AF2.3 and the red structure by model tuned with F2E. We show in red arrows how to rotate from the wrong prediction to the ground truth. (c) The predicted structures of 8GPT against ground truth.
  • Figure 5: The predicted structures of PDB 8DB4 against ground truth. The grey structure is the experimental ground truth, the green structure by AF2.3 and the red structure by model tuned with F2E.
  • ...and 3 more figures

Theorems & Definitions (23)

  • Definition 2.1
  • Definition 2.2
  • Definition 2.3
  • Definition 2.4
  • Definition 2.5
  • Definition 2.6
  • Definition 2.7
  • Definition 3.1
  • Definition 3.2
  • Definition 4.1
  • ...and 13 more