add ability to use taylor series linear attention in atom encoder and decoder

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -61,6 +61,7 @@
 )
 
 from alphafold3_pytorch.attention import Attention
+from taylor_series_linear_attention import TaylorSeriesLinearAttn
 
 import einx
 from einops import rearrange, repeat, reduce, einsum, pack, unpack
@@ -1334,7 +1335,12 @@ def __init__(
         attn_window_size = None,
         attn_pair_bias_kwargs: dict = dict(),
         num_register_tokens = 0,
-        serial = False
+        serial = False,
+        use_linear_attn = False,
+        linear_attn_kwargs = dict(
+            heads = 8,
+            dim_head = 16
+        )
     ):
         super().__init__()
         dim_single_cond = default(dim_single_cond, dim)
@@ -1343,6 +1349,15 @@ def __init__(
 
         for _ in range(depth):
 
+            linear_attn = None
+
+            if use_linear_attn:
+                linear_attn = TaylorSeriesLinearAttn(
+                    dim = dim,
+                    prenorm = True,
+                    **linear_attn_kwargs
+                )
+
             pair_bias_attn = AttentionPairBias(
                 dim = dim,
                 dim_pairwise = dim_pairwise,
@@ -1368,6 +1383,7 @@ def __init__(
             )
 
             layers.append(ModuleList([
+                linear_attn,
                 conditionable_pair_bias,
                 conditionable_transition
             ]))
@@ -1408,7 +1424,10 @@ def forward(
 
         # main transformer
 
-        for attn, transition in self.layers:
+        for linear_attn, attn, transition in self.layers:
+
+            if exists(linear_attn):
+                noised_repr = linear_attn(noised_repr, mask = mask) + noised_repr
 
             attn_out = attn(
                 noised_repr,
@@ -1527,7 +1546,12 @@ def __init__(
         serial = False,
         atom_encoder_kwargs: dict = dict(),
         atom_decoder_kwargs: dict = dict(),
-        token_transformer_kwargs: dict = dict()
+        token_transformer_kwargs: dict = dict(),
+        use_linear_attn = False,
+        linear_attn_kwargs: dict = dict(
+            heads = 8,
+            dim_head = 16
+        )
     ):
         super().__init__()
 
@@ -1584,6 +1608,8 @@ def __init__(
             depth = atom_encoder_depth,
             heads = atom_encoder_heads,
             serial = serial,
+            use_linear_attn = use_linear_attn,
+            linear_attn_kwargs = linear_attn_kwargs,
             **atom_encoder_kwargs
         )
 
@@ -1624,6 +1650,8 @@ def __init__(
             depth = atom_decoder_depth,
             heads = atom_decoder_heads,
             serial = serial,
+            use_linear_attn = use_linear_attn,
+            linear_attn_kwargs = linear_attn_kwargs,
             **atom_decoder_kwargs
         )
 

pyproject.toml

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

tests/test_af3.py

@@ -176,16 +176,17 @@ def test_msa_module():
 
     assert pairwise.shape == pairwise_out.shape
 
-
-def test_diffusion_transformer():
+@pytest.mark.parametrize('use_linear_attn', (False, True))
+def test_diffusion_transformer(use_linear_attn):
 
     single = torch.randn(2, 16, 384)
     pairwise = torch.randn(2, 16, 16, 128)
     mask = torch.randint(0, 2, (2, 16)).bool()
 
     diffusion_transformer = DiffusionTransformer(
         depth = 2,
-        heads = 16
+        heads = 16,
+        use_linear_attn = use_linear_attn
     )
 
     single_out = diffusion_transformer(