HP2L Taught an AI to Think Like a Radiologist — Step by Step, Disease by Disease

The following is a complete PyTorch implementation of HP2L as described in Sections 3.1 through 3.4 of the paper. It covers the full three-level hierarchical ViT backbone with PTB and VTB sub-blocks, the cross-attention Prompt Learning module, the EMA-based Prototype Learning module, the level-wise classification heads, the combined training loss with hierarchy consistency penalty, and a runnable smoke test on synthetic 3D brain MRI data.

# =============================================================================
# HP2L: Hierarchical Prompt and Prototype Learning for Brain Disorder Diagnosis
# Paper: "A hierarchical prompt and prototype learning framework for brain
#         disorder classification"
# Authors: Yuxiao Liu et al., ShanghaiTech / Henan Provincial People's Hospital
# Journal: Medical Image Analysis 112 (2026) 104063
# =============================================================================

from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Dict, Optional, Tuple


# ─── SECTION 1: Utility — Multi-Head Self-Attention and Feed-Forward ──────────

class MultiHeadSelfAttention(nn.Module):
    """Standard multi-head self-attention as used in both PTB and VTB.
    
    For PTB, the input sequence includes the prompt token.
    For VTB, the prompt token is excluded — only patch and CLS tokens.
    All mathematical operations follow Equations 4-8 in the paper.
    """
    def __init__(self, D: int, H: int):
        super().__init__()
        self.H = H
        self.d_h = D // H
        self.scale = math.sqrt(self.d_h)
        self.W_Q = nn.Linear(D, D, bias=False)
        self.W_K = nn.Linear(D, D, bias=False)
        self.W_V = nn.Linear(D, D, bias=False)
        self.W_O = nn.Linear(D, D, bias=False)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, D = x.shape
        Q = self.W_Q(x).view(B, N, self.H, self.d_h).transpose(1, 2)
        K = self.W_K(x).view(B, N, self.H, self.d_h).transpose(1, 2)
        V = self.W_V(x).view(B, N, self.H, self.d_h).transpose(1, 2)
        attn = self.dropout(torch.softmax(Q @ K.transpose(-2, -1) / self.scale, dim=-1))
        out = (attn @ V).transpose(1, 2).contiguous().view(B, N, D)
        return self.W_O(out)


class FeedForwardNetwork(nn.Module):
    """Position-wise FFN with GELU activation, as in standard ViT blocks."""
    def __init__(self, D: int, ffn_dim: int = None):
        super().__init__()
        ffn_dim = ffn_dim or D * 4
        self.net = nn.Sequential(
            nn.Linear(D, ffn_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(ffn_dim, D),
            nn.Dropout(0.1),
        )

    def forward(self, x): return self.net(x)


# ─── SECTION 2: PTB — Prompting Transformer Block ────────────────────────────

class PromptingTransformerBlock(nn.Module):
    """PTB: processes the full token sequence including the prompt token.
    
    Implements the operation described in Equations 1-9:
    Z^l = [x_cls, x_pro, x_1, ..., x_N]
    After PTB, the updated prompt token x_pro' is extracted and forwarded
    to the Prompt Learning module; [CLS] and patch tokens go to VTB.
    """
    def __init__(self, D: int, H: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(D)
        self.attn = MultiHeadSelfAttention(D, H)
        self.ln2 = nn.LayerNorm(D)
        self.ffn = FeedForwardNetwork(D)

    def forward(self, x_cls: torch.Tensor, x_pro: torch.Tensor,
                X: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Parameters
        ----------
        x_cls : (B, 1, D) classification token
        x_pro : (B, 1, D) prompt token
        X     : (B, N, D) patch tokens

        Returns updated (x_cls', x_pro', X')
        """
        B, N, D = X.shape
        # Concatenate full sequence including prompt token
        Z = torch.cat([x_cls, x_pro, X], dim=1)          # (B, N+2, D)
        Z = Z + self.attn(self.ln1(Z))
        Z = Z + self.ffn(self.ln2(Z))
        x_cls_p = Z[:, 0:1, :]                             # updated CLS
        x_pro_p = Z[:, 1:2, :]                             # updated prompt
        X_p = Z[:, 2:, :]                                   # updated patch tokens
        return x_cls_p, x_pro_p, X_p


# ─── SECTION 3: VTB — Vanilla Transformer Block ───────────────────────────────

class VanillaTransformerBlock(nn.Module):
    """VTB: refines CLS and patch tokens without the prompt token.
    
    Implements Equation 10: [x_cls^{l+1}, X^{l+1}] = B_van([x_cls', X']).
    By intentionally excluding the prompt token from this stage, the VTB
    allows image features to develop independently of diagnostic priors,
    creating a clean separation of concerns between prompting and refinement.
    """
    def __init__(self, D: int, H: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(D)
        self.attn = MultiHeadSelfAttention(D, H)
        self.ln2 = nn.LayerNorm(D)
        self.ffn = FeedForwardNetwork(D)

    def forward(self, x_cls: torch.Tensor,
                X: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Operates on CLS + patch tokens only (no prompt token)."""
        Z = torch.cat([x_cls, X], dim=1)                   # (B, N+1, D)
        Z = Z + self.attn(self.ln1(Z))
        Z = Z + self.ffn(self.ln2(Z))
        return Z[:, 0:1, :], Z[:, 1:, :]                 # updated CLS, X


# ─── SECTION 4: Prompt Learning Module ───────────────────────────────────────

class PromptLearningModule(nn.Module):
    """Refines the prompt token via cross-attention over class prototypes.
    
    Implements Equations 11-15:
      Q = x_pro' W_Q  (query from updated prompt token)
      K = P^l W_K     (keys from class prototype tokens)
      V = P^l W_V     (values from class prototype tokens)
      alpha_i = Softmax_i(QK^T / sqrt(d))
      x_pro^{l+1} = sum_i alpha_i V_i
    
    This mechanism allows the prompt for the next level to be a
    dynamically weighted blend of class-specific semantic anchors,
    enabling evidence-driven correction rather than fixed propagation.
    """
    def __init__(self, D: int, d: int = None):
        super().__init__()
        d = d or D
        self.W_Q = nn.Linear(D, d, bias=False)
        self.W_K = nn.Linear(D, d, bias=False)
        self.W_V = nn.Linear(D, d, bias=False)
        self.out_proj = nn.Linear(d, D, bias=False)
        self.scale = math.sqrt(d)
        self.ln = nn.LayerNorm(D)

    def forward(self, x_pro_prime: torch.Tensor,
                prototypes: torch.Tensor) -> torch.Tensor:
        """
        Parameters
        ----------
        x_pro_prime : (B, 1, D) updated prompt token from PTB
        prototypes  : (C_l, D) class prototype tokens for current level

        Returns
        -------
        x_pro_next : (B, 1, D) refined prompt token for next level
        """
        B = x_pro_prime.shape[0]
        C = prototypes.shape[0]
        # Expand prototypes to batch dimension
        P = prototypes.unsqueeze(0).expand(B, -1, -1)          # (B, C, D)
        Q = self.W_Q(x_pro_prime)                               # (B, 1, d)
        K = self.W_K(P)                                         # (B, C, d)
        V = self.W_V(P)                                         # (B, C, d)
        # Attention weights alpha: how much to attend to each class prototype
        alpha = torch.softmax(Q @ K.transpose(-2, -1) / self.scale, dim=-1)
        attended = alpha @ V                                     # (B, 1, d)
        x_pro_next = x_pro_prime + self.out_proj(attended)      # residual update
        return self.ln(x_pro_next)


# ─── SECTION 5: Prototype Learning with EMA Update ───────────────────────────

class PrototypeLearning(nn.Module):
    """Maintains per-class prototype tokens at each diagnostic level.
    
    Prototypes are updated after each mini-batch using exponential moving
    average (EMA) as described in Equations 16-17:
    
      x_bar_c = mean of CLS tokens for class c in current batch
      p_c^t = alpha * p_c^{t-1} + (1 - alpha) * x_bar_c
    
    When class c is absent from the batch, its prototype retains its
    previous value unchanged. This provides stability for rare classes
    that may appear in only a fraction of training batches.
    
    Parameters
    ----------
    C : number of disease classes at this level
    D : token dimension (768)
    alpha : EMA momentum coefficient (default 0.99 per paper)
    """
    def __init__(self, C: int, D: int, alpha: float = 0.99):
        super().__init__()
        self.C = C
        self.alpha = alpha
        # Prototype tokens initialized from N(0,1) as in Section 3.3
        self.register_buffer('prototypes', torch.randn(C, D))

    def get_prototypes(self) -> torch.Tensor:
        """Return current prototype embeddings: (C, D)."""
        return self.prototypes

    @torch.no_grad()
    def update(self, x_cls_batch: torch.Tensor,
                labels: torch.Tensor) -> None:
        """EMA prototype update for one training batch.
        
        Parameters
        ----------
        x_cls_batch : (B, D) CLS tokens for the current batch
        labels      : (B, C) multi-label ground-truth at this level
                      (binary matrix — a subject can have multiple labels)
        """
        for c in range(self.C):
            # Find samples with ground-truth label c present
            mask = labels[:, c].bool()
            if mask.sum() > 0:
                x_bar_c = x_cls_batch[mask].mean(dim=0)     # mean CLS for class c
                self.prototypes[c] = (
                    self.alpha * self.prototypes[c]
                    + (1 - self.alpha) * x_bar_c
                )
                # If class absent: prototype unchanged (no-op)


# ─── SECTION 6: Level-Specific ViT Block ──────────────────────────────────────

class HierarchicalViTBlock(nn.Module):
    """One level of the HP2L backbone: PTB + Prompt Learning + VTB.
    
    The complete inference loop at level l:
    1. PTB: process [CLS, prompt, patches] together → updated tokens
    2. Prompt Learning: refine prompt via cross-attention over prototypes
    3. VTB: refine CLS and patch tokens without prompt token
    4. Classification head: predict y_hat^l from updated CLS token
    5. Prototype update (during training): EMA update of prototype tokens
    """
    def __init__(self, D: int, H: int, C_l: int):
        super().__init__()
        self.ptb = PromptingTransformerBlock(D, H)
        self.prompt_learn = PromptLearningModule(D)
        self.vtb = VanillaTransformerBlock(D, H)
        self.prototype_learn = PrototypeLearning(C_l, D)
        # Classification head: project CLS token to C_l logits
        self.cls_head = nn.Linear(D, C_l)
        self.C_l = C_l

    def forward(
        self,
        x_cls: torch.Tensor,
        x_pro: torch.Tensor,
        X: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
    ) -> Dict:
        """
        Parameters
        ----------
        x_cls  : (B, 1, D) CLS token from previous level
        x_pro  : (B, 1, D) prompt token from previous level
        X      : (B, N, D) patch tokens
        labels : (B, C_l) ground-truth labels for EMA update (training only)

        Returns dict with next-level tokens, logits, and updated CLS.
        """
        # (1) Prompting Transformer Block
        x_cls_p, x_pro_p, X_p = self.ptb(x_cls, x_pro, X)

        # (2) Prompt Learning: refine prompt via prototype cross-attention
        prototypes = self.prototype_learn.get_prototypes()
        x_pro_next = self.prompt_learn(x_pro_p, prototypes)

        # (3) Vanilla Transformer Block: refine CLS and patches
        x_cls_next, X_next = self.vtb(x_cls_p, X_p)

        # (4) Classification head: predict logits from updated CLS
        logits = self.cls_head(x_cls_next.squeeze(1))             # (B, C_l)

        # (5) Prototype EMA update (training only)
        if labels is not None and self.training:
            self.prototype_learn.update(x_cls_next.squeeze(1).detach(), labels)

        return {
            'x_cls': x_cls_next,
            'x_pro': x_pro_next,
            'X': X_next,
            'logits': logits,
        }


# ─── SECTION 7: Image Feature Encoder (3D CNN → token sequence) ──────────────

class BrainMRIEncoder(nn.Module):
    """3D CNN encoder: maps (B, 5, H, W, D) → (B, 32, token_dim).
    
    Performs four spatial downsampling stages as in Section 4.2,
    producing a dense feature grid of size 32 x 12 x 14 x 12.
    The 32 spatial tokens are linearly projected to D=768.
    Input: 5 MRI sequences (T1, T2, FLAIR, DWI, ADC) concatenated as channels.
    Missing modalities are zero-filled before concatenation.
    """
    def __init__(self, in_channels: int = 5, D: int = 768):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv3d(in_channels, 32, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(8, 32), nn.GELU(),
            nn.Conv3d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(16, 64), nn.GELU(),
            nn.Conv3d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(16, 128), nn.GELU(),
            nn.Conv3d(128, 32, kernel_size=3, stride=2, padding=1),
            nn.GroupNorm(8, 32), nn.GELU(),
        )
        # After 4 downsampling stages: spatial dims -> /16 each axis
        # Project 32 feature channels to token dimension D
        self.token_proj = nn.Linear(32, D)
        self.D = D

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Parameters
        ----------
        x : (B, 5, H, W, D_depth) multi-modal 3D brain MRI

        Returns
        -------
        tokens : (B, N, D) flattened spatial tokens with N = HxWxD / 16^3
        """
        feat = self.cnn(x)               # (B, 32, H', W', D')
        B, C, H, W, Dp = feat.shape
        tokens = feat.permute(0, 2, 3, 4, 1).reshape(B, H * W * Dp, C)
        return self.token_proj(tokens)  # (B, N, D)


# ─── SECTION 8: HP2L Full Framework ──────────────────────────────────────────

class HP2L(nn.Module):
    """HP2L: Hierarchical Prompt and Prototype Learning for BD Diagnosis.
    
    Implements the full three-level hierarchical classification framework
    described in Sections 3.1-3.4 and Algorithm 1 of the paper.
    
    Three diagnostic levels:
      Level 1: C1 = 4  (vascular, occupying, developmental, normal)
      Level 2: C2 = 6  (hemorrhage, infarction, WMH, tumor, degeneration,
                         deformity)
      Level 3: C3 = 16 (23 BDs; 15 disease subtypes + normal)
    
    At inference: single forward pass produces predictions for all levels.
    No additional post-processing required for clinical deployment.
    """
    def __init__(self,
                 class_hierarchy: List[int] = [4, 6, 16],
                 D: int = 768,
                 H: int = 12,
                 in_channels: int = 5):
        super().__init__()
        self.L = len(class_hierarchy)
        self.D = D

        # 3D CNN feature encoder
        self.encoder = BrainMRIEncoder(in_channels, D)

        # Learnable initial CLS and prompt tokens
        self.x_cls_init = nn.Parameter(torch.randn(1, 1, D))
        self.x_pro_init = nn.Parameter(torch.randn(1, 1, D))

        # One HierarchicalViTBlock per diagnostic level
        self.levels = nn.ModuleList([
            HierarchicalViTBlock(D, H, C_l)
            for C_l in class_hierarchy
        ])

        # Label smoothing factor as per Section 4.2
        self.label_smoothing = 0.1

    def forward(self, images: torch.Tensor,
                labels_per_level: Optional[List[torch.Tensor]] = None
                ) -> Dict:
        """
        Parameters
        ----------
        images          : (B, 5, H, W, D_depth) multi-modal 3D brain MRI
        labels_per_level: list of L label tensors, each (B, C_l), for
                          prototype EMA update during training (optional)

        Returns
        -------
        dict with:
          'logits_per_level' : list of (B, C_l) tensors
          'cls_per_level'    : list of (B, 1, D) CLS tokens for attention viz
        """
        B = images.shape[0]
        # Extract patch token sequence from 3D CNN encoder
        X = self.encoder(images)                             # (B, N, D)

        # Initialize CLS and prompt tokens, expand to batch size
        x_cls = self.x_cls_init.expand(B, -1, -1).clone()  # (B, 1, D)
        x_pro = self.x_pro_init.expand(B, -1, -1).clone()  # (B, 1, D)

        logits_per_level = []
        cls_per_level = []

        for l, level_block in enumerate(self.levels):
            labels_l = labels_per_level[l] if labels_per_level else None
            out = level_block(x_cls, x_pro, X, labels=labels_l)

            # Pass updated tokens to next level
            x_cls = out['x_cls']
            x_pro = out['x_pro']
            X = out['X']

            logits_per_level.append(out['logits'])
            cls_per_level.append(x_cls)

        return {
            'logits_per_level': logits_per_level,
            'cls_per_level': cls_per_level,
        }


# ─── SECTION 9: Loss Functions ────────────────────────────────────────────────

def level_classification_loss(logits: torch.Tensor,
                              labels: torch.Tensor,
                              smoothing: float = 0.1) -> torch.Tensor:
    """Level-wise sigmoid binary cross-entropy with label smoothing (Eq. 18).
    
    Uses BCEWithLogits applied element-wise over all class labels.
    Label smoothing prevents overconfidence on noisy multi-center labels.
    """
    labels_smooth = labels.float() * (1 - smoothing) + 0.5 * smoothing
    return F.binary_cross_entropy_with_logits(logits, labels_smooth)


def hierarchy_consistency_loss(
    logits_per_level: List[torch.Tensor],
    labels_per_level: List[torch.Tensor],
    parent_map: List[Dict[int, int]],
) -> torch.Tensor:
    """Hierarchy-consistency penalty across adjacent diagnostic levels (Eq. 19).
    
    For each true child label c, penalize max(0, p_parent - p_child).
    This asymmetric penalty never suppresses fine-grained confidence
    that exceeds coarse confidence — it only fires when fine-grained
    evidence lags behind coarse evidence for a true positive class.
    
    Parameters
    ----------
    logits_per_level : list of (B, C_l) logit tensors
    labels_per_level : list of (B, C_l) ground-truth label tensors
    parent_map       : list of dicts mapping child class index -> parent
                       class index at the preceding level
    """
    consist_loss = torch.tensor(0.0, device=logits_per_level[0].device)
    B = logits_per_level[0].shape[0]

    for l in range(len(logits_per_level) - 1):
        probs_coarse = torch.sigmoid(logits_per_level[l])    # (B, C_l)
        probs_fine = torch.sigmoid(logits_per_level[l + 1]) # (B, C_{l+1})
        labels_fine = labels_per_level[l + 1].float()

        for child_idx, parent_idx in parent_map[l].items():
            if child_idx >= probs_fine.shape[1] or parent_idx >= probs_coarse.shape[1]:
                continue
            p_child = probs_fine[:, child_idx]
            p_parent = probs_coarse[:, parent_idx]
            y_child = labels_fine[:, child_idx]
            penalty = y_child * torch.clamp(p_parent - p_child, min=0)
            consist_loss = consist_loss + penalty.mean()

    return consist_loss


def hp2l_total_loss(
    logits_per_level: List[torch.Tensor],
    labels_per_level: List[torch.Tensor],
    parent_map: List[Dict[int, int]],
    lambda_levels: List[float] = [1.0, 1.0, 1.0],
    lambda_consist: float = 0.5,
    label_smoothing: float = 0.1,
) -> Dict[str, torch.Tensor]:
    """Full training objective (Eq. 20): level-wise BCE + consistency penalty.
    
    Returns dict with 'total', 'cls_losses', and 'consist' for logging.
    """
    cls_losses = [
        level_classification_loss(logits_per_level[l], labels_per_level[l],
                                   label_smoothing)
        for l in range(len(logits_per_level))
    ]
    total_cls = sum(lambda_levels[l] * cls_losses[l] for l in range(len(cls_losses)))
    consist = hierarchy_consistency_loss(logits_per_level, labels_per_level, parent_map)
    total = total_cls + lambda_consist * consist
    return {'total': total, 'cls_losses': cls_losses, 'consist': consist}


# ─── SECTION 10: Smoke Test — Synthetic Brain MRI Data ───────────────────────

def _smoke_test():
    """End-to-end smoke test of HP2L on synthetic 3D brain MRI data.
    
    Verifies:
      - Forward pass through full 3-level hierarchical backbone
      - EMA prototype updates during training
      - Loss computation with hierarchy-consistency penalty
      - Gradient flow through all components
    """
    print("=" * 65)
    print("HP2L Smoke Test — Synthetic 3D Brain MRI Data")
    print("Paper: Liu et al., Medical Image Analysis 112 (2026) 104063")
    print("=" * 65)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(ff"\nDevice: {device}")

    # Small model config for fast smoke test
    B = 4                           # batch size
    H_vol, W_vol, D_vol = 64, 64, 32  # reduced spatial dimensions
    D_model = 256                   # token dimension (768 in paper)
    num_heads = 8
    class_hierarchy = [4, 6, 16]  # 3-level BD hierarchy from paper

    # Synthetic multi-modal 3D brain MRI: (B, 5 modalities, H, W, D)
    images = torch.randn(B, 5, H_vol, W_vol, D_vol, device=device)

    # Synthetic multi-label ground-truth for each diagnostic level
    labels = [
        torch.randint(0, 2, (B, C), device=device).float()
        for C in class_hierarchy
    ]

    # Simple parent-child mapping for consistency loss
    # Level 0->1: coarse class -> intermediate (example mapping)
    parent_map = [
        {0: 0, 1: 0, 2: 0, 3: 1, 4: 1, 5: 2},         # level 1->2
        {i: i // 3 for i in range(16)},                 # level 2->3
    ]

    # Instantiate HP2L with small model config
    model = HP2L(
        class_hierarchy=class_hierarchy,
        D=D_model,
        H=num_heads,
        in_channels=5,
    ).to(device)
    model.train()

    total_params = sum(p.numel() for p in model.parameters())
    print(ff"Total parameters: {total_params:,}")
    print(ff"Input shape: {list(images.shape)}")
    print(ff"Class hierarchy: {class_hierarchy}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2)

    # Forward pass
    out = model(images, labels_per_level=labels)
    logits_per_level = out['logits_per_level']

    # Compute full training objective
    loss_dict = hp2l_total_loss(
        logits_per_level, labels, parent_map,
        lambda_levels=[1.0, 1.0, 1.0],
        lambda_consist=0.5,
        label_smoothing=0.1,
    )

    optimizer.zero_grad()
    loss_dict['total'].backward()
    optimizer.step()

    print(ff"\n{'─'*45}")
    print(ff"Total loss:          {loss_dict['total'].item():.4f}")
    print(ff"Level-1 cls loss:    {loss_dict['cls_losses'][0].item():.4f}")
    print(ff"Level-2 cls loss:    {loss_dict['cls_losses'][1].item():.4f}")
    print(ff"Level-3 cls loss:    {loss_dict['cls_losses'][2].item():.4f}")
    print(ff"Consistency loss:    {loss_dict['consist'].item():.4f}")
    print(ff"Logit shapes: {[list(l.shape) for l in logits_per_level]}")
    print(f"{'─'*45}")

    # Verify inference mode (no labels, no prototype update)
    model.eval()
    with torch.no_grad():
        out_inf = model(images)
    probs = [torch.sigmoid(l) for l in out_inf['logits_per_level']]
    print(ff"\nInference level-3 mean probability: {probs[2].mean().item():.3f}")
    print("Smoke test passed. HP2L forward and backward cycles OK.")
    print("See Algorithm 1 in Liu et al. 2026 for the full inference workflow.")
    print("=" * 65)


if __name__ == '__main__':
    _smoke_test()