Add TM head and logging for predicted TM score in AlphaFold2 model

- Introduced TM head with configurable parameters in model configurations.
- Updated AlphaFold2 class to include TM logits and predicted TM score.
- Enhanced training and evaluation scripts to log predicted TM metrics.
- Modified YAML configuration files to support new TM head settings.
- Added tests to validate TM head functionality and logging behavior.

Author: pablo-reyes8

config/experiments/af2_canonical.yaml

@@ -67,6 +67,9 @@ model:
   masked_msa_num_classes: 23
   masked_msa_head_enabled: true
   plddt_head_enabled: true
+  tm_num_bins: 64
+  tm_max_error: 31.5
+  tm_head_enabled: false
   torsion_head_enabled: true
 
 loss:

config/experiments/af2_low_vram.yaml

@@ -70,6 +70,9 @@ model:
   masked_msa_num_classes: 23
   masked_msa_head_enabled: true
   plddt_head_enabled: true
+  tm_num_bins: 64
+  tm_max_error: 31.5
+  tm_head_enabled: false
   torsion_head_enabled: true
 
 loss:

config/experiments/af2_poc.yaml

@@ -64,6 +64,9 @@ model:
   masked_msa_num_classes: 23
   masked_msa_head_enabled: true
   plddt_head_enabled: true
+  tm_num_bins: 64
+  tm_max_error: 31.5
+  tm_head_enabled: false
   torsion_head_enabled: true
 
 loss:

model/alphafold2.py

@@ -35,6 +35,7 @@ class AlphaFold2(nn.Module):
       - torsion angles
       - pLDDT
       - distogram logits
+      - optional TM logits and pTM
     """
     @staticmethod
     def _normalize_ablation_id(ablation):
@@ -72,6 +73,7 @@ def resolve_ablation_defaults(cls, ablation):
                 "distogram_head_enabled": False,
                 "masked_msa_head_enabled": False,
                 "plddt_head_enabled": False,
+                "tm_head_enabled": False,
                 "torsion_head_enabled": False,
             },
             4: {
@@ -105,6 +107,8 @@ def __init__(
         dist_bins=64,
         masked_msa_num_classes=23,
         plddt_bins=50,
+        tm_num_bins=64,
+        tm_max_error=31.5,
         n_torsions=7,
         num_res_blocks_torsion=2,
         recycle_min_bin=3.25,
@@ -132,6 +136,7 @@ def __init__(
         distogram_head_enabled=True,
         masked_msa_head_enabled=True,
         plddt_head_enabled=True,
+        tm_head_enabled=False,
         torsion_head_enabled=True):
 
         super().__init__()
@@ -164,6 +169,7 @@ def __init__(
         distogram_head_enabled = ablation_defaults.get("distogram_head_enabled", distogram_head_enabled)
         masked_msa_head_enabled = ablation_defaults.get("masked_msa_head_enabled", masked_msa_head_enabled)
         plddt_head_enabled = ablation_defaults.get("plddt_head_enabled", plddt_head_enabled)
+        tm_head_enabled = ablation_defaults.get("tm_head_enabled", tm_head_enabled)
         torsion_head_enabled = ablation_defaults.get("torsion_head_enabled", torsion_head_enabled)
 
         self.ablation = self._normalize_ablation_id(ablation)
@@ -185,6 +191,7 @@ def __init__(
         self.distogram_head_enabled = bool(distogram_head_enabled)
         self.masked_msa_head_enabled = bool(masked_msa_head_enabled)
         self.plddt_head_enabled = bool(plddt_head_enabled)
+        self.tm_head_enabled = bool(tm_head_enabled)
         self.torsion_head_enabled = bool(torsion_head_enabled)
 
 
@@ -222,6 +229,7 @@ def __init__(
         self.plddt_head = PlddtHead(c_s=c_s, num_bins=plddt_bins)
         self.distogram_head = DistogramHead(c_z=c_z, num_bins=dist_bins)
         self.masked_msa_head = MaskedMsaHead(c_m=c_m, num_classes=masked_msa_num_classes)
+        self.tm_head = TMHead(c_z=c_z, num_bins=tm_num_bins, max_error=tm_max_error)
         self.torsion_head = TorsionHead(c_s=c_s, n_torsions=n_torsions , num_res_blocks = num_res_blocks_torsion)
         self.recycling_embedder = RecyclingEmbedder(
             c_m=c_m,
@@ -251,6 +259,7 @@ def __init__(
         zero_init_linear(self.plddt_head.mlp[-1])
         zero_init_linear(self.distogram_head.linear)
         zero_init_linear(self.masked_msa_head.linear)
+        zero_init_linear(self.tm_head.linear)
 
         self._freeze_module(self.evoformer, enabled=self.evoformer_enabled)
         self._freeze_module(self.extra_msa_stack, enabled=self.extra_msa_stack_enabled)
@@ -261,6 +270,7 @@ def __init__(
         self._freeze_module(self.distogram_head, enabled=self.distogram_head_enabled)
         self._freeze_module(self.masked_msa_head, enabled=self.masked_msa_head_enabled)
         self._freeze_module(self.plddt_head, enabled=self.plddt_head_enabled)
+        self._freeze_module(self.tm_head, enabled=self.tm_head_enabled)
         self._freeze_module(self.torsion_head, enabled=self.torsion_head_enabled)
 
     @staticmethod
@@ -420,6 +430,10 @@ def forward(
             masked_msa_logits = None
             if self.masked_msa_head_enabled:
                 masked_msa_logits = self.masked_msa_head(m[:, :original_msa_depth])
+            if self.tm_head_enabled:
+                tm_logits, ptm = self.tm_head(z, residue_mask=seq_mask)
+            else:
+                tm_logits, ptm = None, None
 
             # single repr + structure
             s0 = self.single_proj(m)
@@ -484,6 +498,8 @@ def forward(
                 "plddt": plddt,
                 "distogram_logits": distogram_logits,
                 "masked_msa_logits": masked_msa_logits,
+                "tm_logits": tm_logits,
+                "ptm": ptm,
             }
 
             if recycle_idx < num_recycles:

model/alphafold2_heads.py

@@ -1,8 +1,8 @@
 """Prediction heads built on top of the shared AlphaFold representations.
 
 The classes in this module project internal sequence or pair features into the
-single representation, pLDDT logits, and distogram logits used by the model
-output dictionary and downstream loss computation.
+single representation, pLDDT logits, distogram logits, masked-MSA logits, and
+an optional predicted-TM head used for confidence-style reporting.
 """
 
 import torch
@@ -75,3 +75,92 @@ def __init__(self, c_m=256, num_classes=23):
     def forward(self, m):
         logits = self.linear(self.ln(m))
         return logits
+
+
+def compute_predicted_tm_score(
+    tm_logits: torch.Tensor,
+    *,
+    residue_mask: torch.Tensor | None = None,
+    bin_centers: torch.Tensor | None = None,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    """AlphaFold pTM lower bound from pairwise error logits.
+
+    Parameters
+    ----------
+    tm_logits : [B, L, L, num_bins]
+        Logits over aligned-error bins derived from the final pair representation.
+    residue_mask : [B, L], optional
+        Valid residues to include in the domain / chain subset.
+    bin_centers : [num_bins], optional
+        Representative error values for each aligned-error bin.
+    eps : float
+        Small numerical constant.
+    """
+
+    if tm_logits.ndim != 4:
+        raise ValueError(f"tm_logits must have shape [B, L, L, C], got {tuple(tm_logits.shape)}")
+
+    batch_size, length, _, num_bins = tm_logits.shape
+    if bin_centers is None:
+        if num_bins <= 1:
+            bin_width = 0.5
+        else:
+            bin_width = 31.5 / float(num_bins - 1)
+        bin_centers = torch.arange(num_bins, device=tm_logits.device, dtype=tm_logits.dtype)
+        bin_centers = bin_width * (bin_centers + 0.5)
+    else:
+        bin_centers = bin_centers.to(device=tm_logits.device, dtype=tm_logits.dtype)
+        if bin_centers.numel() != num_bins:
+            raise ValueError(
+                f"bin_centers must have {num_bins} entries, got {bin_centers.numel()}"
+            )
+
+    if residue_mask is None:
+        residue_mask = torch.ones(batch_size, length, device=tm_logits.device, dtype=tm_logits.dtype)
+    else:
+        residue_mask = residue_mask.to(device=tm_logits.device, dtype=tm_logits.dtype)
+
+    num_res = residue_mask.sum(dim=-1).clamp_min(1.0)
+    d0 = 1.24 * torch.clamp(num_res, min=19.0).sub(15.0).pow(1.0 / 3.0) - 1.8
+    d0 = d0.clamp_min(0.5)
+
+    probs = F.softmax(tm_logits, dim=-1)
+    tm_kernel = 1.0 / (1.0 + (bin_centers.view(1, 1, 1, -1) / (d0.view(-1, 1, 1, 1) + eps)) ** 2)
+    expected_tm = (probs * tm_kernel).sum(dim=-1)
+
+    per_alignment = (expected_tm * residue_mask[:, None, :]).sum(dim=-1) / num_res.view(-1, 1)
+    per_alignment = per_alignment.masked_fill(residue_mask <= 0, float("-inf"))
+
+    ptm = per_alignment.max(dim=-1).values
+    has_valid = residue_mask.sum(dim=-1) > 0
+    ptm = torch.where(has_valid, ptm, torch.zeros_like(ptm))
+    return ptm
+
+
+class TMHead(nn.Module):
+    def __init__(self, c_z=128, num_bins=64, max_error=31.5):
+        super().__init__()
+        self.num_bins = int(num_bins)
+        self.max_error = float(max_error)
+        self.ln = nn.LayerNorm(c_z)
+        self.linear = nn.Linear(c_z, self.num_bins)
+
+        if self.num_bins <= 1:
+            bin_width = 0.5
+        else:
+            bin_width = self.max_error / float(self.num_bins - 1)
+        bin_centers = bin_width * (torch.arange(self.num_bins, dtype=torch.float32) + 0.5)
+        self.register_buffer("bin_centers", bin_centers, persistent=False)
+
+    def compute_ptm(self, tm_logits, residue_mask=None):
+        return compute_predicted_tm_score(
+            tm_logits,
+            residue_mask=residue_mask,
+            bin_centers=self.bin_centers,
+        )
+
+    def forward(self, z, residue_mask=None):
+        logits = self.linear(self.ln(z))
+        ptm = self.compute_ptm(logits, residue_mask=residue_mask)
+        return logits, ptm

scripts/ablations/run_suite.py

@@ -148,6 +148,7 @@ def _write_comparison_tables(rows: list[dict], output_dir: Path) -> None:
         "msa_loss",
         "plddt_loss",
         "torsion_loss",
+        "ptm_logged",
         "rmsd_logged",
         "tm_score_logged",
         "gdt_ts_logged",
@@ -200,6 +201,7 @@ def main(argv: Sequence[str] | None = None) -> None:
                 "msa_loss": stats.get("msa_loss"),
                 "plddt_loss": stats.get("plddt_loss"),
                 "torsion_loss": stats.get("torsion_loss"),
+                "ptm_logged": stats.get("ptm_logged"),
                 "rmsd_logged": stats.get("rmsd_logged"),
                 "tm_score_logged": stats.get("tm_score_logged"),
                 "gdt_ts_logged": stats.get("gdt_ts_logged"),

tests/test_ablation_suite.py

@@ -187,6 +187,7 @@ def test_alphafold2_ablation_defaults_are_explicit_and_baseline_safe():
     assert AlphaFold2.resolve_ablation_defaults(2)["recycle_single_enabled"] is False
     assert AlphaFold2.resolve_ablation_defaults(3)["masked_msa_head_enabled"] is False
     assert AlphaFold2.resolve_ablation_defaults(3)["plddt_head_enabled"] is False
+    assert AlphaFold2.resolve_ablation_defaults(3)["tm_head_enabled"] is False
     assert AlphaFold2.resolve_ablation_defaults(4)["use_block_specific_params"] is True
     assert AlphaFold2.resolve_ablation_defaults(5)["recycle_single_enabled"] is False
     assert AlphaFold2.resolve_ablation_defaults(5)["evoformer_enabled"] is False

tests/test_forward_model.py

@@ -4,6 +4,8 @@
 
 import torch
 
+from model.alphafold2 import AlphaFold2
+
 
 def test_alphafold2_forward_smoke(toy_model, toy_batch):
     with torch.no_grad():
@@ -28,6 +30,8 @@ def test_alphafold2_forward_smoke(toy_model, toy_batch):
     assert outputs["plddt"].shape == (batch_size, length)
     assert outputs["distogram_logits"].shape == (batch_size, length, length, 64)
     assert outputs["masked_msa_logits"].shape == (batch_size, toy_batch["msa_tokens"].shape[1], length, 23)
+    assert outputs["tm_logits"] is None
+    assert outputs["ptm"] is None
 
     for value in outputs.values():
         if torch.is_tensor(value):
@@ -41,6 +45,45 @@ def test_alphafold2_forward_smoke(toy_model, toy_batch):
     assert torch.all((outputs["plddt"] >= 0.0) & (outputs["plddt"] <= 100.0))
 
 
+def test_alphafold2_tm_head_can_be_enabled(toy_batch):
+    torch.manual_seed(11)
+    model = AlphaFold2(
+        n_tokens=27,
+        c_m=256,
+        c_z=128,
+        c_s=256,
+        max_relpos=32,
+        pad_idx=0,
+        num_evoformer_blocks=1,
+        num_structure_blocks=1,
+        transition_expansion_evoformer=2,
+        transition_expansion_structure=2,
+        use_block_specific_params=False,
+        dist_bins=64,
+        plddt_bins=50,
+        tm_num_bins=64,
+        tm_head_enabled=True,
+        n_torsions=3,
+        num_res_blocks_torsion=1,
+    ).eval()
+
+    with torch.no_grad():
+        outputs = model(
+            seq_tokens=toy_batch["seq_tokens"],
+            msa_tokens=toy_batch["msa_tokens"],
+            seq_mask=toy_batch["seq_mask"],
+            msa_mask=toy_batch["msa_mask"],
+            ideal_backbone_local=toy_batch["ideal_backbone_local"],
+        )
+
+    batch_size, length = toy_batch["seq_tokens"].shape
+    assert outputs["tm_logits"].shape == (batch_size, length, length, 64)
+    assert outputs["ptm"].shape == (batch_size,)
+    assert torch.isfinite(outputs["tm_logits"]).all()
+    assert torch.isfinite(outputs["ptm"]).all()
+    assert torch.all((outputs["ptm"] >= 0.0) & (outputs["ptm"] <= 1.0))
+
+
 def test_alphafold_loss_orchestrator_returns_finite_components(toy_model, toy_batch, toy_criterion):
     with torch.no_grad():
         outputs = toy_model(