able to restrict atom attention to intramolecular in atom encoder / decoder, so in the case of sequence-local attention, ligand is not like erroneously conducting attention with one end of some polypeptide

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1555,7 +1555,8 @@ def forward(
         *,
         single_repr: Float['b n ds'],
         pairwise_repr: Float['b n n dp'] | Float['b nw w (w*2) dp'],
-        mask: Bool['b n'] | None = None
+        mask: Bool['b n'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None
     ):
         w = self.attn_window_size
         has_windows = exists(w)
@@ -1596,7 +1597,8 @@ def forward(
                 noised_repr,
                 cond = single_repr,
                 pairwise_repr = pairwise_repr,
-                mask = mask
+                mask = mask,
+                windowed_mask = windowed_mask
             )
 
             if serial:
@@ -1806,7 +1808,8 @@ def forward(
         single_inputs_repr: Float['b n dsi'],
         pairwise_trunk: Float['b n n dpt'],
         pairwise_rel_pos_feats: Float['b n n dpr'],
-        molecule_atom_lens: Int['b n']
+        molecule_atom_lens: Int['b n'],
+        atom_parent_ids: Int['b m'] | None = None
     ):
         w = self.atoms_per_window
         device = noised_atom_pos.device
@@ -1887,11 +1890,22 @@ def forward(
 
         atompair_feats = self.atompair_feats_mlp(atompair_feats) + atompair_feats
 
+        # take care of restricting atom attention to be intra molecular, if the atom_parent_ids were passed in
+
+        windowed_mask = None
+
+        if exists(atom_parent_ids):
+            atom_parent_ids_rows = pad_and_window(atom_parent_ids, w)
+            atom_parent_ids_columns = concat_previous_window(atom_parent_ids_rows, dim_seq = 1, dim_window = 2)
+
+            windowed_mask = einx.equal('b n i, b n j -> b n i j', atom_parent_ids_rows, atom_parent_ids_columns)
+
         # atom encoder
 
         atom_feats = self.atom_encoder(
             atom_feats,
             mask = atom_mask,
+            windowed_mask = windowed_mask,
             single_repr = atom_feats_cond,
             pairwise_repr = atompair_feats
         )
@@ -1929,6 +1943,7 @@ def forward(
         atom_feats = self.atom_decoder(
             atom_decoder_input,
             mask = atom_mask,
+            windowed_mask = windowed_mask,
             single_repr = atom_feats_cond,
             pairwise_repr = atompair_feats
         )
@@ -2154,6 +2169,7 @@ def forward(
         pairwise_trunk: Float['b n n dpt'],
         pairwise_rel_pos_feats: Float['b n n dpr'],
         molecule_atom_lens: Int['b n'],
+        atom_parent_ids: Int['b m'] | None = None,
         return_denoised_pos = False,
         additional_molecule_feats: Float[f'b n {ADDITIONAL_MOLECULE_FEATS}'] | None = None,
         add_smooth_lddt_loss = False,
@@ -2181,6 +2197,7 @@ def forward(
                 atom_feats = atom_feats,
                 atom_mask = atom_mask,
                 atompair_feats = atompair_feats,
+                atom_parent_ids = atom_parent_ids,
                 mask = mask,
                 single_trunk_repr = single_trunk_repr,
                 single_inputs_repr = single_inputs_repr,
@@ -3222,6 +3239,7 @@ def forward(
         atom_ids: Int['b m'] | None = None,
         atompair_ids: Int['b m m'] | Int['b nw w1 w2'] | None = None,
         atom_mask: Bool['b m'] | None = None,
+        atom_parent_ids: Int['b m'] | None = None,
         token_bonds: Bool['b n n'] | None = None,
         msa: Float['b s n d'] | None = None,
         msa_mask: Bool['b s'] | None = None,
@@ -3426,6 +3444,7 @@ def forward(
                 num_sample_steps = num_sample_steps,
                 atom_feats = atom_feats,
                 atompair_feats = atompair_feats,
+                atom_parent_ids = atom_parent_ids,
                 atom_mask = atom_mask,
                 mask = mask,
                 single_trunk_repr = single,
@@ -3483,6 +3502,7 @@ def forward(
                     atom_pos,
                     atom_mask,
                     atom_feats,
+                    atom_parent_ids,
                     atompair_feats,
                     mask,
                     pairwise_mask,
@@ -3504,6 +3524,7 @@ def forward(
                         atom_pos,
                         atom_mask,
                         atom_feats,
+                        atom_parent_ids,
                         atompair_feats,
                         mask,
                         pairwise_mask,
@@ -3547,6 +3568,7 @@ def forward(
                 add_bond_loss = diffusion_add_bond_loss,
                 atom_feats = atom_feats,
                 atompair_feats = atompair_feats,
+                atom_parent_ids = atom_parent_ids,
                 atom_mask = atom_mask,
                 mask = mask,
                 single_trunk_repr = single,

alphafold3_pytorch/attention.py

@@ -222,6 +222,7 @@ def forward(
         seq: Float['b i d'],
         mask: Bool['b n']| None = None,
         context: Float['b j d'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None,
         attn_bias: Float['... i j'] | Float['... nw w (w*2)'] | None = None
 
     ) -> Float['b i d']:
@@ -239,6 +240,7 @@ def forward(
             q, k, v,
             attn_bias = attn_bias,
             mask = mask,
+            windowed_mask = windowed_mask,
             memory_kv = self.memory_kv
         )
 
@@ -324,6 +326,7 @@ def local_attn(
         k: Float['b h n d'],
         v: Float['b h n d'],
         mask: Bool['b n'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None,
         attn_bias: Float['... n n'] | Float['... nw w (w*2)'] | None = None,
         memory_kv: Float['2 h m d'] | None = None
     ) -> Float['b h n d']:
@@ -386,6 +389,9 @@ def local_attn(
             if exists(attn_bias):
                 attn_bias = pad_at_dim(attn_bias, (num_mem_kv, 0), value = 0.)
 
+            if exists(windowed_mask):
+                windowed_mask = pad_at_dim(windowed_mask, (num_mem_kv, 0), value = True)
+
             if exists(mask):
                 mask = pad_at_dim(mask, (num_mem_kv, 0), value = True)
 
@@ -400,13 +406,27 @@ def local_attn(
             assert attn_bias.ndim == sim.ndim
             sim = sim + attn_bias
 
+        # windowed masking - for masking out atoms not belonging to the same molecule / polypeptide / nucleic acid in sequence-local attention
+
+        if exists(windowed_mask):
+            sim = einx.where(
+                'b n i j, b h n i j, -> b h n i j',
+                windowed_mask, sim, max_neg_value(sim)
+            )
+
+        # mask out buckets of padding
+
         sim = einx.where(
             'b n j, b h n i j, -> b h n i j',
             mask, sim, max_neg_value(sim)
         )
 
+        # local attention
+
         attn = sim.softmax(dim = -1)
 
+        # aggregate
+
         out = einsum(attn, v, "... i j, ... j d -> ... i d")
 
         # un-window the output
@@ -426,6 +446,7 @@ def forward(
         k: Float['b h j d'],
         v: Float['b h j d'],
         mask: Bool['b j'] | None = None,
+        windowed_mask: Bool['b nw w (w*2)'] | None = None,
         attn_bias: Float['... i j'] | Float['... nw w (w*2)'] | None = None,
         memory_kv: Float['2 h m d'] | None = None
     ) -> Float['b h i d']:
@@ -439,7 +460,7 @@ def forward(
         # todo (handle attn bias efficiently)
 
         if self.is_local_attn:
-            return self.local_attn(q, k, v, mask = mask, attn_bias = attn_bias, memory_kv = memory_kv)
+            return self.local_attn(q, k, v, mask = mask, windowed_mask = windowed_mask, attn_bias = attn_bias, memory_kv = memory_kv)
 
         assert not exists(is_windowed_attn_bias) or not is_windowed_attn_bias
 

pyproject.toml

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

tests/test_af3.py

@@ -412,9 +412,11 @@ def test_distogram_head():
 
 @pytest.mark.parametrize('window_atompair_inputs', (True, False))
 @pytest.mark.parametrize('stochastic_frame_average', (True, False))
+@pytest.mark.parametrize('atom_transformer_intramolecular_attn', (True, False))
 def test_alphafold3(
     window_atompair_inputs: bool,
-    stochastic_frame_average: bool
+    stochastic_frame_average: bool,
+    atom_transformer_intramolecular_attn: bool
 ):
     seq_len = 16
     atoms_per_window = 27
@@ -434,6 +436,11 @@ def test_alphafold3(
     additional_molecule_feats = torch.randn(2, seq_len, 9)
     molecule_ids = torch.randint(0, 32, (2, seq_len))
 
+    atom_parent_ids = None
+
+    if atom_transformer_intramolecular_attn:
+        atom_parent_ids = torch.ones(2, atom_seq_len).long()
+
     template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
     template_mask = torch.ones((2, 2)).bool()
 
@@ -478,6 +485,7 @@ def test_alphafold3(
         atom_inputs = atom_inputs,
         molecule_ids = molecule_ids,
         molecule_atom_lens = molecule_atom_lens,
+        atom_parent_ids = atom_parent_ids,
         atompair_inputs = atompair_inputs,
         additional_molecule_feats = additional_molecule_feats,
         token_bonds = token_bonds,