Building a Brain CT Classifier with Multiple Instance Learning

Overview

For my undergraduate thesis, I built a multi-label classification system for brain CT scans. The challenge: 859 patient scans, multiple possible pathologies per scan, and no slice-level annotations — only patient-level labels.

This post walks through the key design decisions: encoder choice, Multiple Instance Learning (MIL), the dual-head architecture, and how I handled severe class imbalance.

Why ConvNeXt-V2?

ConvNeXt-V2 Tiny offered an excellent accuracy-to-parameter trade-off for our constrained dataset. Unlike Vision Transformers, it doesn't require massive pretraining data to generalize well on fine-grained medical features.

Multiple Instance Learning

The core challenge: we have labels at the patient level ("this patient has an epidural hemorrhage") but not at the slice level ("slice 34 contains the hemorrhage"). MIL treats each patient as a bag of slices (instances), and learns to make bag-level predictions by aggregating instance-level features.

class MILAttention(nn.Module):
    def __init__(self, feature_dim: int):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Linear(feature_dim, 128),
            nn.Tanh(),
            nn.Linear(128, 1),
        )

    def forward(self, features):
        # features: [N_slices, feature_dim]
        scores = self.attention(features)          # [N, 1]
        weights = torch.softmax(scores, dim=0)     # [N, 1]
        aggregated = (weights * features).sum(0)   # [feature_dim]
        return aggregated, weights

The attention weights tell us which slices the model found important — a built-in form of explainability.

Dual-Head Architecture

Rather than flattening everything into one multi-label output, I used two specialized heads:

• Pathology head: predicts hemorrhage type (epidural, subdural, subarachnoid, intraparenchymal, intraventricular)
• Location head: predicts anatomical region (frontal, temporal, parietal, occipital, etc.)

This factored design improved performance on both tasks compared to a single unified head.

Handling Class Imbalance

With real clinical data, some pathologies appear in <5% of cases. Standard BCE loss fails catastrophically here. I used Asymmetric Loss (Ridnik et al., 2021):

class AsymmetricLoss(nn.Module):
    def __init__(self, gamma_neg=4, gamma_pos=0, clip=0.05):
        super().__init__()
        self.gamma_neg = gamma_neg
        self.gamma_pos = gamma_pos
        self.clip = clip

    def forward(self, logits, targets):
        # Asymmetrically down-weights easy negatives
        ...

This reduced false negatives for rare pathologies by ~23%.

Results & Lessons

• The MIL attention maps were clinically interpretable — matching radiologist annotations ~68% of the time
• Dual-head design outperformed single-head by 4.2% macro-AUC
• The biggest bottleneck was data quality, not model architecture

Working with real clinical data taught me more about the practical constraints of medical AI than any benchmark dataset could.