minimal viable working implementation of atomic resolution confidence heads

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3137,7 +3137,8 @@ def __init__(
             S_noise = 1.003,
         ),
         augment_kwargs: dict = dict(),
-        stochastic_frame_average = False
+        stochastic_frame_average = False,
+        confidence_head_atom_resolution = False
     ):
         super().__init__()
 
@@ -3315,6 +3316,7 @@ def __init__(
             num_plddt_bins = num_plddt_bins,
             num_pde_bins = num_pde_bins,
             num_pae_bins = num_pae_bins,
+            atom_resolution = confidence_head_atom_resolution,
             **confidence_head_kwargs
         )
 
@@ -3434,11 +3436,11 @@ def forward(
         molecule_atom_indices: Int['b n'] | None = None, # the 'token centre atoms' mentioned in the paper, unsure where it is used in the architecture
         num_sample_steps: int | None = None,
         atom_pos: Float['b m 3'] | None = None,
-        distance_labels: Int['b n n'] | None = None,
-        pae_labels: Int['b n n'] | None = None,
-        pde_labels: Int['b n n'] | None = None,
-        plddt_labels: Int['b n'] | None = None,
-        resolved_labels: Int['b n'] | None = None,
+        distance_labels: Int['b n n'] | Int['b m m'] | None = None,
+        pae_labels: Int['b n n'] | Int['b m m'] | None = None,
+        pde_labels: Int['b n n'] | Int['b m m'] | None = None,
+        plddt_labels: Int['b n'] | Int['b m'] | None = None,
+        resolved_labels: Int['b n'] | Int['b m'] | None = None,
         return_loss_breakdown = False,
         return_loss: bool = None,
         return_present_sampled_atoms: bool = False,
@@ -3710,6 +3712,8 @@ def forward(
                 pairwise_repr = pairwise.detach(),
                 pred_atom_pos = confidence_head_atom_pos_input.detach(),
                 molecule_atom_indices = molecule_atom_indices,
+                molecule_atom_lens = molecule_atom_lens,
+                atom_feats = atom_feats,
                 mask = mask,
                 return_pae_logits = True
             )
@@ -3884,24 +3888,37 @@ def forward(
                 pairwise_repr = pairwise.detach(),
                 pred_atom_pos = denoised_atom_pos.detach(),
                 molecule_atom_indices = molecule_atom_indices,
+                molecule_atom_lens = molecule_atom_lens,
                 mask = mask,
+                atom_feats = atom_feats,
                 return_pae_logits = return_pae_logits
             )
 
+            # determine which mask to use for labels depending on atom resolution or not for confidence head
+
+            label_mask = mask
+
+            if self.confidence_head.atom_resolution:
+                label_mask = atom_mask
+
+            label_pairwise_mask = einx.logical_and('... i, ... j -> ... i j', label_mask, label_mask)
+
+            # cross entropy losses
+
             if exists(pae_labels):
-                pae_labels = torch.where(pairwise_mask, pae_labels, ignore)
+                pae_labels = torch.where(label_pairwise_mask, pae_labels, ignore)
                 pae_loss = F.cross_entropy(ch_logits.pae, pae_labels, ignore_index = ignore)
 
             if exists(pde_labels):
-                pde_labels = torch.where(pairwise_mask, pde_labels, ignore)
+                pde_labels = torch.where(label_pairwise_mask, pde_labels, ignore)
                 pde_loss = F.cross_entropy(ch_logits.pde, pde_labels, ignore_index = ignore)
 
             if exists(plddt_labels):
-                plddt_labels = torch.where(mask, plddt_labels, ignore)
+                plddt_labels = torch.where(label_mask, plddt_labels, ignore)
                 plddt_loss = F.cross_entropy(ch_logits.plddt, plddt_labels, ignore_index = ignore)
 
             if exists(resolved_labels):
-                resolved_labels = torch.where(mask, resolved_labels, ignore)
+                resolved_labels = torch.where(label_mask, resolved_labels, ignore)
                 resolved_loss = F.cross_entropy(ch_logits.resolved, resolved_labels, ignore_index = ignore)
 
             confidence_loss = pae_loss + pde_loss + plddt_loss + resolved_loss

pyproject.toml

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

tests/test_af3.py

@@ -459,12 +459,14 @@ def test_distogram_head():
 @pytest.mark.parametrize('missing_atoms', (True, False))
 @pytest.mark.parametrize('atom_transformer_intramolecular_attn', (True, False))
 @pytest.mark.parametrize('num_molecule_mods', (0, 5))
+@pytest.mark.parametrize('confidence_head_atom_resolution', (True, False))
 def test_alphafold3(
     window_atompair_inputs: bool,
     stochastic_frame_average: bool,
     missing_atoms: bool,
     atom_transformer_intramolecular_attn: bool,
-    num_molecule_mods: int
+    num_molecule_mods: int,
+    confidence_head_atom_resolution: bool
 ):
     seq_len = 16
     atoms_per_window = 27
@@ -509,10 +511,11 @@ def test_alphafold3(
     distogram_atom_indices = molecule_atom_lens - 1
     molecule_atom_indices = molecule_atom_lens - 1
 
-    pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
-    pde_labels = torch.randint(0, 64, (2, seq_len, seq_len))
-    plddt_labels = torch.randint(0, 50, (2, seq_len))
-    resolved_labels = torch.randint(0, 2, (2, seq_len))
+    label_len = atom_seq_len if confidence_head_atom_resolution else seq_len
+    pae_labels = torch.randint(0, 64, (2, label_len, label_len))
+    pde_labels = torch.randint(0, 64, (2, label_len, label_len))
+    plddt_labels = torch.randint(0, 50, (2, label_len))
+    resolved_labels = torch.randint(0, 2, (2, label_len))
 
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
@@ -538,7 +541,8 @@ def test_alphafold3(
             token_transformer_depth = 1,
             atom_decoder_depth = 1,
         ),
-        stochastic_frame_average = stochastic_frame_average
+        stochastic_frame_average = stochastic_frame_average,
+        confidence_head_atom_resolution = confidence_head_atom_resolution
     )
 
     loss, breakdown = alphafold3(