Training Pipeline

The ProMaya training loop optimizes the entire 5-module deep learning framework end-to-end (excluding the XGBoost head, which is fit post-hoc).

The Hybrid Loss Function

Because ProMaya fuses modalities from vastly different scales (Atoms to Language Models), a simple Binary Cross-Entropy (BCE) loss is insufficient. The network is trained using a complex, multi-objective loss function:

\[ L_{total} = L_{BCE} + 0.1 \cdot L_{contrastive} + 0.05 \cdot L_{LSMD} + 0.05 \cdot L_{IDR} + 0.2 \cdot L_{focal} \]

Breakdown:

\(L_{BCE}\) (Primary): Standard binary cross-entropy on the interaction probability.
\(L_{focal}\) (Mining): Focal loss (\(\gamma=2\)) to heavily penalize the network for getting the "Extreme Difficulty" docking-derived negatives wrong.
\(L_{contrastive}\) (Alignment): InfoNCE loss (\(\tau=0.07\)) applied during Module 2 to ensure the Atom, Residue, and Sequence latent spaces align.
\(L_{LSMD}\) (Physics): \(L_1\) loss enforcing that the abstract embeddings preserve the physical Local Surface Mass Density distribution.
\(L_{IDR}\) (Disorder): BCE loss ensuring the network maintains an accurate representation of Intrinsic Disorder Regions.

Optimization Strategy

Optimizer: AdamW (Learning Rate \(= 1\times10^{-4}\), Weight Decay \(= 1\times10^{-2}\))
Scheduler: Cosine Annealing (\(T_{max} = 100\), \(\eta_{min} = 1\times10^{-6}\))
Batch Size: 32
Gradient Clipping: Norm 1.0 (Crucial to prevent explosions from the Heterogeneous Graph Transformer)
Mixed Precision: FP16 (bfloat16) training enabled to drastically reduce VRAM usage.

src/training/loop.py

def train_epoch(model, dataloader, criterion, optimizer, device, epoch):
    model.train()
    total_loss = 0.0

    for batch_idx, (protein_a, protein_b, labels) in enumerate(dataloader):
        protein_a = move_to_device(protein_a, device)
        protein_b = move_to_device(protein_b, device)
        labels = labels.to(device)

        # Forward pass (end-to-end through Modules 1-5)
        predictions, outputs = model(protein_a, protein_b)

        # Calculate hybrid loss
        loss, loss_dict = criterion(predictions, labels, outputs)

        optimizer.zero_grad()
        loss.backward()

        # Gradient clipping prevents HGT explosions
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

Computational Scale

Training the full ProMaya model is computationally demanding due to the explicit atomic graphs and 3D point clouds.

~32M Trainable Params

24-48 Hours to Train

100 Total Epochs

(Note: The 650M parameter ProtTrans model is frozen and pre-cached, so its parameters are not counted in the trainable total).

Next: Review the performance in Evaluation Metrics â†’