Module 4: Cross-Protein Interaction

Modules 1-3 process Protein A (\(P_A\)) and Protein B (\(P_B\)) completely independently. Module 4 is where the actual mechanics of binding are simulated.

Instead of simply concatenating the embeddings of \(P_A\) and \(P_B\), ProMaya employs Multi-Scale Cross-Attention to dynamically map the complementary surfaces, sequences, and atoms between the two proteins.

Multi-Scale Attention

The framework computes attention across three distinct spatial scales:

1. Residue-Residue

8 Heads, 2 Blocks
Models co-evolutionary patterns and complementary amino acid motifs (e.g., a positively charged pocket on P_A interacting with a negatively charged loop on P_B).

2. Surface-Surface

4 Heads, 2 Blocks
Models strict geometric shape complementarity (the "lock and key" mechanism), ensuring steric clashes are avoided.

3. Sparse Atom-Atom

4 Heads, LSMD-Filtered
Due to computational constraints, \(O(N^2)\) atomic attention is impossible. ProMaya filters atoms by their LSMD score (only keeping atoms \(> \mu + 0.5\sigma\)). This performs precise chemical compatibility checks only on the densest, most likely binding interfaces.

Pair Embedding Formulation

After cross-attention, the updated embeddings for each protein (\(e'_A\) and \(e'_B\)) are pooled. To capture both the magnitude and the direction of the interaction, the final pair embedding is constructed using a standard relation concatenation strategy:

\[ e_{pair} = [ e'_A \ || \ e'_B \ || \ |e'_A - e'_B| \ || \ (e'_A \odot e'_B) \ || \ e_{globalA} \ || \ e_{globalB} ] \]

This massive 2048-dimensional vector encapsulates every scale of the predicted protein-protein interaction and is passed to the final classification module.

src/models/interaction.py

import torch
import torch.nn as nn

class MultiScaleInteraction(nn.Module):
    def __init__(self, d_model=256, nhead_res=8, nhead_surf=4, nhead_atom=4):
        super().__init__()
        self.res_attn = nn.MultiheadAttention(d_model, nhead_res, batch_first=True)
        self.surf_attn = nn.MultiheadAttention(d_model, nhead_surf, batch_first=True)
        self.atom_attn = nn.MultiheadAttention(d_model, nhead_atom, batch_first=True)

    def forward(self, nodes_A, nodes_B):
        # 1. Residue-Residue Interaction
        res_A, _ = self.res_attn(query=nodes_A['residue'], key=nodes_B['residue'], value=nodes_B['residue'])
        res_B, _ = self.res_attn(query=nodes_B['residue'], key=nodes_A['residue'], value=nodes_A['residue'])

        # 2. Surface-Surface Interaction
        surf_A, _ = self.surf_attn(query=nodes_A['surface'], key=nodes_B['surface'], value=nodes_B['surface'])
        surf_B, _ = self.surf_attn(query=nodes_B['surface'], key=nodes_A['surface'], value=nodes_A['surface'])

        # 3. Sparse Atom-Atom Interaction (Assume nodes_A/B['atom'] is already LSMD filtered)
        atom_A, _ = self.atom_attn(query=nodes_A['atom'], key=nodes_B['atom'], value=nodes_B['atom'])
        atom_B, _ = self.atom_attn(query=nodes_B['atom'], key=nodes_A['atom'], value=nodes_A['atom'])

        # Mean pooling across elements
        e_A = torch.cat([res_A.mean(dim=1), surf_A.mean(dim=1), atom_A.mean(dim=1)], dim=-1) # 768-d
        e_B = torch.cat([res_B.mean(dim=1), surf_B.mean(dim=1), atom_B.mean(dim=1)], dim=-1) # 768-d

        # Relation Concatenation
        e_pair = torch.cat([e_A, e_B, torch.abs(e_A - e_B), e_A * e_B], dim=-1)
        return e_pair

Next: See how this vector is scored in Hybrid Classification â†’