throw in recurrent depth for pairformer stack

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1021,6 +1021,7 @@ def __init__(
         dim_single = 384,
         dim_pairwise = 128,
         depth = 48,
+        recurrent_depth = 1, # effective depth will be depth * recurrent_depth
         pair_bias_attn_dim_head = 64,
         pair_bias_attn_heads = 16,
         dropout_row_prob = 0.25,
@@ -1058,6 +1059,12 @@ def __init__(
 
         self.layers = layers
 
+        # https://arxiv.org/abs/2405.16039 and https://arxiv.org/abs/2405.15071
+        # although possibly recycling already takes care of this
+
+        assert recurrent_depth > 0
+        self.recurrent_depth = recurrent_depth
+
         self.num_registers = num_register_tokens
         self.has_registers = num_register_tokens > 0
 
@@ -1093,16 +1100,17 @@ def forward(
 
         # main transformer block layers
 
-        for (
-            pairwise_block,
-            pair_bias_attn,
-            single_transition
-        ) in self.layers:
+        for _ in range(self.recurrent_depth):
+            for (
+                pairwise_block,
+                pair_bias_attn,
+                single_transition
+            ) in self.layers:
 
-            pairwise_repr = pairwise_block(pairwise_repr = pairwise_repr, mask = mask)
+                pairwise_repr = pairwise_block(pairwise_repr = pairwise_repr, mask = mask)
 
-            single_repr = pair_bias_attn(single_repr, pairwise_repr = pairwise_repr, mask = mask) + single_repr
-            single_repr = single_transition(single_repr) + single_repr
+                single_repr = pair_bias_attn(single_repr, pairwise_repr = pairwise_repr, mask = mask) + single_repr
+                single_repr = single_transition(single_repr) + single_repr
 
         # splice out registers
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.10"
+version = "0.1.11"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_af3.py

@@ -156,14 +156,18 @@ def test_centre_random_augmentation():
     assert augmented_coords.shape == coords.shape
 
 
-def test_pairformer():
+@pytest.mark.parametrize('recurrent_depth', (1, 2))
+def test_pairformer(
+    recurrent_depth
+):
     single = torch.randn(2, 16, 384)
     pairwise = torch.randn(2, 16, 16, 128)
     mask = torch.randint(0, 2, (2, 16)).bool()
 
     pairformer = PairformerStack(
         depth = 4,
         num_register_tokens = 4,
+        recurrent_depth = recurrent_depth
     )
 
     single_out, pairwise_out = pairformer(