Module 5: Hybrid Classification

The final stage of the ProMaya framework transforms the dense 2048-dimensional pair embedding from Module 4 into a single, calibrated interaction probability \([0, 1]\).

To maximize generalization and prevent overfitting on the complex graph embeddings, ProMaya utilizes a Hybrid Neural-XGBoost classification layer.

1. Neural Feed-Forward Network (FFN)

The pair embedding first passes through a 3-layer dense neural network to distill the interaction signature:

Linear Layer 1: \(2048 \rightarrow 512\)
Activation: BatchNorm + GELU
Regularization: Dropout (0.3)
Linear Layer 2: \(512 \rightarrow 128\)
Activation: BatchNorm + GELU
Regularization: Dropout (0.2)
Output: The 128-dimensional Interaction Signature
Linear Layer 3: \(128 \rightarrow 1\) (Used during initial backpropagation)

2. XGBoost Ensemble

While the FFN is trained end-to-end via gradient descent, the final production prediction utilizes an XGBoost (Extreme Gradient Boosting) classifier operating on the frozen 128-d Interaction Signature.

Tree-based ensembles are remarkably robust against outliers and scale disparities, providing a highly calibrated final probability.

Ensemble Size: 500 decision trees
Max Depth: 6
Calibration: Isotonic regression is applied to ensure the output probability \([0, 1]\) directly corresponds to the statistical likelihood of interaction.

src/models/classification.py

import torch
import torch.nn as nn

class HybridClassifier(nn.Module):
    def __init__(self, in_features=2048):
        super().__init__()
        self.ffn = nn.Sequential(
            nn.Linear(in_features, 512),
            nn.BatchNorm1d(512),
            nn.GELU(),
            nn.Dropout(0.3),

            nn.Linear(512, 128),
            nn.BatchNorm1d(128),
            nn.GELU(),
            nn.Dropout(0.2)
        )
        # End-to-end training logit
        self.output_layer = nn.Linear(128, 1)
        # Pre-trained XGBoost model for production
        self.xgb_model = None 

    def forward(self, e_pair):
        # Extract 128-d Interaction Signature
        signature = self.ffn(e_pair)

        # Neural prediction (used for backprop)
        nn_logits = self.output_layer(signature)

        if not self.training and self.xgb_model is not None:
            # In production inference, override with XGBoost probability
            xgb_prob = self.xgb_model.predict_proba(signature.cpu().numpy())[:, 1]
            # Convert back to logit scale if needed
            nn_logits = torch.logit(torch.tensor(xgb_prob, device=e_pair.device)).unsqueeze(1)

        return nn_logits, signature

[TIP] Why the Hybrid Approach? Deep neural networks are excellent at feature extraction (Modules 1-4) but can suffer from poor probability calibration (being overly confident). XGBoost excels at tabular, structured feature classification. Combining them gives ProMaya state-of-the-art discrimination and calibration.

Next: Explore how the model is trained in Dataset Construction â†’