HP2L Shows How AI Can Now Think Like a Radiologist and Diagnose 23 Brain Disorders Step by Step

The following is a complete PyTorch implementation of HP2L covering Sections 3.1 through 3.4 of the paper. It includes the full three-level hierarchical ViT backbone with Prompting Transformer Block and Vanilla Transformer Block sub-units, 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  Multi-Head Self-Attention and Feed-Forward Utilities

class MultiHeadSelfAttention(nn.Module):
    """Standard multi-head self-attention used in both PTB and VTB.

    For PTB the input sequence includes the prompt token.
    For VTB the prompt token is excluded and only patch and CLS tokens remain.
    All math follows Equations 4 through 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  Prompting Transformer Block (PTB)

class PromptingTransformerBlock(nn.Module):
    """PTB processes the full token sequence including the prompt token.

    Implements the operation described in Equations 1 through 9.
    Z_l = [x_cls, x_pro, x_1, ..., x_N]
    After PTB the updated prompt token x_pro_prime 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]:
        """
        x_cls  (B, 1, D) classification token
        x_pro  (B, 1, D) prompt token
        X      (B, N, D) patch tokens
        Returns updated (x_cls_prime, x_pro_prime, X_prime)
        """
        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 patches
        return x_cls_p, x_pro_p, X_p


# SECTION 3  Vanilla Transformer Block (VTB)

class VanillaTransformerBlock(nn.Module):
    """VTB refines CLS and patch tokens without the prompt token.

    Implements Equation 10  [x_cls_next, X_next] = B_van([x_cls_prime, X_prime]).
    Excluding the prompt token here lets image features develop independently
    of diagnostic priors, creating a clean separation of concerns.
    """
    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 plus patch tokens only. No prompt token involved."""
        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 and X


# SECTION 4  Prompt Learning Module

class PromptLearningModule(nn.Module):
    """Refines the prompt token through cross-attention over class prototypes.

    Implements Equations 11 through 15.
      Q = x_pro_prime times W_Q    (query from updated prompt token)
      K = P_l times W_K            (keys from class prototype tokens)
      V = P_l times W_V            (values from class prototype tokens)
      alpha_i = Softmax_i(QK^T / sqrt(d))
      x_pro_next = sum_i alpha_i V_i

    This lets the prompt for the next level be a dynamically weighted blend
    of class-specific semantic anchors, enabling evidence-driven correction
    instead of fixed propagation of whatever the prior level decided.
    """
    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:
        """
        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]
        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)
        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 update after each mini-batch using exponential moving average
    as described in Equations 16 through 17.

      x_bar_c = mean of CLS tokens for class c in current batch
      p_c_t   = alpha times p_c_{t-1} + (1 - alpha) times 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 all training batches.

    Parameters
      C      number of disease classes at this level
      D      token dimension (768 in the paper)
      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 of shape (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.

        x_cls_batch  (B, D) CLS tokens for the current batch
        labels       (B, C) multi-label ground-truth at this level
                     (binary matrix where a subject can have multiple labels)
        """
        for c in range(self.C):
            mask = labels[:, c].bool()
            if mask.sum() > 0:
                x_bar_c = x_cls_batch[mask].mean(dim=0)
                self.prototypes[c] = (
                    self.alpha * self.prototypes[c]
                    + (1 - self.alpha) * x_bar_c
                )
                # If class is absent from batch the prototype holds its value


# SECTION 6  Level-Specific ViT Block

class HierarchicalViTBlock(nn.Module):
    """One complete level of the HP2L backbone combining PTB, Prompt Learning, and VTB.

    The full inference loop at level l:
    1. PTB  process [CLS, prompt, patches] together to get updated tokens
    2. Prompt Learning  refine prompt through cross-attention over prototypes
    3. VTB  refine CLS and patch tokens without the prompt token present
    4. Classification head  predict logits from updated CLS token
    5. Prototype update  EMA update of prototype tokens (training only)
    """
    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)
        self.cls_head       = nn.Linear(D, C_l)
        self.C_l            = C_l

    def forward(self, x_cls, x_pro, X, labels=None) -> Dict:
        """
        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.
        """
        # Step 1  Prompting Transformer Block
        x_cls_p, x_pro_p, X_p = self.ptb(x_cls, x_pro, X)

        # Step 2  Prompt Learning through prototype cross-attention
        prototypes = self.prototype_learn.get_prototypes()
        x_pro_next = self.prompt_learn(x_pro_p, prototypes)

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

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

        # Step 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  3D CNN Image Feature Encoder

class BrainMRIEncoder(nn.Module):
    """3D CNN encoder that maps (B, 5, H, W, D) to (B, 32, token_dim).

    Performs four spatial downsampling stages as in Section 4.2.
    Input is 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(),
        )
        self.token_proj = nn.Linear(32, D)
        self.D = D

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x       (B, 5, H, W, D_depth) multi-modal 3D brain MRI
        Returns (B, N, D) flattened spatial tokens
        """
        feat = self.cnn(x)
        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)


# SECTION 8  HP2L Full Framework

class HP2L(nn.Module):
    """HP2L  Hierarchical Prompt and Prototype Learning for Brain Disorder Diagnosis.

    Implements the full three-level hierarchical classification framework
    described in Sections 3.1 through 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 represented as 15 disease subtypes plus normal)

    Single forward pass at inference produces predictions for all three levels.
    No additional post-processing is 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

        self.encoder    = BrainMRIEncoder(in_channels, D)
        self.x_cls_init = nn.Parameter(torch.randn(1, 1, D))
        self.x_pro_init = nn.Parameter(torch.randn(1, 1, D))

        self.levels = nn.ModuleList([
            HierarchicalViTBlock(D, H, C_l)
            for C_l in class_hierarchy
        ])

        self.label_smoothing = 0.1

    def forward(self, images: torch.Tensor,
                labels_per_level: Optional[List[torch.Tensor]] = None) -> Dict:
        """
        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 EMA update
        Returns dict with logits_per_level and cls_per_level
        """
        B = images.shape[0]
        X = self.encoder(images)

        x_cls = self.x_cls_init.expand(B, -1, -1).clone()
        x_pro = self.x_pro_init.expand(B, -1, -1).clone()

        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)

            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, labels_per_level,
                               parent_map) -> torch.Tensor:
    """Hierarchy-consistency penalty across adjacent diagnostic levels (Eq. 19).

    For each true child label c penalize max(0, p_parent minus 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.

    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 to parent index
    """
    consist_loss = torch.tensor(0.0, device=logits_per_level[0].device)

    for l in range(len(logits_per_level) - 1):
        probs_coarse = torch.sigmoid(logits_per_level[l])
        probs_fine   = torch.sigmoid(logits_per_level[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, labels_per_level, parent_map,
                    lambda_levels=[1.0, 1.0, 1.0], lambda_consist=0.5,
                    label_smoothing=0.1) -> Dict:
    """Full training objective (Eq. 20)  level-wise BCE plus 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 on Synthetic 3D Brain MRI Data

def _smoke_test():
    """End-to-end smoke test of HP2L on synthetic 3D brain MRI data.

    Verifies:
      Forward pass through the 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'
    B               = 4
    H_vol, W_vol, D_vol = 64, 64, 32
    D_model         = 256
    num_heads       = 8
    class_hierarchy = [4, 6, 16]

    images = torch.randn(B, 5, H_vol, W_vol, D_vol, device=device)
    labels = [torch.randint(0, 2, (B, C), device=device).float() for C in class_hierarchy]

    parent_map = [
        {0: 0, 1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
        {i: i // 3 for i in range(16)},
    ]

    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(f"\nDevice  {device}")
    print(f"Total parameters  {total_params:,}")
    print(f"Input shape  {list(images.shape)}")
    print(f"Class hierarchy  {class_hierarchy}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2)
    out       = model(images, labels_per_level=labels)
    loss_dict = hp2l_total_loss(out['logits_per_level'], labels, parent_map)

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

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

    model.eval()
    with torch.no_grad():
        out_inf = model(images)
    probs = [torch.sigmoid(l) for l in out_inf['logits_per_level']]
    print(f"\nInference level 3 mean probability  {probs[2].mean().item():.3f}")
    print("Smoke test passed. HP2L forward and backward cycles OK.")
    print("=" * 65)


if __name__ == '__main__':
    _smoke_test()