test the waters with stochastic frame averaging first before going any deeper

Author: lucidrains

README.md

@@ -237,3 +237,14 @@ docker run -v .:/data --gpus all -it af3
     url     = {https://api.semanticscholar.org/CorpusID:238419638}
 }
 ```
+
+```bibtex
+@article{Duval2023FAENetFA,
+    title   = {FAENet: Frame Averaging Equivariant GNN for Materials Modeling},
+    author  = {Alexandre Duval and Victor Schmidt and Alex Hernandez Garcia and Santiago Miret and Fragkiskos D. Malliaros and Yoshua Bengio and David Rolnick},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2305.05577},
+    url     = {https://api.semanticscholar.org/CorpusID:258564608}
+}
+```

alphafold3_pytorch/alphafold3.py

@@ -2108,7 +2108,6 @@ def forward(
         pairwise_trunk: Float['b n n dpt'],
         pairwise_rel_pos_feats: Float['b n n dpr'],
         molecule_atom_lens: Int['b n'],
-        frame_average_fn: Callable[[Float['b n 3']], Float['b n 3']] | None = None,
         return_denoised_pos = False,
         additional_molecule_feats: Float[f'b n {ADDITIONAL_MOLECULE_FEATS}'] | None = None,
         add_smooth_lddt_loss = False,
@@ -2145,11 +2144,6 @@ def forward(
             )
         )
 
-        # frame average the denoised atom positions if needed
-
-        if exists(frame_average_fn):
-            denoised_atom_pos = frame_average_fn(denoised_atom_pos)
-
         # total loss, for accumulating all auxiliary losses
 
         total_loss = 0.
@@ -2905,7 +2899,8 @@ def __init__(
             S_tmax = 50,
             S_noise = 1.003,
         ),
-        augment_kwargs: dict = dict()
+        augment_kwargs: dict = dict(),
+        stochastic_frame_average = False
     ):
         super().__init__()
 
@@ -2918,6 +2913,18 @@ def __init__(
         self.num_augmentations = diffusion_num_augmentations
         self.augmenter = CentreRandomAugmentation(**augment_kwargs)
 
+        # stochastic frame averaging
+        # https://arxiv.org/abs/2305.05577
+
+        self.stochastic_frame_average = stochastic_frame_average
+
+        if stochastic_frame_average:
+            self.frame_average = FrameAverage(
+                dim = 3,
+                stochastic = True,
+                return_stochastic_as_augmented_pos = True
+            )
+
         # input feature embedder
 
         self.input_embedder = InputFeatureEmbedder(
@@ -3330,7 +3337,7 @@ def forward(
 
         if calc_diffusion_loss:
 
-            num_augs = self.num_augmentations
+            num_augs = self.num_augmentations + int(self.stochastic_frame_average)
 
             # take care of augmentation
             # they did 48 during training, as the trunk did the heavy lifting
@@ -3378,8 +3385,25 @@ def forward(
                     )
                 )
 
+                # handle stochastic frame averaging
+
+                if self.stochastic_frame_average:
+                    fa_atom_pos, atom_pos = atom_pos[:1], atom_pos[1:]
+
+                    fa_atom_pos = self.frame_average(
+                        fa_atom_pos,
+                        frame_average_mask = atom_mask[:1]
+                    )
+
+                # normal random augmentations, 48 times in paper
+
                 atom_pos = self.augmenter(atom_pos)
 
+                # concat back the stochastic frame averaged position
+
+                if self.stochastic_frame_average:
+                    atom_pos = torch.cat((fa_atom_pos, atom_pos), dim = 0)
+
             diffusion_loss, denoised_atom_pos, diffusion_loss_breakdown, _ = self.edm(
                 atom_pos,
                 additional_molecule_feats = additional_molecule_feats,

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
     "einx>=0.2.2",
     "ema-pytorch>=0.4.8",
     "environs",
-    "frame-averaging-pytorch>=0.0.17",
+    "frame-averaging-pytorch>=0.0.18",
     "hydra-core",
     "jaxtyping>=0.2.28",
     "lightning>=2.2.5",

tests/test_af3.py

@@ -31,6 +31,9 @@
     atom_ref_pos_to_atompair_inputs
 )
 
+def join(str, delimiter = ','):
+    return delimiter.join(str)
+
 def test_atom_ref_pos_to_atompair_inputs():
     atom_ref_pos = torch.randn(16, 3)
     atom_ref_space_uid = torch.ones(16).long()
@@ -400,9 +403,17 @@ def test_distogram_head():
 
     logits = distogram_head(pairwise_repr)
 
-@pytest.mark.parametrize('window_atompair_inputs', (True, False))
+@pytest.mark.parametrize(
+    join([
+        'window_atompair_inputs',
+        'stochastic_frame_average'
+    ]), [
+        (True, False),
+        (True, False)
+    ])
 def test_alphafold3(
-    window_atompair_inputs: bool
+    window_atompair_inputs: bool,
+    stochastic_frame_average: bool
 ):
     seq_len = 16
     atoms_per_window = 27
@@ -457,6 +468,7 @@ def test_alphafold3(
             token_transformer_depth = 1,
             atom_decoder_depth = 1,
         ),
+        stochastic_frame_average = stochastic_frame_average
     )
 
     loss, breakdown = alphafold3(