Module 2: Cross-Level Multimodal Alignment

After extracting the four distinct data modalities (Atom, Residue, Surface, Sequence), ProMaya must synthesize them. Operating on them independently would miss critical biophysical realitiesâ€”for example, an evolutionarily conserved residue (Sequence) is only a binding hotspot if it is exposed on the geometry (Surface) with the right charge density (Atom).

Module 2 utilizes a 4-Stage Bidirectional Cross-Attention mechanism to achieve this fusion.

graph TD
    A[Atomic Features] -->|Stage 1: Atom â†’ Residue| R[Residue Features]
    R -->|Stage 2: Residue â†’ Surface| S[Surface Point Cloud]
    S -->|Stage 3: Surface â†’ Sequence| Seq[ProtTrans Sequence]
    Seq -->|Stage 4: Sequence â†’ Residue| R

    style A fill:#e1f5fe,stroke:#0b4b7c,stroke-width:2px,color:#0b4b7c
    style R fill:#b2ebf2,stroke:#008080,stroke-width:2px,color:#0b4b7c
    style S fill:#e1f5fe,stroke:#00bcd4,stroke-width:2px,color:#0b4b7c
    style Seq fill:#f0f8ff,stroke:#1e3a5f,stroke-width:2px,color:#0b4b7c

The 4 Stages

Each stage utilizes 2 blocks of 8-head cross-attention.

Stage 1: Atom â†’ Residue

Mechanism: Fuses atomic-level electron density (LSMD) and precise chemical geometry up into the residue-level embeddings.
Biological Purpose: Ensures that the residue's "understanding" of its environment is grounded in actual physical mass density, rather than just abstract amino acid type.

Stage 2: Residue â†’ Surface

Mechanism: Projects the aggregated biochemical properties (charge, hydrophobicity, secondary structure) onto the triangulated 3D surface points.
Biological Purpose: Gives the "dumb" geometric shapes (from PointNet++) explicit chemical meaning, creating a true physiochemical binding landscape.

Stage 3: Surface â†’ Sequence

Mechanism: Takes the 3D spatial geometry and maps it to the 1D sequence embeddings (ProtTrans).
Biological Purpose: Grounds the abstract, language-model learned evolutionary signals in physical 3D reality.

Stage 4: Sequence â†’ Residue

Mechanism: Feeds the fully enriched, geometrically-aware evolutionary context back to the primary residue graph.
Biological Purpose: Completes the loop, ensuring the final residue graph nodes possess both micro-physical (Atomic) and macro-contextual (Sequence) intelligence.

src/models/alignment.py

import torch
import torch.nn as nn

class MultimodalAlignment(nn.Module):
    def __init__(self, d_model: int = 256, nhead: int = 8):
        super().__init__()
        # Cross-attention blocks for the 4 stages
        self.atom_to_res = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.res_to_surf = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.surf_to_seq = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.seq_to_res  = nn.MultiheadAttention(d_model, nhead, batch_first=True)

    def forward(self, atom, res, surf, seq):
        # Stage 1: Atom -> Residue
        # Query: Residue, Key/Value: Atom
        res_aligned, _ = self.atom_to_res(query=res, key=atom, value=atom)
        res = res + res_aligned

        # Stage 2: Residue -> Surface
        surf_aligned, _ = self.res_to_surf(query=surf, key=res, value=res)
        surf = surf + surf_aligned

        # Stage 3: Surface -> Sequence
        seq_aligned, _ = self.surf_to_seq(query=seq, key=surf, value=surf)
        seq = seq + seq_aligned

        # Stage 4: Sequence -> Residue
        res_final, _ = self.seq_to_res(query=res, key=seq, value=seq)
        res = res + res_final

        return atom, res, surf, seq

[TIP] Auxiliary Alignment Losses To ensure the cross-attention layers learn physically meaningful mappings, ProMaya applies three auxiliary loss functions during training: 1. LSMD Alignment (\(L_1\)): Forces the network to preserve physical mass density mapping. 2. InfoNCE Contrastive: Aligns modalities in the latent space. 3. IDR Consistency (BCE): Ensures intrinsic disorder regions are mathematically preserved across structural and sequence representations.

Next: Learn how this aligned data is graphed in the HGT Encoder â†’