researcher pass in constructed atompair inputs, which contains offset relative positions + inverse squared distance + mask itself, which is then all projected by one LinearNoBias within InputFeatureEmbedder

Author: lucidrains

README.md

@@ -37,8 +37,9 @@ seq_len = 16
 atom_seq_len = seq_len * 27
 
 atom_inputs = torch.randn(2, atom_seq_len, 77)
+atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
+
 atom_lens = torch.randint(0, 27, (2, seq_len))
-atompair_feats = torch.randn(2, atom_seq_len, atom_seq_len, 16)
 additional_residue_feats = torch.randn(2, seq_len, 10)
 
 template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
@@ -63,8 +64,8 @@ resolved_labels = torch.randint(0, 2, (2, seq_len))
 loss = alphafold3(
     num_recycling_steps = 2,
     atom_inputs = atom_inputs,
+    atompair_inputs = atompair_inputs,
     residue_atom_lens = atom_lens,
-    atompair_feats = atompair_feats,
     additional_residue_feats = additional_residue_feats,
     msa = msa,
     msa_mask = msa_mask,
@@ -87,8 +88,8 @@ sampled_atom_pos = alphafold3(
     num_recycling_steps = 4,
     num_sample_steps = 16,
     atom_inputs = atom_inputs,
+    atompair_inputs = atompair_inputs,
     residue_atom_lens = atom_lens,
-    atompair_feats = atompair_feats,
     additional_residue_feats = additional_residue_feats,
     msa = msa,
     msa_mask = msa_mask,

alphafold3_pytorch/alphafold3.py

@@ -14,7 +14,9 @@
 ds - feature dimension (single)
 dp - feature dimension (pairwise)
 dap - feature dimension (atompair)
+dapi - feature dimension (atompair input)
 da - feature dimension (atom)
+dai - feature dimension (atom input)
 t - templates
 s - msa
 r - registers
@@ -2446,6 +2448,7 @@ def __init__(
         self,
         *,
         dim_atom_inputs,
+        dim_atompair_inputs = 5,
         atoms_per_window = 27,
         dim_atom = 128,
         dim_atompair = 16,
@@ -2461,6 +2464,8 @@ def __init__(
 
         self.to_atom_feats = LinearNoBias(dim_atom_inputs, dim_atom)
 
+        self.to_atompair_feats = LinearNoBias(dim_atompair_inputs, dim_atompair)
+
         self.atom_repr_to_atompair_feat_cond = nn.Sequential(
             nn.LayerNorm(dim_atom),
             LinearNoBiasThenOuterSum(dim_atom, dim_atompair),
@@ -2501,8 +2506,8 @@ def forward(
         self,
         *,
         atom_inputs: Float['b m dai'],
+        atompair_inputs: Float['b m m dapi'],
         atom_mask: Bool['b m'],
-        atompair_feats: Float['b m m dap'],
         additional_residue_feats: Float[f'b n {ADDITIONAL_RESIDUE_FEATS}'],
         residue_atom_lens: Int['b n'] | None = None,
 
@@ -2513,6 +2518,7 @@ def forward(
         w = self.atoms_per_window
 
         atom_feats = self.to_atom_feats(atom_inputs)
+        atompair_feats = self.to_atompair_feats(atompair_inputs)
 
         atom_feats_cond = self.atom_repr_to_atompair_feat_cond(atom_feats)
         atompair_feats = atom_feats_cond + atompair_feats
@@ -2709,6 +2715,7 @@ def __init__(
         dim_template_model = 64,
         atoms_per_window = 27,
         dim_atom = 128,
+        dim_atompair_inputs = 5,
         dim_atompair = 16,
         dim_input_embedder_token = 384,
         dim_single = 384,
@@ -2810,6 +2817,7 @@ def __init__(
 
         self.input_embedder = InputFeatureEmbedder(
             dim_atom_inputs = dim_atom_inputs,
+            dim_atompair_inputs = dim_atompair_inputs,
             atoms_per_window = atoms_per_window,
             dim_atom = dim_atom,
             dim_atompair = dim_atompair,
@@ -2995,7 +3003,7 @@ def forward(
         self,
         *,
         atom_inputs: Float['b m dai'],
-        atompair_feats: Float['b m m dap'],
+        atompair_inputs: Float['b m m dapi'],
         additional_residue_feats: Float[f'b n {ADDITIONAL_RESIDUE_FEATS}'],
         residue_atom_lens: Int['b n'] | None = None,
         atom_mask: Bool['b m'] | None = None,
@@ -3063,8 +3071,8 @@ def forward(
             atompair_feats
         ) = self.input_embedder(
             atom_inputs = atom_inputs,
+            atompair_inputs = atompair_inputs,
             atom_mask = atom_mask,
-            atompair_feats = atompair_feats,
             additional_residue_feats = additional_residue_feats,
             residue_atom_lens = residue_atom_lens
         )

alphafold3_pytorch/trainer.py

@@ -25,7 +25,7 @@
 class Alphafold3Input(TypedDict):
     atom_inputs:                Float['m dai']
     residue_atom_lens:          Int['n 2']
-    atompair_feats:             Float['m m dap']
+    atompair_inputs:            Float['m m dap']
     additional_residue_feats:   Float['n 10']
     templates:                  Float['t n n dt']
     template_mask:              Bool['t'] | None

pyproject.toml

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

tests/test_af3.py

@@ -353,8 +353,9 @@ def test_input_embedder():
 
     atom_seq_len = 16 * 27
     atom_inputs = torch.randn(2, atom_seq_len, 77)
+    atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
+
     atom_mask = torch.ones((2, atom_seq_len)).bool()
-    atompair_feats = torch.randn(2, atom_seq_len, atom_seq_len, 16)
     additional_residue_feats = torch.randn(2, 16, 10)
 
     embedder = InputFeatureEmbedder(
@@ -364,7 +365,7 @@ def test_input_embedder():
     embedder(
         atom_inputs = atom_inputs,
         atom_mask = atom_mask,
-        atompair_feats = atompair_feats,
+        atompair_inputs = atompair_inputs,
         additional_residue_feats = additional_residue_feats
     )
 
@@ -383,8 +384,9 @@ def test_alphafold3():
     token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
+    atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
+
     atom_lens = torch.randint(0, 27, (2, seq_len))
-    atompair_feats = torch.randn(2, atom_seq_len, atom_seq_len, 16)
     additional_residue_feats = torch.randn(2, seq_len, 10)
 
     template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
@@ -428,7 +430,7 @@ def test_alphafold3():
         num_recycling_steps = 2,
         atom_inputs = atom_inputs,
         residue_atom_lens = atom_lens,
-        atompair_feats = atompair_feats,
+        atompair_inputs = atompair_inputs,
         additional_residue_feats = additional_residue_feats,
         token_bond = token_bond,
         msa = msa,
@@ -451,7 +453,7 @@ def test_alphafold3():
         num_sample_steps = 16,
         atom_inputs = atom_inputs,
         residue_atom_lens = atom_lens,
-        atompair_feats = atompair_feats,
+        atompair_inputs = atompair_inputs,
         additional_residue_feats = additional_residue_feats,
         msa = msa,
         templates = template_feats,
@@ -465,8 +467,8 @@ def test_alphafold3_without_msa_and_templates():
     atom_seq_len = seq_len * 27
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
+    atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
     atom_lens = torch.randint(0, 27, (2, seq_len))
-    atompair_feats = torch.randn(2, atom_seq_len, atom_seq_len, 16)
     additional_residue_feats = torch.randn(2, seq_len, 10)
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
@@ -505,7 +507,7 @@ def test_alphafold3_without_msa_and_templates():
         num_recycling_steps = 2,
         atom_inputs = atom_inputs,
         residue_atom_lens = atom_lens,
-        atompair_feats = atompair_feats,
+        atompair_inputs = atompair_inputs,
         additional_residue_feats = additional_residue_feats,
         atom_pos = atom_pos,
         residue_atom_indices = residue_atom_indices,
@@ -528,8 +530,8 @@ def test_alphafold3_with_packed_atom_repr():
     token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
+    atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
 
-    atompair_feats = torch.randn(2, atom_seq_len, atom_seq_len, 16)
     additional_residue_feats = torch.randn(2, seq_len, 10)
 
     template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
@@ -574,7 +576,7 @@ def test_alphafold3_with_packed_atom_repr():
         num_recycling_steps = 2,
         atom_inputs = atom_inputs,
         residue_atom_lens = residue_atom_lens,
-        atompair_feats = atompair_feats,
+        atompair_inputs = atompair_inputs,
         additional_residue_feats = additional_residue_feats,
         token_bond = token_bond,
         msa = msa,
@@ -596,7 +598,7 @@ def test_alphafold3_with_packed_atom_repr():
         num_sample_steps = 16,
         atom_inputs = atom_inputs,
         residue_atom_lens = residue_atom_lens,
-        atompair_feats = atompair_feats,
+        atompair_inputs = atompair_inputs,
         additional_residue_feats = additional_residue_feats,
         msa = msa,
         templates = template_feats,

tests/test_trainer.py

@@ -34,8 +34,9 @@ def __getitem__(self, idx):
         atom_seq_len = self.atom_seq_len
 
         atom_inputs = torch.randn(atom_seq_len, 77)
+        atompair_inputs = torch.randn(atom_seq_len, atom_seq_len, 5)
+
         residue_atom_lens = torch.randint(0, 27, (seq_len,))
-        atompair_feats = torch.randn(atom_seq_len, atom_seq_len, 16)
         additional_residue_feats = torch.randn(seq_len, 10)
 
         templates = torch.randn(2, seq_len, seq_len, 44)
@@ -57,8 +58,8 @@ def __getitem__(self, idx):
 
         return Alphafold3Input(
             atom_inputs = atom_inputs,
+            atompair_inputs = atompair_inputs,
             residue_atom_lens = residue_atom_lens,
-            atompair_feats = atompair_feats,
             additional_residue_feats = additional_residue_feats,
             templates = templates,
             template_mask = template_mask,