Normal Guidance is what Attention Needs

Deep learning has revolutionized medical image analysis, particularly for 3D imaging modalities such as CT and MRI. Accurate detection and localization of lesions within volumetric scans are critical for diagnosis and treatment planning. However, acquiring detailed slice-level annotations is prohibitively expensive and time-consuming, motivating the use of weakly supervised learning approaches that rely solely on volume-level labels. Multiple Instance Learning (MIL) has emerged as a prominent framework in this context, modeling each 3D scan as a bag of 2D slice instances with only bag-level labels available during training. Attention-based MIL methods, such as ABMIL, assign weights to slices to highlight regions likely contributing to the diagnosis, offering interpretability and localization capabilities. Transformer-based MIL models, like TransMIL, further incorporate dependencies among slices via self-attention mechanisms. Despite these advances, recent studies have shown that a simple center-focused Gaussian baseline, which ignores image content and assigns attention weights based solely on slice position, can outperform complex attention and transformer-based MIL models in brain CT localization tasks. This counterintuitive finding highlights a critical gap: existing MIL attention mechanisms lack sufficient spatial inductive bias to effectively leverage anatomical priors inherent in medical imaging. Addressing this limitation is essential for improving slice-level lesion localization under weak supervision.

The core problem addressed is how to accurately localize lesion-bearing slices within 3D medical images using only coarse, volume-level binary labels during training. Existing attention-based MIL methods generate attention weights per slice but often fail to outperform naive spatial baselines, indicating insufficient incorporation of spatial prior knowledge. This deficiency hampers clinical applicability, as precise slice-level localization is crucial for interpretability, trust, and diagnostic efficiency. The challenge lies in designing mechanisms that effectively integrate spatial inductive biases reflecting clinical knowledge—such as lesions tending to cluster near central slices—while preserving the flexibility to adapt to image-specific features. Moreover, achieving this without sacrificing overall scan-level classification accuracy and under computational constraints remains difficult. The problem is compounded by the presence of multiple, spatially disjoint lesions and varying scan lengths, necessitating robust and generalizable solutions.

The paper introduces several key innovations:

1. Normal Guidance Regularization: A novel framework that guides learned attention weights toward a discrete normal distribution parameterized by the empirical mean and variance of attention per scan. This explicitly encodes spatial prior knowledge that lesions are more likely near the scan center, addressing the lack of inductive bias in prior MIL attention mechanisms.

2. Multi-Head Normal Guidance: An extension for transformer-based MIL that applies separate normal guidance to each attention head, enabling the model to attend to multiple spatially disjoint regions simultaneously. This overcomes the unimodal limitation of single normal guidance and better models complex lesion distributions.

3. Comprehensive Empirical Validation: Extensive experiments on three large-scale, diverse CT datasets totaling over 4 million slices demonstrate the method’s superiority in slice-level localization and competitive scan-level classification, establishing new state-of-the-art results and practical performance ceilings.

These innovations collectively advance weakly supervised 3D medical image analysis by integrating clinically meaningful spatial priors into attention mechanisms, enhancing interpretability and localization accuracy.

Method details:

• �� Input Preparation: Each 3D CT scan is decomposed into a variable number of 2D axial slices. Each slice is encoded into a fixed-length embedding vector using a frozen Vision Transformer (ViT) pretrained on ImageNet, producing instance-level features.

• �� MIL Framework: The set of slice embeddings forms a bag input to MIL models. Two architectures are considered: ABMIL, which uses a learned attention pooling mechanism to aggregate slice embeddings, and TransMIL, which employs multi-head self-attention to model inter-slice dependencies.

• �� Normal Guidance Construction:
• For each scan, compute the empirical mean (μ) and variance (σ²) of the learned attention weights across slices.
• Construct a discrete normal probability density function (PDF) over slice indices using μ and σ².
• Normalize this PDF to form the reference attention distribution.

• �� Regularization Objective:
• Define a divergence metric (forward KL divergence, reverse KL divergence, or squared error) between the learned attention distribution and the reference normal distribution.
• Incorporate this divergence as a regularization term weighted by a hyperparameter λ into the standard binary cross-entropy loss for scan-level classification.
• During backpropagation, apply stop-gradient to the reference distribution to prevent its parameters from updating.

• �� Multi-Head Extension:
• For TransMIL’s multi-head attention, compute separate normal reference distributions per head.
• Average the divergence terms across heads to form the total regularization loss.

• �� Training Setup:
• Use only scan-level binary labels for training, validation, and hyperparameter tuning.
• Evaluate slice-level localization using expert-annotated slice labels only at test time.
• Optimize using stochastic gradient descent with momentum, early stopping based on validation AUROC, and grid search over learning rates and λ.

This methodology effectively biases attention weights toward clinically plausible spatial patterns while preserving data-driven adaptability.

Experimental design:

• �� Datasets: Three large-scale public CT datasets covering different anatomical regions and pathologies:
• Head CT: 21,744 scans, 752,803 slices, intracranial hemorrhage labels.
• Chest CT: 7,279 scans, 1,790,594 slices, pulmonary embolism labels.
• Abdomen CT: 4,711 scans, 1,500,653 slices, abdominal trauma labels.

• �� Baselines: ABMIL, TransMIL, Smooth Operator (a smoothing-based MIL method), and a center-focused Gaussian baseline ignoring image content.

• �� Metrics: Slice-level localization evaluated by AUROC on positive bags only; scan-level classification evaluated by AUROC; AUPRC results provided in appendix.

• �� Upper Bounds: Constructed best-in-class ceilings for localization and classification using instance-level labels and oracle pooling to contextualize results.

• �� Training Details: Frozen ViT encoder with embedding size 768; linear classifier head; batch size 64; 1000 epochs; early stopping on validation AUROC; hyperparameter grid search.

• �� Ablations: Compared divergence types (forward KL, reverse KL, squared error), single vs. multi-head guidance, and regularization strengths.

• �� Reproducibility: Three random train/validation/test splits stratified by patient and class; results averaged with standard deviations reported.

This rigorous setup ensures robust evaluation of Normal Guidance’s effectiveness and generalizability.

Key results:

• �� Slice-Level Localization: Normal Guidance achieves AUROC of 0.871 (Head CT), 0.866 (Chest CT), and 0.663 (Abdomen CT), outperforming ABMIL, TransMIL, Smooth Operator, and the center-focused Gaussian baseline (e.g., Chest CT baseline 0.78 vs. NG 0.866).

• �� Scan-Level Classification: Maintains competitive AUROC (e.g., 0.925 on Head CT), close to upper bounds derived from instance-level labels, indicating no trade-off between localization regularization and classification accuracy.

• �� Multi-Head Normal Guidance: Further improves localization, especially in multi-lesion scenarios, reaching 0.706 AUROC on semi-synthetic data near the best-in-class ceiling of 0.884.

• �� Divergence Choice: Forward KL divergence yields slightly better localization than reverse KL or squared error.

• �� Sensitivity Analysis: Moderate regularization strength optimizes performance; overly strong or weak regularization degrades results.

• �� Qualitative Analysis: Attention maps guided by Normal Guidance are more focused and aligned with expert annotations, enhancing interpretability.

• �� Transformer MIL vs. ABMIL: Transformer-based models marginally improve classification but not localization without Normal Guidance, underscoring the importance of spatial prior regularization.

Applications:

• �� Weakly Supervised Lesion Localization: Enables accurate slice-level lesion identification in 3D medical images using only volume-level labels, reducing annotation burden and facilitating clinical workflows.

• �� Clinical Decision Support: Integrates into radiology pipelines to highlight suspicious slices, aiding radiologists in diagnosis, improving efficiency and trust.

• �� Multi-Organ CT Analysis: Applicable across brain, chest, and abdomen CT scans, supporting diverse pathologies such as hemorrhage, embolism, and trauma.

• �� Extension to Other Modalities: Potential adaptation to MRI, PET, and whole-slide pathology images for broader medical imaging tasks.

• �� Research Tool: Provides a benchmark and framework for developing spatially informed weakly supervised learning methods in medical imaging.

Limitations:

• �� Negative Bag Attention: The method does not explicitly regularize attention distributions for negative bags, leaving ambiguity in attention behavior when no lesion is present, which may affect model interpretability and robustness.

• �� Causal Interpretability: Although attention aligns better with expert annotations, it does not guarantee causal attribution, limiting the reliability of attention as an explanation for model decisions.

• �� Model Expressiveness: Using a frozen ViT encoder and linear classifier restricts representational capacity; end-to-end fine-tuning could improve results but at high computational cost.

• �� Computational Resources: Training transformer-based MIL models with Normal Guidance is resource-intensive, potentially limiting scalability.

• �� Dataset Variability: Performance gaps on abdomen CT may reflect organ coverage variability and incomplete scans, suggesting dataset-specific challenges.