prepare the architecture for missing atoms, by replacing the noised atom pos features with a missing atom feature, and then removing that atom from the denoising loss and other related losses

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1741,6 +1741,8 @@ def __init__(
 
         self.atom_pos_to_atom_feat = LinearNoBias(3, dim_atom)
 
+        self.missing_atom_feat = nn.Parameter(torch.zeros(dim_atom))
+
         self.single_repr_to_atom_feat_cond = nn.Sequential(
             nn.LayerNorm(dim_single),
             LinearNoBias(dim_single, dim_atom)
@@ -1839,7 +1841,8 @@ 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
+        atom_parent_ids: Int['b m'] | None = None,
+        missing_atom_mask: Bool['b m']| None = None
     ):
         w = self.atoms_per_window
         device = noised_atom_pos.device
@@ -1864,7 +1867,16 @@ def forward(
 
         # the most surprising part of the paper; no geometric biases!
 
-        atom_feats = self.atom_pos_to_atom_feat(noised_atom_pos) + atom_feats
+        noised_atom_pos_feats = self.atom_pos_to_atom_feat(noised_atom_pos)
+
+        # for missing atoms, replace the noise atom pos features with a missing embedding
+
+        if exists(missing_atom_mask):
+            noised_atom_pos_feats = einx.where('b m, d, b m d -> b m d', missing_atom_mask, self.missing_atom_feat, noised_atom_pos_feats)
+
+        # sum the noised atom position features to the atom features
+
+        atom_feats = noised_atom_pos_feats + atom_feats
 
         # condition atom feats cond (cl) with single repr
 
@@ -2199,6 +2211,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'],
+        missing_atom_mask: Bool['b m'] | None = None,
         atom_parent_ids: Int['b m'] | None = None,
         return_denoised_pos = False,
         is_molecule_types: Bool[f'b n {IS_MOLECULE_TYPES}'] | None = None,
@@ -2227,6 +2240,7 @@ def forward(
             network_condition_kwargs = dict(
                 atom_feats = atom_feats,
                 atom_mask = atom_mask,
+                missing_atom_mask = missing_atom_mask,
                 atompair_feats = atompair_feats,
                 atom_parent_ids = atom_parent_ids,
                 mask = mask,
@@ -2282,6 +2296,11 @@ def forward(
 
         losses = losses * loss_weights
 
+        # if there are missing atoms, update the atom mask to not include them in the loss
+
+        if exists(missing_atom_mask):
+            atom_mask = atom_mask & ~ missing_atom_mask
+
         # account for atom mask
 
         mse_loss = losses[atom_mask].mean()
@@ -3337,6 +3356,7 @@ def forward(
         atompair_ids: Int['b m m'] | Int['b nw {self.w} {self.w*2}'] | None = None,
         is_molecule_mod: Bool['b n num_mods'] | None = None,
         atom_mask: Bool['b m'] | None = None,
+        missing_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,
@@ -3656,6 +3676,7 @@ def forward(
                 (
                     atom_pos,
                     atom_mask,
+                    missing_atom_mask,
                     atom_feats,
                     atom_parent_ids,
                     atompair_feats,
@@ -3679,6 +3700,7 @@ def forward(
                     for t in (
                         atom_pos,
                         atom_mask,
+                        missing_atom_mask,
                         atom_feats,
                         atom_parent_ids,
                         atompair_feats,
@@ -3730,6 +3752,7 @@ def forward(
                 atom_feats = atom_feats,
                 atompair_feats = atompair_feats,
                 atom_parent_ids = atom_parent_ids,
+                missing_atom_mask = missing_atom_mask,
                 atom_mask = atom_mask,
                 mask = mask,
                 single_trunk_repr = single,

pyproject.toml

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

tests/test_af3.py

@@ -417,11 +417,13 @@ def test_distogram_head():
 
 @pytest.mark.parametrize('window_atompair_inputs', (True, False))
 @pytest.mark.parametrize('stochastic_frame_average', (True, False))
+@pytest.mark.parametrize('missing_atoms', (True, False))
 @pytest.mark.parametrize('atom_transformer_intramolecular_attn', (True, False))
 @pytest.mark.parametrize('num_molecule_mods', (0, 5))
 def test_alphafold3(
     window_atompair_inputs: bool,
     stochastic_frame_average: bool,
+    missing_atoms: bool,
     atom_transformer_intramolecular_attn: bool,
     num_molecule_mods: int
 ):
@@ -449,6 +451,10 @@ def test_alphafold3(
     if num_molecule_mods > 0:
         is_molecule_mod = torch.zeros(2, seq_len, num_molecule_mods).uniform_(0, 1) < 0.1
 
+    missing_atom_mask = None
+    if missing_atoms:
+        missing_atom_mask = torch.randint(0, 2, (2, atom_seq_len)).bool()
+
     atom_parent_ids = None
 
     if atom_transformer_intramolecular_attn:
@@ -501,6 +507,7 @@ def test_alphafold3(
         molecule_atom_lens = molecule_atom_lens,
         atom_parent_ids = atom_parent_ids,
         atompair_inputs = atompair_inputs,
+        missing_atom_mask = missing_atom_mask,
         is_molecule_types = is_molecule_types,
         is_molecule_mod = is_molecule_mod,
         additional_molecule_feats = additional_molecule_feats,