add ability to use conditionally routed attention from CoLT5 paper for alleviating sequence-local issue with atoms

Author: lucidrains

README.md

@@ -263,3 +263,11 @@ docker run -v .:/data --gpus all -it af3
     url     = {https://api.semanticscholar.org/CorpusID:247187905}
 }
 ```
+
+```bibtex
+@inproceedings{Ainslie2023CoLT5FL,
+    title   = {CoLT5: Faster Long-Range Transformers with Conditional Computation},
+    author  = {Joshua Ainslie and Tao Lei and Michiel de Jong and Santiago Ontan'on and Siddhartha Brahma and Yury Zemlyanskiy and David Uthus and Mandy Guo and James Lee-Thorp and Yi Tay and Yun-Hsuan Sung and Sumit Sanghai},
+    year    = {2023}
+}
+```

alphafold3_pytorch/alphafold3.py

@@ -41,6 +41,8 @@
 
 from taylor_series_linear_attention import TaylorSeriesLinearAttn
 
+from colt5_attention import ConditionalRoutedAttention
+
 import einx
 from einops import rearrange, repeat, reduce, einsum, pack, unpack
 from einops.layers.torch import Rearrange
@@ -1460,7 +1462,15 @@ def __init__(
         linear_attn_kwargs = dict(
             heads = 8,
             dim_head = 16
+        ),
+        use_colt5_attn = False,
+        colt5_attn_kwargs = dict(
+            heavy_dim_head = 64,
+            heavy_heads = 8,
+            num_heavy_tokens_q = 512,
+            num_heavy_tokens_kv = 512
         )
+
     ):
         super().__init__()
         self.attn_window_size = attn_window_size
@@ -1481,6 +1491,15 @@ def __init__(
                     **linear_attn_kwargs
                 )
 
+            colt5_attn = None
+
+            if use_colt5_attn:
+                colt5_attn = ConditionalRoutedAttention(
+                    dim = dim,
+                    has_light_attn = False,
+                    **colt5_attn_kwargs
+                )
+
             pair_bias_attn = AttentionPairBias(
                 dim = dim,
                 dim_pairwise = dim_pairwise,
@@ -1508,6 +1527,7 @@ def __init__(
 
             layers.append(ModuleList([
                 linear_attn,
+                colt5_attn,
                 conditionable_pair_bias,
                 conditionable_transition
             ]))
@@ -1560,11 +1580,14 @@ def forward(
 
         # main transformer
 
-        for linear_attn, attn, transition in self.layers:
+        for linear_attn, colt5_attn, attn, transition in self.layers:
 
             if exists(linear_attn):
                 noised_repr = linear_attn(noised_repr, mask = mask) + noised_repr
 
+            if exists(colt5_attn):
+                noised_repr = colt5_attn(noised_repr, mask = mask) + noised_repr
+
             attn_out = attn(
                 noised_repr,
                 cond = single_repr,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.51"
+version = "0.1.52"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -25,6 +25,7 @@ classifiers=[
 dependencies = [
     "beartype",
     "biopython>=1.83",
+    "CoLT5-attention>=0.11.0",
     "einops>=0.8.0",
     "einx>=0.2.2",
     "ema-pytorch>=0.5.0",

tests/test_af3.py

@@ -206,7 +206,8 @@ def test_msa_module():
     assert pairwise.shape == pairwise_out.shape
 
 @pytest.mark.parametrize('use_linear_attn', (False, True))
-def test_diffusion_transformer(use_linear_attn):
+@pytest.mark.parametrize('use_colt5_attn', (False, True))
+def test_diffusion_transformer(use_linear_attn, use_colt5_attn):
 
     single = torch.randn(2, 16, 384)
     pairwise = torch.randn(2, 16, 16, 128)
@@ -215,7 +216,8 @@ def test_diffusion_transformer(use_linear_attn):
     diffusion_transformer = DiffusionTransformer(
         depth = 2,
         heads = 16,
-        use_linear_attn = use_linear_attn
+        use_linear_attn = use_linear_attn,
+        use_colt5_attn = use_colt5_attn
     )
 
     single_out = diffusion_transformer(