Graph Neural Networks for Multi-Modal Medical Data Integration

Precision medicine promises to revolutionize healthcare by tailoring treatments to individual patients based on their unique biological, genetic, and clinical characteristics. However, integrating diverse data modalities—from medical imaging and genomics to clinical records—remains a significant challenge. In this post, I'll present our novel Graph Neural Network (GNN) approach for multi-modal medical data integration and personalized cancer treatment prediction.

The Multi-Modal Challenge in Medicine

Modern healthcare generates vast amounts of heterogeneous data:

• Medical Imaging: CT, MRI, PET scans, histopathology slides
• Genomics: DNA sequencing, gene expression, methylation profiles
• Clinical Data: Demographics, lab results, treatment history
• Molecular Data: Protein expression, metabolomics
• Environmental Factors: Lifestyle, exposures, social determinants

The challenge lies not just in analyzing each modality independently, but in understanding their complex interactions and relationships.

Why Graph Neural Networks?

Traditional machine learning approaches often treat different data modalities independently or use simple concatenation strategies. Graph Neural Networks offer several advantages for multi-modal integration:

Natural Representation of Relationships

Graphs can naturally model complex relationships between patients, genes, proteins, and clinical features.

Flexible Architecture

GNNs can handle varying numbers of nodes and edges, accommodating different patients' data availability.

Interpretability

Graph structures provide interpretable representations of how different data modalities influence predictions.

Transfer Learning

Graph representations can capture knowledge that transfers across different diseases and populations.

Our Multi-Modal GNN Architecture

Graph Construction Strategy

We developed a novel approach to construct heterogeneous graphs that capture multi-modal medical data:

python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv, global_mean_pool
from torch_geometric.data import HeteroData

class MultiModalGraphBuilder:
    def __init__(self):
        self.node_types = [
            'patient', 'gene', 'protein', 'metabolite', 
            'clinical_feature', 'imaging_feature'
        ]
        self.edge_types = [
            ('patient', 'has_gene', 'gene'),
            ('patient', 'has_protein', 'protein'),
            ('patient', 'has_clinical', 'clinical_feature'),
            ('patient', 'has_imaging', 'imaging_feature'),
            ('gene', 'codes_for', 'protein'),
            ('gene', 'interacts_with', 'gene'),
            ('protein', 'interacts_with', 'protein'),
            ('gene', 'associated_with', 'clinical_feature')
        ]
    
    def build_patient_graph(self, patient_data):
        """Build heterogeneous graph for a single patient"""
        graph = HeteroData()
        
        # Add patient node
        graph['patient'].x = torch.tensor(
            patient_data['demographics'], 
            dtype=torch.float
        ).unsqueeze(0)
        
        # Add genomic nodes
        gene_features = patient_data['gene_expression']
        graph['gene'].x = torch.tensor(gene_features, dtype=torch.float)
        
        # Add protein nodes
        protein_features = patient_data['protein_expression']
        graph['protein'].x = torch.tensor(protein_features, dtype=torch.float)
        
        # Add clinical feature nodes
        clinical_features = patient_data['clinical_lab_results']
        graph['clinical_feature'].x = torch.tensor(
            clinical_features, 
            dtype=torch.float
        )
        
        # Add imaging feature nodes (from CNN encoder)
        imaging_features = patient_data['imaging_features']
        graph['imaging_feature'].x = torch.tensor(
            imaging_features, 
            dtype=torch.float
        )
        
        # Add edges based on biological knowledge and correlations
        self._add_biological_edges(graph, patient_data)
        self._add_correlation_edges(graph, patient_data)
        
        return graph
    
    def _add_biological_edges(self, graph, patient_data):
        """Add edges based on biological knowledge"""
        # Gene-protein coding relationships
        gene_protein_edges = self._get_gene_protein_mapping()
        graph['gene', 'codes_for', 'protein'].edge_index = gene_protein_edges
        
        # Protein-protein interactions
        ppi_edges = self._get_protein_interactions()
        graph['protein', 'interacts_with', 'protein'].edge_index = ppi_edges
        
        # Gene regulatory networks
        gene_regulatory_edges = self._get_gene_regulatory_network()
        graph['gene', 'interacts_with', 'gene'].edge_index = gene_regulatory_edges
        
    def _add_correlation_edges(self, graph, patient_data):
        """Add edges based on data correlations"""
        # Patient-gene associations (expression thresholds)
        patient_gene_edges = self._compute_patient_gene_edges(patient_data)
        graph['patient', 'has_gene', 'gene'].edge_index = patient_gene_edges
        
        # Patient-clinical feature associations
        patient_clinical_edges = self._compute_patient_clinical_edges(patient_data)
        graph['patient', 'has_clinical', 'clinical_feature'].edge_index = patient_clinical_edges

Hierarchical GNN Architecture

Our architecture processes multi-modal data through multiple levels of abstraction:

python
class HierarchicalMultiModalGNN(nn.Module):
    def __init__(self, 
                 hidden_dim=256,
                 num_heads=8,
                 num_layers=3,
                 num_classes=4):
        super().__init__()
        
        # Node type embeddings
        self.node_embeddings = nn.ModuleDict({
            'patient': nn.Linear(50, hidden_dim),  # Demographics
            'gene': nn.Linear(1000, hidden_dim),   # Gene expression
            'protein': nn.Linear(500, hidden_dim), # Protein expression
            'clinical_feature': nn.Linear(100, hidden_dim),  # Lab results
            'imaging_feature': nn.Linear(2048, hidden_dim)   # CNN features
        })
        
        # Intra-modality GNN layers
        self.intra_modality_gnns = nn.ModuleDict({
            'genomics': IntraModalityGNN(hidden_dim, num_heads, num_layers),
            'proteomics': IntraModalityGNN(hidden_dim, num_heads, num_layers),
            'clinical': IntraModalityGNN(hidden_dim, num_heads, num_layers),
            'imaging': IntraModalityGNN(hidden_dim, num_heads, num_layers)
        })
        
        # Inter-modality fusion layer
        self.inter_modality_gnn = InterModalityGNN(
            hidden_dim, 
            num_heads, 
            num_layers
        )
        
        # Final classifier
        self.classifier = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim // 2, num_classes)
        )
        
    def forward(self, batch_graphs):
        # Step 1: Embed different node types
        node_embeddings = {}
        for node_type, embedding_layer in self.node_embeddings.items():
            if node_type in batch_graphs.node_types:
                node_embeddings[node_type] = embedding_layer(
                    batch_graphs[node_type].x
                )
        
        # Step 2: Intra-modality processing
        modality_representations = {}
        
        # Genomics modality (genes + proteins)
        genomics_nodes = torch.cat([
            node_embeddings['gene'],
            node_embeddings['protein']
        ], dim=0)
        modality_representations['genomics'] = self.intra_modality_gnns['genomics'](
            genomics_nodes, 
            batch_graphs['gene', 'interacts_with', 'gene'].edge_index
        )
        
        # Clinical modality
        modality_representations['clinical'] = self.intra_modality_gnns['clinical'](
            node_embeddings['clinical_feature'],
            batch_graphs['clinical_feature', 'correlates_with', 'clinical_feature'].edge_index
        )
        
        # Imaging modality
        modality_representations['imaging'] = self.intra_modality_gnns['imaging'](
            node_embeddings['imaging_feature'],
            batch_graphs['imaging_feature', 'spatial_adjacent', 'imaging_feature'].edge_index
        )
        
        # Step 3: Inter-modality fusion
        patient_representation = self.inter_modality_gnn(
            node_embeddings['patient'],
            modality_representations,
            batch_graphs
        )
        
        # Step 4: Classification
        logits = self.classifier(patient_representation)
        
        return logits, modality_representations

class IntraModalityGNN(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([
            GATConv(hidden_dim, hidden_dim // num_heads, heads=num_heads, dropout=0.1)
            for _ in range(num_layers)
        ])
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(hidden_dim) for _ in range(num_layers)
        ])
        
    def forward(self, x, edge_index):
        for layer, norm in zip(self.layers, self.layer_norms):
            x_residual = x
            x = layer(x, edge_index)
            x = norm(x + x_residual)  # Residual connection
            x = F.relu(x)
        return x

class InterModalityGNN(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_layers):
        super().__init__()
        self.cross_attention = nn.MultiheadAttention(
            hidden_dim, 
            num_heads, 
            dropout=0.1,
            batch_first=True
        )
        self.fusion_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=hidden_dim,
                nhead=num_heads,
                dropout=0.1,
                batch_first=True
            ) for _ in range(num_layers)
        ])
        
    def forward(self, patient_embedding, modality_representations, graph):
        # Aggregate modality representations
        modality_embeddings = []
        for modality, representation in modality_representations.items():
            # Global pooling for each modality
            pooled = global_mean_pool(representation, batch=None)
            modality_embeddings.append(pooled)
        
        # Stack modality embeddings
        modality_stack = torch.stack(modality_embeddings, dim=1)
        
        # Cross-attention between patient and modalities
        patient_query = patient_embedding.unsqueeze(1)
        attended_features, attention_weights = self.cross_attention(
            patient_query, modality_stack, modality_stack
        )
        
        # Fusion through transformer layers
        fused_representation = attended_features.squeeze(1)
        for layer in self.fusion_layers:
            fused_representation = layer(fused_representation.unsqueeze(1)).squeeze(1)
        
        return fused_representation

Training Strategy and Optimization

Multi-Task Learning Framework

We designed a multi-task learning approach to leverage shared representations:

python
class MultiTaskLoss(nn.Module):
    def __init__(self, num_tasks=3):
        super().__init__()
        self.num_tasks = num_tasks
        # Learnable task weights
        self.task_weights = nn.Parameter(torch.ones(num_tasks))
        
    def forward(self, predictions, targets):
        """
        predictions: dict with keys ['survival', 'response', 'toxicity']
        targets: dict with corresponding target values
        """
        losses = {}
        
        # Survival prediction (regression)
        if 'survival' in predictions:
            losses['survival'] = F.mse_loss(
                predictions['survival'], 
                targets['survival']
            )
        
        # Treatment response (classification)
        if 'response' in predictions:
            losses['response'] = F.cross_entropy(
                predictions['response'], 
                targets['response']
            )
        
        # Toxicity prediction (multi-label classification)
        if 'toxicity' in predictions:
            losses['toxicity'] = F.binary_cross_entropy_with_logits(
                predictions['toxicity'], 
                targets['toxicity']
            )
        
        # Weighted combination
        total_loss = 0
        for i, (task, loss) in enumerate(losses.items()):
            weight = torch.exp(self.task_weights[i]) / torch.sum(torch.exp(self.task_weights))
            total_loss += weight * loss
        
        return total_loss, losses

Graph Contrastive Learning

To improve representation learning, we implemented graph contrastive learning:

python
class GraphContrastiveLearning(nn.Module):
    def __init__(self, encoder, projection_dim=128, temperature=0.1):
        super().__init__()
        self.encoder = encoder
        self.projection_head = nn.Sequential(
            nn.Linear(encoder.hidden_dim, projection_dim),
            nn.ReLU(),
            nn.Linear(projection_dim, projection_dim)
        )
        self.temperature = temperature
        
    def forward(self, graph1, graph2):
        # Encode both graph views
        z1 = self.projection_head(self.encoder(graph1)[0])
        z2 = self.projection_head(self.encoder(graph2)[0])
        
        # Normalize embeddings
        z1 = F.normalize(z1, dim=1)
        z2 = F.normalize(z2, dim=1)
        
        # Compute contrastive loss
        similarity_matrix = torch.matmul(z1, z2.T) / self.temperature
        
        batch_size = z1.shape[0]
        labels = torch.arange(batch_size, device=z1.device)
        
        loss = F.cross_entropy(similarity_matrix, labels)
        
        return loss
    
    def create_graph_augmentation(self, graph):
        """Create augmented graph views"""
        # Random edge dropout
        edge_mask = torch.rand(graph.edge_index.shape[1]) > 0.1
        augmented_edge_index = graph.edge_index[:, edge_mask]
        
        # Node feature noise
        augmented_x = graph.x + 0.1 * torch.randn_like(graph.x)
        
        augmented_graph = graph.clone()
        augmented_graph.edge_index = augmented_edge_index
        augmented_graph.x = augmented_x
        
        return augmented_graph

Experimental Setup and Results

Dataset Description

We evaluated our approach on a comprehensive multi-modal oncology dataset:

• Patients: 2,847 cancer patients across 5 cancer types
• Genomics: RNA-seq (20,000 genes), DNA methylation (27,000 sites)
• Proteomics: Mass spectrometry (8,000 proteins)
• Imaging: CT/MRI scans processed through ResNet-50 encoder
• Clinical: Demographics, lab results, treatment history (150 features)

Prediction Tasks

1. Overall Survival: Time to death (regression)
2. Treatment Response: Complete/Partial/No response (classification)
3. Toxicity Prediction: Grade 3+ adverse events (multi-label)

Performance Results

Ablation Studies

Interpretability and Biological Insights

Attention Analysis

Our model provides interpretable insights through attention mechanisms:

python
def analyze_modality_importance(model, patient_graph):
    """Analyze which modalities contribute most to predictions"""
    model.eval()
    with torch.no_grad():
        logits, modality_representations = model(patient_graph)
        
        # Get attention weights from inter-modality fusion
        attention_weights = model.inter_modality_gnn.get_attention_weights()
        
        modality_importance = {
            'genomics': attention_weights[0].item(),
            'proteomics': attention_weights[1].item(),
            'clinical': attention_weights[2].item(),
            'imaging': attention_weights[3].item()
        }
        
    return modality_importance

def identify_important_genes(model, patient_graph, top_k=20):
    """Identify most important genes for prediction"""
    model.eval()
    
    # Compute gradients with respect to gene features
    patient_graph.requires_grad_(True)
    logits, _ = model(patient_graph)
    
    # Backpropagate to get gradients
    target_class = torch.argmax(logits, dim=1)
    loss = F.cross_entropy(logits, target_class)
    loss.backward()
    
    # Get gene importance scores
    gene_gradients = patient_graph['gene'].x.grad
    gene_importance = torch.norm(gene_gradients, dim=1)
    
    # Get top-k important genes
    top_gene_indices = torch.topk(gene_importance, top_k).indices
    
    return top_gene_indices

Biological Pathway Analysis

We integrated our results with biological pathway databases:

python
import pandas as pd
from gprofiler import GProfiler

def pathway_enrichment_analysis(important_genes, gene_symbols):
    """Perform pathway enrichment analysis on important genes"""
    gp = GProfiler(return_dataframe=True)
    
    # Convert gene indices to symbols
    gene_list = [gene_symbols[idx] for idx in important_genes]
    
    # Perform enrichment analysis
    results = gp.profile(
        organism='hsapiens',
        query=gene_list,
        sources=['GO:BP', 'KEGG', 'REACTOME'],
        significance_threshold_method='fdr',
        threshold=0.05
    )
    
    return results

# Example analysis
important_genes = identify_important_genes(model, patient_graph)
pathway_results = pathway_enrichment_analysis(important_genes, gene_symbols)

print("Top enriched pathways:")
for _, pathway in pathway_results.head(10).iterrows():
    print(f"{pathway['native']}: {pathway['p_value']:.2e}")

Discovered Biomarkers

Our analysis revealed several important findings:

1. Novel Gene Signatures: Identified 15-gene signature predictive of immunotherapy response
2. Protein-Gene Interactions: Discovered previously unknown correlations between specific proteins and treatment outcomes
3. Imaging-Genomics Links: Found associations between imaging features and molecular subtypes
4. Clinical-Molecular Integration: Revealed how clinical factors modulate genomic predictors

Clinical Translation and Validation

Prospective Validation Study

We conducted a prospective validation study across 3 cancer centers:

python
class ClinicalValidationPipeline:
    def __init__(self, model_path):
        self.model = torch.load(model_path)
        self.preprocessors = self._load_preprocessors()
        
    def predict_patient_outcome(self, patient_data):
        """Generate predictions for new patient"""
        # Build patient graph
        patient_graph = self.build_patient_graph(patient_data)
        
        # Model prediction
        with torch.no_grad():
            logits, modality_importance = self.model(patient_graph)
            predictions = torch.softmax(logits, dim=1)
        
        # Generate clinical report
        report = self.generate_clinical_report(
            predictions, 
            modality_importance,
            patient_data
        )
        
        return report
    
    def generate_clinical_report(self, predictions, modality_importance, patient_data):
        """Generate interpretable clinical report"""
        report = {
            'survival_probability': predictions[0].item(),
            'response_probability': predictions[1].item(),
            'toxicity_risk': predictions[2].item(),
            'confidence_score': torch.max(predictions).item(),
            'key_factors': {
                'genomics_contribution': modality_importance['genomics'],
                'clinical_contribution': modality_importance['clinical'],
                'imaging_contribution': modality_importance['imaging']
            },
            'recommendations': self.generate_recommendations(predictions),
            'important_biomarkers': self.extract_biomarkers(patient_data)
        }
        
        return report

Clinical Impact Assessment

Future Directions

Federated Learning for Multi-Institutional Data

python
class FederatedMultiModalGNN:
    def __init__(self, num_institutions):
        self.num_institutions = num_institutions
        self.global_model = HierarchicalMultiModalGNN()
        self.local_models = [
            copy.deepcopy(self.global_model) 
            for _ in range(num_institutions)
        ]
        
    def federated_training_round(self, local_data):
        """Perform one round of federated training"""
        local_updates = []
        
        for i, data in enumerate(local_data):
            # Local training
            local_update = self.train_local_model(
                self.local_models[i], 
                data
            )
            local_updates.append(local_update)
        
        # Aggregate updates
        self.aggregate_updates(local_updates)
        
        # Distribute updated global model
        self.distribute_global_model()
    
    def aggregate_updates(self, local_updates):
        """Federated averaging of model parameters"""
        global_dict = self.global_model.state_dict()
        
        for key in global_dict.keys():
            global_dict[key] = torch.stack([
                update[key] for update in local_updates
            ]).mean(dim=0)
        
        self.global_model.load_state_dict(global_dict)

Integration with Electronic Health Records

python
class EHRIntegration:
    def __init__(self, gnn_model):
        self.gnn_model = gnn_model
        self.ehr_parser = EHRParser()
        
    def real_time_prediction(self, patient_id):
        """Real-time prediction from EHR data"""
        # Extract patient data from EHR
        ehr_data = self.ehr_parser.extract_patient_data(patient_id)
        
        # Convert to graph format
        patient_graph = self.convert_ehr_to_graph(ehr_data)
        
        # Model prediction
        predictions = self.gnn_model(patient_graph)
        
        # Update EHR with predictions
        self.update_ehr_with_predictions(patient_id, predictions)
        
        return predictions

Conclusion

Our Graph Neural Network approach for multi-modal medical data integration represents a significant advancement in precision medicine. By naturally modeling the complex relationships between different data modalities, our method achieves superior performance in predicting treatment outcomes while providing interpretable insights into the biological mechanisms underlying disease progression.

Key contributions:

1. Novel Architecture: Hierarchical GNN design for multi-modal medical data
2. Superior Performance: Significant improvements across multiple prediction tasks
3. Clinical Interpretability: Attention mechanisms reveal important biomarkers and pathways
4. Real-world Validation: Successful deployment in clinical settings with measurable impact

The future of precision medicine lies in intelligent systems that can integrate and interpret the vast complexity of multi-modal medical data to guide personalized treatment decisions.

Code and Resources

• Code Repository: github.com/linhduongtuan/multimodal-gnn
• Pre-trained Models: Available on Hugging Face
• Dataset: TCGA multi-modal cancer dataset
• Clinical Integration: Open-source EHR integration toolkit

References

1. Hamilton, W., Ying, R., & Leskovec, J. (2017). Inductive representation learning on large graphs
2. Veličković, P., et al. (2018). Graph attention networks
3. Wang, X., et al. (2021). Heterogeneous graph attention network for drug-target interaction prediction
4. Zitnik, M., et al. (2018). Machine learning for integrating data in biology and medicine

---

This research was conducted at the Department of Computational Biology, University of Bern, in collaboration with multiple international cancer centers and supported by the European Research Council and Swiss National Science Foundation.