Normal Guidance is what Attention Needs
arXiv:2605.27306v1 [cs.LG] 26 May 2026
Ethan Harvey∗ Department of Computer Science Tufts University [email protected]
Dennis Johan Loevlie∗ Department of Computer Science Tufts University [email protected]
Michael C. Hughes Department of Computer Science Tufts University [email protected]
Abstract We consider training classifiers for 3D medical images using only one binary label for the entire volume rather than a label for each 2D slice. In such weakly supervised settings, can we learn accurate classifiers for slice-level predictions? Attention-based multiple instance learning (MIL) can produce an attention score for every slice. Yet recent work demonstrates that a simple center-focused baseline that ignores image content can outperform attention-based and transformer-based MIL at slice-level classification of 3D brain scans. We show this baseline also outperforms existing MIL at slice-level classification of thoracic and abdominal CT scans. Motivated by this baseline, we propose Normal Guidance, a regularization technique that encourages the learned attention distribution to follow a bell-shaped curve. Across three medical imaging datasets totaling over 4 million 2D slices, we show our Normal Guidance enables attention-based and transformer-based MIL methods to deliver significantly better slice-level localization than the state-of-theart while remaining competitive at whole-scan classification.
1
Introduction
Deep learning has become a go-to approach for prediction tasks involving 3D medical images [Wang et al., 2017, Saab et al., 2019, Wu et al., 2021, Castro-Macı́as et al., 2024]. The primary clinical task is often to classify the presence or absence of lesions or disease in the overall organ being imaged (e.g. brain or lungs). A complementary task is localization, where the goal is to indicate which specific 2D slices of the 3D volume show evidence for the overall prediction. Localization models predict a (probabilistic) binary label for each slice. Models with high-quality localization allow clinicians to audit predictions, hopefully improving interpretability, trust, and accountability. Yet in many applied tasks, especially with 3D medical scans, the expense of acquiring expert labels means that model development can only be “weakly” supervised: coarse scan-level binary labels are available for each 3D scan in the training set, but not fine-grained labels for individual slices. An unsolved technical challenge is thus to develop effective slice-level localization from weak supervision, while preserving competitive whole-scan classification. Deep multiple instance learning (MIL) is a strategy for such weakly-supervised learning [Quellec et al., 2017, Saab et al., 2019]. MIL methods can take as input a variable-sized set of instances (a “bag”), each with their own feature vector, and produce one predicted binary label. For localization ∗ Equal contribution
Code: github.com/tufts-ml/normal-guidance
Preprint.
with deep MIL, attention-based MIL [Ilse et al., 2018] has become the foundation of many methods in the last decade due to its by-design interpretability: each instance j in bag i has an attention weight aij ∈ [0, 1] indicating its relative importance for the classification task. Overall, localization in MIL remains underexplored [Castro-Macı́as et al., 2024]. Past works have used attention weights for localization primarily through qualitative results [Ilse et al., 2018, Lu et al., 2021, Shao et al., 2021, Keshvarikhojasteh et al., 2024] that visually highlight regions of interest (ROIs) with high attention values. A few works also quantitatively evaluate localization [Li et al., 2021, Fourkioti et al., 2024, Castro-Macı́as et al., 2024]. Recently, Harvey et al. [2026] applied MIL to several datasets of 3D brain scans, including computed tomography (CT) and magnetic resonance imaging (MRI). Harvey et al. showed a simple center-focused baseline that ignores all image content can outperform the learned attention from modern MIL for localization. The attention values for 3D scan i in this baseline are set as aij ∝ NormPDF(j| S2i , 1) where Si is the number of 2D axial slices in the 3D scan (“bag”) and integer j indexes each 2D slice (“instance”) from 1 to Si in spatial order. This discretized univariate Normal PDF results in bell-shaped attention when plotted as a function of j, with a peak at the middle slice. It reflects the clinical intuition that the center of the brain is more relevant than the edges for prediction tasks like detecting the presence of lesions or white matter disease. In this paper, we offer three contributions to improve the development of weakly-supervised classifiers of 3D medical images, focusing on CT scans due to large, high-quality available datasets: • We show that the simple center-focused baseline from Harvey et al. [2026] outperforms modern transformer-based MIL at slice localization tasks that span several body parts, including thoracic CT (lungs) and abdominal CT (kidneys, spleen, liver and bowel). This result expands the brainspecific results of past work and suggests current MIL methods are not adequate for localization in 3D medical imaging even given large labeled data. Better inductive biases are clearly needed. • Motivated by this baseline, we propose Normal Guidance for the attention mechanism within deep MIL. We use a normal distribution parameterized by the empirical mean and variance of the learned attention weights to regularize the learned attention weights. This inductive bias encourages the attention as a function of j to form a smooth bell-curve shape. Normal Guidance does not force all scans to focus on the exact center; instead image-specific content can inform attention when helpful. We demonstrate that Normal Guidance enables attention-based and transformerbased MIL to perform better than the simple baseline for localization (Tab. 1) while remaining competitive at whole scan classification (Tab. 2). • We develop methods for quantifying a best-in-class ceiling for localization (Sec. 4.3) and scanlevel classification (Sec. 4.4). These are approximate practical upper bounds for what our chosen MIL architectures (pooled frozen embeddings with linear classifier heads) are capable of given the available training data. These methods unfairly using instance-level labels to establish bestcase performance for weakly-supervised MIL. While other MIL methods remain far from the localization ceiling on all tasks, we show that MIL with Normal Guidance comes much closer (within 0.03 AUROC of ceiling) on two CT tasks. Our supporting evidence for these contributions is drawn from experiments on three open-access medical imaging datasets of CT scans totaling over 4 million 2D slices. The datasets span different body parts (head, chest, abdomen), and the classification tasks cover acute injuries, embolisms, and lesions. Each dataset includes human expert annotations for every slice. All non-ceiling models are trained strictly in the weak-supervised regime, using only scan-level labels for training, early stopping, and hyperparameter selection; slice-level labels are only used to evaluate localization quality. We hope this work improves weakly-supervised localization for 3D imaging tasks and advances understanding of what inductive biases are needed for deep MIL to move from the lab to the clinic.
2
Background and Related Works
MIL [Dietterich et al., 1997, Maron and Lozano-Pérez, 1997, Quellec et al., 2017] is a branch of weakly supervised learning. In MIL, the training dataset D = {(xi , yi )}N i=1 consists of N bags, where each bag xi = {xi,1 , . . . , xi,Si } contains a set of Si instances. Although instance-level labels exist {yi,1 , . . . , yi,Si } where ∀j : yij ∈ {0, 1}, only the bag-level label yi ∈ {0, 1} is observed during training. The standard MIL assumption [Dietterich et al., 1997, Raff and Holt, 2023] is that a 2
bag is negative if and only if all instances are negative (yi = 0 ⇔ ∀j ∈ {1, . . . , Si } : yij = 0) and positive if and only if at least one instance is positive (yi = 1 ⇔ ∃j ∈ {1, . . . , Si } : yij = 1). MIL architectures typically consist of three parts: an encoder, pooling operation, and classifier. First, an encoder f (·) : RC×H×W → RM maps each C-channel 2D image xij ∈ RC×H×W to an instance-level embedding hij = f (xij ) ∈ RM . Second, a pooling operation σ(·) : RM ×Si → RM maps the set of instance-level embeddings hi = {hi,1 , . . . , hi,Si } to a single bag-level embedding zi = σ(hi ) ∈ RM . Finally, a classifier g(·) : RM → [0, 1] maps the bag-level embedding to a predicted probability ŷi = g(zi ) ∈ [0, 1]. Permutation-invariant pooling. Early MIL methods [Pinheiro and Collobert, 2015, Zhu et al., 2017, Feng and Zhou, 2017] used simple, non-learnable pooling operations to aggregate instancelevel embeddings. Ilse et al. [2018] proposed learnable attention-based MIL (ABMIL) to be both data-driven and interpretable. Extensions to ABMIL have maintained permutation-invariance (e.g., CLAM [Lu et al., 2021], DSMIL [Li et al., 2021], MAD-MIL [Keshvarikhojasteh et al., 2024]). Correlated MIL. Permutation-invariance is inappropriate in cases with clear spatial dependencies across instances, such as bags of patches drawn from the same high-resolution 2D image or bags of 2D slices from a 3D volume. To address this, recent work has proposed correlated MIL [Shao et al., 2021] where transformer-based MIL methods (e.g., TransMIL [Shao et al., 2021], CAMIL [Fourkioti et al., 2024]) use multi-head self-attention to incorporate dependencies between instances. Smooth operator. To improve localization while accounting for dependencies, Castro-Macı́as et al. [2024] proposed a smoothing operator that smoothes the learned embeddings of neighboring instances. Their module achieved state-of-the-art localization results among all tested MIL methods on 3D neuroimage and whole slide imaging (WSI) classification tasks. Centered Gaussian. Harvey et al. [2026] raised questions about the quality of modern MIL for localization in brain scans. They examined a simple baseline: a discretized Gaussian distribution centered on the middle 2D slice of a 3D brain scan, aij ∝ NormPDF(j| S2i , 1). They found this baseline outperforms the smoothing operator and other attention-based and transformer-based MIL at localizing in 3D neuroimages. Alternatives for attention regularization. In MIL for WSI, prior works have proposed regularization techniques for learned attention weights. Sharma et al. [2021] proposed cluster-to-cluster, which adds a Kullback-Leibler (KL) divergence term between the attention weights of patches in a cluster and the uniform distribution. Zhang et al. [2025] introduced attention entropy maximization, which subtracts an entropy term on the attention weights to penalize excessive attention concentration. Peled et al. [2026] integrates 2D spatial context into pairwise multi-head attention through a learnable distance-decayed prior. Unlike these methods, our Normal Guidance encourages a univariate learned attention distribution to follow a bell-shaped curve. For further expanded discussion of related work, see App. G.
3
Methods
3.1
Guided Attention
We formulate a guided attention framework designed to regularize the learned attention weights ai by penalizing their divergence from a discrete reference distribution ri over the Si instances in bag i. The guided objective that we minimize during training is defined as: PN LGuided := N1 i=1 ℓBCE (yi , ŷi ) + λD(ri , ai ) (1) where ℓBCE is the standard binary cross entropy loss, λ ≥ 0 is a hyperparameter that controls the strength of the regularization term, and D(ri , ai ) is a divergence measure quantifying the “distance” between the reference distribution and the learned attention weights for bag i. To construct the reference ri , we recommend a Normal Guidance procedure described in Sec. 3.2. We apply D(ri , ai ) to all bags (positive and negative), with each ri constructed from the attention values of bag i. Designing an ri that depends on bag label is possible, but we leave this to future work. Some tasks could imagine using instance-level labels to inform the reference ri (see App. F), but this requires available labels and breaks the assumptions of MIL. 3
1
5
9
14 18 22 27 31 35 40
Squared Error (mean-seeking)
1
Attention weight aij
ABMIL
5
9
14 18 22 27 31 35 40
Forward KL (mode-covering)
1
ABMIL
5
9
14 18 22 27 31 35 40
Reverse KL (mode-seeking)
1
5
9
14 18 22 27 31 35 40
Attention Weights after Training
ABMIL
ABMIL Squared error Forward KL Reverse KL
0.06 0.04 0.02 0.00
0
10
20 Slice index j
30
40
0
10
20 Slice index j
30
40
0
10
20 Slice index j
30
40
0
10
20 Slice index j
30
40
Figure 1: Normal Guidance on one 3D CT scan under different divergences D. Left panels: using squared error (orange), forward KL (green), and reverse KL (red). Arrows indicate direction of change for each attention weight aij to reduce D(·) via a gradient step; line length indicates magnitude of the regularization term for each attention weight. Right: Attention weights ai after training with standard ABMIL or with Normal Guidance. Normal Guidance improves the distribution’s unimodality and focus, though recovering the true ROI (red region in top bar) remains imperfect. Among possible choices for divergence D(·), we study three variants with different behavior, described below and compared side-by-side in Fig. 1: PSi • Squared error (mean-seeking). The squared error D(ri , ai ) = j=1 (rij − aij )2 is symmetric, penalizes large errors more than small errors, and generally encourages ai to be mean-seeking. PSi rij • Forward KL (mode-covering). The forward KL DKL (ri ∥ai ) = j=1 rij log aij is zeroavoiding with respect to ri , resulting in mode-covering behavior. As illustrated in Fig. 1, it asymmetrically penalizes the model most if aij is small when rij is large. See Murphy [2012] for more details. PSi aij • Reverse KL (mode-seeking). The reverse KL DKL (ai ∥ri ) = j=1 is zero-forcing aij log rij with respect to ai , resulting in mode-seeking behavior. As illustrated in Fig. 1, it asymmetrically penalizes the model most if aij is large when rij is small. See Murphy [2012] for more details. 3.2
Normal Guidance
In many correlated MIL tasks, instances are spatially or temporally dependent. For example, in medical imaging, if a tumor is visible in one slice, it is likely to be present in adjacent slices. To incorporate this inductive bias into the guided attention framework, we propose Normal Guidance. First, we compute the empirical mean and variance of the attended slice position under the learned distribution ai : P P 2 PSi Si Si 2 E[J] = j=1 j · aij , Var(J) = j · a − j · a . (2) ij ij j=1 j=1 Next, these values define a Normal PDF that we evaluate at integers j from 1 to Si and renormalize to obtain the reference distribution ri : r̂ij = NormPDF(j|E[J], Var(J)),
r̂ij rij = PSi
k=1 r̂ik
.
(3)
When performing gradient descent on the loss in Eq. (1), we use the stop-gradient operation on ri to require updates to ai that improve agreement with ri . We recompute ri before each backward pass. 3.3
Multi-Head Normal Guidance
While Normal Guidance incorporates local dependencies, it encourages learning unimodal attention weights. This can be restrictive in correlated MIL settings where multiple disjoint regions are relevant (e.g., multiple tumors across non-adjacent slices). In such cases, a single normal distribution may over-smooth across regions or focus on only one region. To overcome this unimodal restriction, we extend our method to Multi-Head Normal Guidance. For multi-head self-attention, we use a separate unimodal reference distribution via Eq. (3) for each head of the final self-attention layer (see Fig. 2). The overall regularization term is computed by averaging the divergence across all heads. Because each head is guided toward a separate bell-shaped curve, 4
the combined effect can be intuitively thought of as guided by a Gaussian mixture model. This extension allows the network to simultaneously attend to multiple spatially separate regions.
4
Experiments
4.1
Datasets
Head CT. The RSNA 2019 Brain CT Hemorrhage Challenge [Flanders et al., 2020] includes 752,803 slices from 21,744 CT scans (20-60 slices per scan) with labels for presence/absence of any intracranial hemorrhage (ICH). We partition the released training set (no other release has labels) into our own training, validation, and test sets as described below. Castro-Macı́as et al. [2024] train and test on a smaller subset of this data (1,150 scans total). Chest CT. The RSNA Pulmonary Embolism CT dataset [Colak et al., 2021] has 1,790,594 slices from 7,279 CT scans (63-1,083 slices per scan) with labels for presence/absence of any pulmonary embolism. We partition the released training set (no other release has labels) as described below. Abdomen CT. The RSNA 2023 Abdominal Trauma AI Challenge [Hermans et al., 2024] includes 1,500,653 slices from 4,711 CT scans (41-1,727 slices per scan) with labels for presence/absence of any abdominal trauma (AT). We partition the released training set (no other release has labels) as described below. Partition by patient. For each medical dataset (Head CT, Chest CT, and Abdomen CT), we randomly assign 3D scans into training, validation, and testing sets using a 4:1:1 ratio, ensuring each patient’s data belongs to exactly one set to avoid leakage. We stratify by class to ensure comparable class frequencies. We repeat this process across three random seeds, so every learning algorithm is evaluated on the same 3 distinct partitions into training, validation, and test sets. Expert annotations. All 3 CT scan datasets include a binary label for each whole CT scan and an instance-level binary label for each 2D axial slice in every scan. The Head CT dataset was annotated by 60 neuroradiologists; the Chest CT dataset was annotated by a group of more than 80 expert thoracic radiologists; and the Abdomen CT dataset was annotated by radiologists with subspecialty training and/or professional experience in abdominal trauma. In our work, slice-level labels are only used to evaluate localization and establish best-in-class performance ceilings. Semi-synthetic. The Shifted Mean MIL dataset [Harvey et al., 2025, 2026] was created to assess generalization capabilities of MIL methods. We follow prior work’s generative process: we draw yi ∼ Bern(0.5) so 50% of bags are positive and Si ∼ Unif({20, . . . , 60}) so bags have between 20 and 60 instances. If a bag is negative, all M features of all Si instances are drawn from a mean 0, variance 1 Gaussian. If a bag is positive, we select a contiguous block of R = 12 instances and draw only the first of M features in this block from a Gaussian with mean ∆ = 0.5; all other features use a mean 0, variance 1 Gaussian. We call this semi-synthetic because the range of Si and the value of R are set to match the Head CT dataset, and the number of features M = 768 is set to match the ViT embedding size we use. We sample 10,000 bags for training, 2,500 bags for validation, and 1,000 bags for testing. We repeat this process with three random seeds; each seed samples a different training, validation, and test set. Further details are in App. E.
1
Attention weight aij
0.12 0.10 0.08
5
9 14 18 22 27 31 35 40
Multi-Head Normal Guidance
Head 1 Head 2 Head 3 Head 4
Head 5 Head 6 Head 7 Head 8
10
20 Slice index j
1
5
9 14 18 22 27 31 35 40
Attention Weights after Training TransMIL Normal Guidance
0.06 0.04 0.02 0.00
0
30
40
0
10
20 Slice index j
30
40
Figure 2: Multi-Head Normal Guidance. Left: discretized normal distribution for each attention head in TransMIL (we set H = 8 throughout). Right: Mean attention weight aij for each head after training, for standard TransMIL and TransMIL plus Multi-Head Normal Guidance. 5
Probability density of positive slices
Semi-Synthetic
Head
Chest
Abdomen
2.5 2.0 1.5 1.0 0.5 0.0
0.0
0.2
0.4 0.6 0.8 Normalized slice index Sj i
1.0
0.0
0.2
0.4 0.6 0.8 Normalized slice index Sj i
1.0
0.0
0.2
0.4 0.6 0.8 Normalized slice index Sj i
1.0
0.0
0.2
0.4 0.6 0.8 Normalized slice index Sj i
1.0
Figure 3: Probability density of positive slices in training set. Because slices in the middle of the scan are more likely to be positive the centered Gaussian performs well for localization.
4.2
Experimental Goals, Metrics, Baselines, and Training Details
We are primarily interested in assessing the localization performance of MIL methods trained with weak supervision. Following past work [Castro-Macı́as et al., 2024], we measure this performance only on positive-labeled bags in the test set. For each bag i such that yi = 1, at all Si slices indexed by j we assess the predicted attention value aij versus the ground truth label yij via the AUROC, using macro averaging over bags. While our focus is localization, we also wish to verify that scanlevel classification remains strong. We thus report scan-level (“bag”-level) classification quality in terms of AUROC. For both localization and bag-level classification, AUPRC results are in App. B. The primary reported number (AUROC or AUPRC) is the average over the 3 test partitions of each task. The number after the “±“ symbol indicates the empirical standard deviation of the 3 values. Baselines. Among MIL methods, we consider max pooling MIL, mean pooling MIL, attentionbased MIL (ABMIL; Ilse et al., 2018), transformer-based MIL (TransMIL; Shao et al., 2021), and the Smooth Operator (known as SmAP and SmTP from Castro-Macı́as et al. [2024]). We chose not to compare to CAMIL [Lu et al., 2021], DSMIL [Li et al., 2021], and CLAM [Fourkioti et al., 2024] because of their poor localization results on the Head CT dataset reported in Castro-Macı́as et al., where these methods were outperformed by the Smooth Operator in all tested localization tasks. For localization only, we also compare to the center-focused Gaussian baseline from Harvey et al.. This method predicts a centered bell-shaped curve across slices for every scan. It is not informed at all by image content and cannot produce scan-level predictions. See App. D.2 for details about localization experiments. For all experiments with Normal Guidance, for simplicity we set D to the forward KL and fix regularization strength λ = 1; we conduct a sensitivity analysis of these choices later in Fig. 4. We configure baseline methods by following past work, providing reproducible details in App. D. Encoder. For all medical imaging experiments, we use a frozen ViT-B/16 [Dosovitskiy et al., 2021] pre-trained on ImageNet [Deng et al., 2009] as the raw feature encoder, with embedding size M = 768. Past work [Harvey et al., 2026] found MIL with a frozen ViT encoder to be competitive with other encoders like ConvNeXt or MedSAM. Their frozen ViT with MIL pipeline was also competitive with fully fine-tuned 3D CNNs while being far more affordable to train. For all methods and tasks, we feed a final scan-level embedding to a linear binary classifier head. Training. All methods are optimized using mini-batch stochastic gradient descent (SGD) with a momentum parameter of 0.9 and batch size of 64. We train for 1,000 epochs and use bag-level validation AUROC for early stopping. For all methods, we grid search two key hyperparameters: learning rate in {0.1, 0.01, 0.001, 0.0001} and L1 regularization strength for all pooling and classifier weights in {1.0, 0.1, 0.01, 0.001, 0.0001, 1e-5, 1e-6, 0.0}. 4.3
Best-in-Class Ceiling for Instance-Level Localization
On any benchmark, it is important to understand how close the current methods are to “saturation”, meaning the best possible performance on that task by a model in that class given the designated inputs. The key assumptions of our approach are using a linear classification head on frozen ViT embeddings. Here we try to establish a “best-in-class ceiling” for each task that represents a realistic upper bound for any correlated MIL method using our assumptions. This is not merely a classifier that handles each instance separately with a “one feature in, one prediction out” architecture, used as an upper bound in some past MIL works [Correia et al., 2020, Guillaumin et al., 2010]. Instead, our 6
instance-level ceiling for CT consumes a local neighborhood of embeddings to account for spatial dependencies in the correlated MIL setting. Ceiling for CT tasks. We train an instance-level classifier using ground truth instance-level labels yij . This classifier takes as input for instance index j a local window of R consecutive embedding vectors centered on hij and outputs an instance-level predicted probability ŷij . We use a 1D convolutional layer with kernel size R over instance representations provided by the frozen ViT, followed by a linear classifier head. For each dataset, we set R equal to the mean number of consecutive positive instances, allowing gaps of up to 3 negative instances in a “consecutive” block. This sets R = 12 for Head CT, R = 35 for Chest CT, and R = 29 for Abdomen, according to Tab. A.2. Ceiling for semi-synthetic. For Shifted Mean MIL data, we derive the probability p(yij =1|hi , Si , yi =1): the chance that instance j comes from the Gaussian with shifted mean given its features and the fact that its bag is positive. This event yij =1 is equivalent to index j falling in the chosen latent segment of R consecutive instances. The generative model is framed in terms of the starting index ui of this segment, so let Uij be the set of possible values of ui that could make j positive. Adding up the probability of these relevant values, we can write p(yij =1|hi , Si , yi =1) as X X p(hi |ui , Si , yi =1)p(ui |Si , yi =1) = p(ui |hi , Si , yi =1) = . (4) p(hi |Si , yi =1) ui ∈Uij
ui ∈Uij
This uses Bayes’ rule. Each term in Eq (4) can be computed in closed-form, as explained in App. E. This oracle upper bound for instance-level performance has not been derived in prior work. 4.4
Best-in-Class Ceiling for Bag-Level Classification
Bag ceiling for CT tasks. We establish an upper bound for bag-level performance of MIL by using oracle instance-level labels yij to pool embeddings, instead of learned attention. This of course has an unfair advantage compared to weakly-supervised MIL since it relies on instance-level annotations even at test time and not just images alone. However, our intention is to establish a ceiling for imageonly methods, since we hope the ideal attention matches expert labels. Procedurally, for each positive bag, we pool instance-level embeddings using a uniform distribution over that bag’s positively labeled instances. For each negative bag, we randomly sample a positive bag and use its instance-level label distribution, adjusting for differing slice counts via interpolation. We then pool the negative bag’s embeddings uniformly over all “positive” instances. This enforces similar ROI distributions regardless of class to create scan-level representations, which are used to train a scan-level linear binary classifier. As an alternative to enforcing similar ROIs, we tried averaging over all negative bag slices uniformly but found unintended shortcuts were possible due position-specific feature distributions. If negative bags pool over the entire volume but positive bags concentrate on specific regions, a simple “how dark is this image?” feature could detect the more dark slices at index 1 or Si in negative bags and yield near-perfect classification that is too optimistic. Bag ceiling for semi-synthetic. Since the true data-generating process is known, we use the Bayes estimator derived in Eq. (10) of Harvey et al. [2026] as an upper bound on bag-level performance.
5
Results
5.1
Localization performance
Localization results are in Tab. 1 and B.1. The major findings from these results are listed below. The Centered Gaussian baseline outperforms existing MIL for localization. Tab. 1 and B.1 show that the simple baseline from Harvey et al. [2026] not only outperforms the attention mechanism for localization on the Head CT task but also on the thoracic and abdominal tasks. On Chest CT, no existing MIL scores above 0.69 while the Centered baseline scores 0.78! Normal Guidance can outperform this baseline: a naive center bias is not always best. The centered Gaussian aij ∝ NormPDF(j| S2i , 1) encodes the inductive bias that positive instances are most likely in the center of the bag (see probability density of positive slices in Fig. 3). Yet our Normal Guidance (NG) and Multi-Head Normal Guidance (MHNG) deliver notable gains in AUROC over this baseline of at least +0.019 on Head CT, +0.08 on Chest CT, and +0.06 on Abdomen CT. This 7
reassuringly indicates that learning attention from image features can outperform the image-agnostic Centered Gaussian baseline. 5.2
Bag-level classification performance
Bag-level (scan-level) classification is reported in Tab. 2 and B.2. Recall that Normal Guidance is specifically designed to improve localization. Here our goal is to verify how this regularization impacts bag-level classification and interpret how methods compare to our new ceilings. Normal Guidance keeps scan-level classification competitive. Across all tested CT tasks, our NG and MHNG score better than or within tolerance (within 0.01 AUROC) of their corresponding unregularized or smoothed alternative. On semi-synthetic bag-level classification, ABMIL + Smooth Operator scores highest (0.771) but TransMIL with MHNG is competitive at 0.766. On two CT tasks, the best MIL methods are near the bag-level ceiling. The best MIL method achieves very good results on Head CT (within 0.002 AUROC of best-in-class ceiling) and good results on Chest CT (within 0.022 AUROC of the best-in-class ceiling). For the semi-synthetic dataset, the best MIL method is within 0.042 AUROC of its ceiling. For the Abdomen CT task, the gap of about 0.1 AUROC may be real. Yet we guess the ceiling is overestimated; this task has much more variability in the organs covered than other tasks, and 6.6% of scans do not fully cover one or more organs. Using instance-level labels from an incomplete positive scan to pool a complete negative scan (or vice versa) could create shortcut features that overestimate performance. Transformer-based MIL only slightly improves bag-level performance. Transformer-based methods improve AUROC by up to 0.05 on the Abdomen CT task over ABMIL-based counterparts, but otherwise don’t seem to add appreciable gain on the other 3 tasks. 5.3
Sensitivity Analysis
Our Normal Guidance procedure’s success could depend on the choice of divergence D and the regularization strength λ. Thus far, these values have been fixed as forward KL and λ = 1. The two panels of Fig. 4 show how performance on the Head CT task depends on these variables. First, the left panel of Fig. 4 suggests a modest preference for Forward KL for both localization and wholescan classification. The right panel of Fig. 4 suggests that λ = 100 = 1 is a reasonable choice for localization, but we might do slightly better by tuning λ. Thus, the localization success of NG and MHNG reported in Sec. 5.1 might be improvable. 5.4
Qualitative Side-by-Side Comparisons of Learned Attention
See App. C for visual comparisons of the learned attention distributions from different MIL methods on a handful of randomly-selected CT scans from each of our 3 datasets. In each figure, we either compare ABMIL, ABMIL with Sm, and ABMIL with NG; or TransMIL, TransMIL with Sm, and TransMIL with MHNG. As a primary takeaway, we can visually confirm that Normal Guidance often produces more focused and unimodal attention distributions than its counterparts. Some plots,
Table 1: Localization AUROC on test set, aka slice-level or instance-level classification, using positive bags only. Takeaway: On real CT images, Normal Guidance outperforms all other MIL and is only MIL to significantly beat the image-free Centered Gaussian baseline. Method Centered Gaussian Max pooling Mean pooling ABMIL Smooth Operator Normal Guidance (ours) TransMIL Smooth Operator Multi-Head Normal Guidance (ours) Best-in-Class Ceiling (see Sec. 4.3)
Semi-Synthetic 0.689±0.015 0.504±0.002 0.500±0.000 0.546±0.054 0.693±0.030 0.541±0.055 0.593±0.131 0.592±0.127 0.706±0.076 0.884±0.004
8
Head CT 0.850±0.001 0.286±0.004 0.500±0.000 0.736±0.052 0.831±0.003 0.871±0.002 0.792±0.031 0.823±0.019 0.869±0.002 0.896±0.003
Chest CT 0.780±0.002 0.461±0.003 0.500±0.000 0.558±0.086 0.688±0.043 0.866±0.005 0.585±0.077 0.654±0.035 0.869±0.002 0.871±0.004
Abdomen CT 0.573±0.031 0.453±0.011 0.500±0.000 0.526±0.041 0.564±0.012 0.663±0.006 0.579±0.003 0.582±0.043 0.634±0.007 0.717±0.011
Table 2: Bag-level AUROC on test set for whole-scan classification. Takeaway: Normal Guidance is competitive with state-of-the-art MIL alternatives here, while improving localization (Tab. 1). Method Max pooling Mean pooling ABMIL Smooth Operator Normal Guidance (ours) TransMIL Smooth Operator Multi-Head Normal Guidance (ours) Best-in-Class Ceiling (see Sec. 4.4)
Semi-Synthetic 0.624±0.009 0.752±0.002 0.751±0.003 0.771±0.008 0.751±0.003 0.763±0.016 0.768±0.025 0.766±0.006 0.810±0.006
Head CT 0.888±0.009 0.920±0.012 0.919±0.009 0.925±0.012 0.925±0.010 0.925±0.014 0.926±0.013 0.926±0.012 0.927±0.012
Chest CT 0.656±0.004 0.669±0.018 0.664±0.017 0.671±0.023 0.678±0.026 0.663±0.023 0.670±0.015 0.662±0.015 0.700±0.014
Abdomen CT 0.643±0.029 0.625±0.027 0.639±0.044 0.648±0.043 0.651±0.026 0.684±0.047 0.664±0.050 0.677±0.029 0.776±0.033
Localization 0.834±0.007 0.871±0.002 0.869±0.003
Classification 0.924±0.012 0.925±0.010 0.919±0.011
Test AUROC
0.87
Divergence Squared error Forward KL Reverse KL
0.86 0.85 ABMIL with Normal Guidance 10 1
100 Regularization strength
101
Figure 4: Sensitivity analysis. Left: Effect of divergence D on localization and bag-level classification, with λ = 100 . Right: Effect of regularization strength λ on localization (D: Forward KL). All results here report test set AUROC on the Head CT task, using ABMIL with Normal Guidance.
especially the right panels of Fig. C.3, show cases where all methods fail to capture the key block of positive slices, indicating plenty of room for further localization method development.
6
Conclusion
The attention mechanism was originally proposed in MIL as a method that incorporates interpretability [Ilse et al., 2018]. Yet this work shows that for localization in 3D medical images, MIL alone cannot outperform a simple baseline due to a lack of inductive bias. We proposed Normal Guidance, a regularization technique for attention-based MIL models that encourages attention to be bell-curve shaped. We investigated three CT datasets totaling over 4 million 2D slices and show that Normal Guidance outperforms all other attention-based MIL methods as well as the simple slicelevel baselines, achieving a new state-of-the-art for instance-level MIL results on all datasets while maintaining high bag-level performance. Limitations. (1) We only use positive bags to evaluate slice-level localization, following past work [Castro-Macı́as et al., 2024]. We do not evaluate the quality of attention in negative bags. This highlights an open problem in the literature: defining what attention should look like for a negative bag. (2) Jain and Wallace [2019] argue that attention weights should not be treated as faithful explanations of classifier decisions, since permuting attention can yield similar predictions. Our results show that Normal Guidance can help learned attention track expert annotations more closely, but this does not by itself enforce that the attended slices are causally responsible for the bag-level prediction. (3) We focus on deep MIL with frozen encoders and linear classifiers, again following past work [Lu et al., 2021, Shao et al., 2021, Fourkioti et al., 2024, Castro-Macı́as et al., 2024]. Parameter-efficient fine-tuning could deliver more effective encoders for this domain, but at prohibitive computational cost. Even with frozen encoders, model training and hyperparameter optimization takes roughly 9 hours for each TransMIL run on the head CT dataset (not to mention repeating this across 32 hyperparameter configurations and 3 train/test partitions). Future work. Beyond 3D images, MIL is often used for whole slide images. Alternative inductive biases that account for the 2D spatial dependencies in WSI are critical for understanding localization results in that domain. A second direction is conditioning ri on the bag label yi : the right reference 9
(Gaussian, Uniform, etc.) for negative bags is worth investigating. Finally, the same bell-shaped prior (NG) may also suit MIL tasks beyond medical imaging, such as video anomaly detection, where evidence is typically contiguous in time [Sultani et al., 2018].
Acknowledgments This work is supported by the U.S. National Institutes of Health (grant # R01NS134859) and the Alzheimer’s Drug Discovery Foundation. Author MCH is also supported in part by the U.S. National Science Foundation (NSF) via IIS CAREER grant # 2338962. We are grateful for resources and support from the Tufts High-Performance Computing Cluster. This paper’s content is solely the responsibility of the authors and does not necessarily represent the official views of the NIH or NSF.
References Hassan Baker and Austin J. Brockmeier. Patch2Loc: Learning to Localize Patches for Unsupervised Brain Lesion Detection. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2026. Finn Behrendt, Debayan Bhattacharya, Julia Krüger, Roland Opfer, and Alexander Schlaefer. Patched Diffusion Models for Unsupervised Anomaly Detection in Brain MRI. In Medical Imaging with Deep Learning, 2024. Francisco M. Castro-Macı́as, Pablo Morales-Álvarez, Yunan Wu, Rafael Molina, and Aggelos K. Katsaggelos. Sm: enhanced localization in Multiple Instance Learning for medical imaging classification. In Advances in Neural Information Processing Systems (NeurIPS), 2024. Errol Colak, Felipe C. Kitamura, Stephen B. Hobbs, Carol C. Wu, Matthew P. Lungren, Luciano M. Prevedello, Jayashree Kalpathy-Cramer, Robyn L. Ball, George Shih, Anouk Stein, Safwan S. Halabi, Emre Altinmakas, Meng Law, Parveen Kumar, Karam A. Manzalawi, Dennis Charles Nelson Rubio, Jacob W. Sechrist, Pauline Germaine, Eva Castro Lopez, Tomas Amerio, Pushpender Gupta, Manoj Jain, Fernando U. Kay, Cheng Ting Lin, Saugata Sen, Jonathan Wesley Revels, Carola C. Brussaard, and John Mongan For the RSNA-STR Annotators and Dataset Curation Contributors. The RSNA pulmonary embolism CT dataset. Radiology: Artificial Intelligence, 3 (2):e200254, 2021. Joana Correia, Isabel Trancoso, and Bhiksha Raj. Automatic in-the-wild dataset annotation with deep generalized multiple instance learning. In Proceedings of the Twelfth Language Resources and Evaluation Conference, 2020. Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2009. Thomas G. Dietterich, Richard H. Lathrop, and Tomás Lozano-Pérez. Solving the multiple instance problem with axis-parallel rectangles. Artificial Intelligence, 89(1-2):31–71, 1997. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations (ICLR), 2021. Ji Feng and Zhi-Hua Zhou. Deep MIML Network. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2017. Adam E. Flanders, Luciano M. Prevedello, George Shih, Safwan S. Halabi, Jayashree KalpathyCramer, Robyn Ball, John T. Mongan, Anouk Stein, Felipe C. Kitamura, Matthew P. Lungren, Geetika Choudhary, Luciano Cala, Luı́s Coelho, Mads Mogensen, Fátima Morón, Eric Miller, Ichiro Ikuta, Vahe Zohrabian, Oran McDonnell, Christoph Lincoln, Luciano Shah, Devon Joyner, Ashish Agarwal, Richard K. Lee, and Jayashree Nath. Construction of a Machine Learning Dataset through Collaboration: The RSNA 2019 Brain CT Hemorrhage Challenge. Radiology: Artificial Intelligence, 2(3), 2020. 10
Olga Fourkioti, Matt De Vries, and Chris Bakal. CAMIL: Context-Aware Multiple Instance Learning for Cancer Detection and Subtyping in Whole Slide Images. In International Conference on Learning Representations (ICLR), 2024. Matthieu Guillaumin, Jakob Verbeek, and Cordelia Schmid. Multiple Instance Metric Learning from Automatically Labeled Bags of Faces. In Proceedings of the European Conference on Computer Vision (ECCV), 2010. Ethan Harvey, Dennis Johan Loevlie, and Michael C. Hughes. Synthetic Data Reveals Generalization Gaps in Correlated Multiple Instance Learning. arXiv preprint arXiv:2510.25759, 2025. Ethan Harvey, Dennis Johan Loevlie, Amir Ali Satani, Wansu Chen, David M. Kent, and Michael C. Hughes. A Multi-Dataset Benchmark of Multiple Instance Learning for 3D Neuroimage Classification. In Conference on Health, Inference, and Learning (CHIL), 2026. Sebastiaan Hermans, Zixuan Hu, Robyn L. Ball, Hui Ming Lin, Luciano M. Prevedello, Ferco H. Berger, Ibrahim Yusuf, Jeffrey D. Rudie, Maryam Vazirabad, Adam E. Flanders, George Shih, John Mongan, Savvas Nicolaou, Brett S. Marinelli, Melissa A. Davis, Kirti Magudia, Ervin Sejdić, and Errol Colak. RSNA 2023 Abdominal Trauma AI Challenge: Review and Outcomes. Radiology: Artificial Intelligence, 7(1):e240334, 2024. Zhe Huang, Benjamin S. Wessler, and Michael C. Hughes. Detecting Heart Disease from MultiView Ultrasound Images via Supervised Attention Multiple Instance Learning. In Machine Learning for Healthcare Conference (MLHC), 2023. Maximilian Ilse, Jakub Tomczak, and Max Welling. Attention-based Deep Multiple Instance Learning. In International Conference on Machine Learning (ICML), 2018. Sarthak Jain and Byron C. Wallace. Attention is not Explanation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), 2019. Hassan Keshvarikhojasteh, Josien P. W. Pluim, and Mitko Veta. Multi-head Attention-based Deep Multiple Instance Learning. In Proceedings of the MICCAI Workshop on Computational Pathology, 2024. Bin Li, Yin Li, and Kevin W. Eliceiri. Dual-stream Multiple Instance Learning Network for Whole Slide Image Classification with Self-supervised Contrastive Learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021. Kechun Liu, Wenjun Wu, Joann G. Elmore, and Linda G. Shapiro. Semantics-Aware Attention Guidance for Diagnosing Whole Slide Images. In International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2024. Ming Y. Lu, Drew F. K. Williamson, Tiffany Y. Chen, Richard J. Chen, Matteo Barbieri, and Faisal Mahmood. Data efficient and weakly supervised computational pathology on whole slide images. Nature Biomedical Engineering, 5(6):555–570, 2021. Oded Maron and Tomás Lozano-Pérez. A Framework for Multiple-Instance Learning. In Advances in Neural Information Processing Systems (NeurIPS), 1997. Kevin P. Murphy. Machine Learning: A Probabilistic Perspective, chapter 21.2.2 Forward or reverse KL? MIT Press, 2012. John Muschelli. Recommendations for processing head CT data. Frontiers in Neuroinformatics, 13: 61, 2019. Sharon Peled, Yosef E Maruvka, and Moti Freiman. PSA-MIL: A Probabilistic Spatial AttentionBased Multiple Instance Learning for Whole Slide Image Classification. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), 2026. Pedro O. Pinheiro and Ronan Collobert. From Image-Level to Pixel-Level Labeling With Convolutional Networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2015. 11
Gwenolé Quellec, Guy Cazuguel, Béatrice Cochener, and Mathieu Lamard. Multiple-Instance Learning for Medical Image and Video Analysis. IEEE Reviews in Biomedical Engineering, 10:213–234, 2017. Edward Raff and James Holt. Reproducibility in Multiple Instance Learning: A Case For Algorithmic Unit Tests. In Advances in Neural Information Processing Systems (NeurIPS), 2023. Khaled Saab, Jared Dunnmon, Roger Goldman, Alex Ratner, Hersh Sagreiya, Christopher Ré, and Daniel Rubin. Doubly Weak Supervision of Deep Learning Models for Head CT. In International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2019. Zhuchen Shao, Hao Bian, Yang Chen, Yifeng Wang, Jian Zhang, Xiangyang Ji, and Yongbing Zhang. TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification. In Advances in Neural Information Processing Systems (NeurIPS), 2021. Yash Sharma, Aman Shrivastava, Lubaina Ehsan, Christopher A. Moskaluk, Sana Syed, and Donald Brown. Cluster-to-Conquer: A Framework for End-to-End Multi-Instance Learning for Whole Slide Image Classification. In Medical Imaging with Deep Learning, 2021. Waqas Sultani, Chen Chen, and Mubarak Shah. Real-World Anomaly Detection in Surveillance Videos. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018. Xiaosong Wang, Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald M. Summers. ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017. Yunan Wu, Arne Schmidt, Enrique Hernández-Sánchez, Rafael Molina, and Aggelos K. Katsaggelos. Combining Attention-Based Multiple Instance Learning and Gaussian Processes for CT Hemorrhage Detection. In International Conference on Medical Image Computing and Computer Assisted Intervention (MICCAI), 2021. Shaolei Zhang and Yang Feng. Modeling Concentrated Cross-Attention for Neural Machine Translation with Gaussian Mixture Model. In Findings of the Association for Computational Linguistics: EMNLP 2021, 2021. Yunlong Zhang, Honglin Li, Yuxuan Sun, Zhongyi Shui, Jingxiong Li, Chenglu Zhu, and Lin Yang. AEM: Attention Entropy Maximization for Multiple Instance Learning based Whole Slide Image Classification. In International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2025. Wentao Zhu, Qi Lou, Yeeleng Scott Vang, and Xiaohui Xie. Deep Multi-instance Networks with Sparse Label Assignment for Whole Mammogram Classification. In International Conference on Medical Image Computing and Computer Assisted Intervention (MICCAI), 2017.
12
A
Dataset Details
A.1
Preprocessing of CT scans
For all medical imaging datasets (Head CT, Chest CT, and Abdomen CT), we convert images into Hounsfield Units (HU) using each image’s rescale slope and intercept; resize each 2D slice to 224 × 224 pixels; and normalize images with the training set mean and standard deviation of each channel. For the Head CT dataset, we use intensity windowing to exclude the skull, other bones, and calcifications (only including -100 to 300 HU) [Muschelli, 2019]. A.2
Descriptive Statistics
Below, we include three tables that summarize the CT datasets: • the number of blocks (ROIs) in a bag, • the number of instances in a typical block (ROI), and • the fraction of all instances in the bag contained in a typical block (ROI). 1 block 2 blocks 3 blocks 4 blocks 5 blocks 6 blocks 7 blocks Head CT 8296/8882 569/8882 17/8882 – – – – Chest CT 1437/2211 497/2211 185/2211 65/2211 16/2211 7/2211 4/2211 Abdomen CT 304/365 53/365 7/365 1/365 – – – Table A.1: Number of blocks of consecutive positive instances in a bag, allowing gaps of up to 3 negative instances in a “consecutive” block.
Mean Min 5th 25th 50th 75th 95th Max Head CT 12 1 2 6 11 17 23 43 Chest CT 35 1 1 7 16 44 137 353 Abdomen CT 29 1 2 6 15 40 95 242 Table A.2: Number of positive instances in a block. We allow gaps of up to 3 negative instances in a “consecutive” block.
Mean Min 5th 25th 50th 75th 95th Max Head CT 0.350 0.017 0.053 0.172 0.342 0.513 0.688 1.000 Chest CT 0.093 0.001 0.004 0.023 0.059 0.137 0.283 0.499 Abdomen CT 0.119 0.002 0.006 0.027 0.063 0.174 0.401 0.870 Table A.3: Fraction of all instances in a block (region of interest or ROI). We allow gaps of up to 3 negative instances in a “consecutive” block.
13
B
AUPRC Results
Table B.1: Localization results for the AUPRC metric, using the test set of each task. Corresponding AUROC results are in Tab. 1. Method Centered Gaussian Max pooling Mean pooling ABMIL Smooth Operator Normal Guidance (ours) TransMIL Smooth Operator Multi-Head Normal Guidance (ours) Best-in-Class Ceiling (see Sec. 4.3)
Semi-Synthetic 0.569±0.012 0.384±0.003 0.332±0.003 0.424±0.045 0.574±0.035 0.403±0.062 0.487±0.140 0.467±0.147 0.578±0.086 0.825±0.005
Head CT 0.710±0.007 0.302±0.007 0.359±0.009 0.634±0.049 0.713±0.007 0.744±0.010 0.659±0.032 0.689±0.007 0.740±0.010 0.781±0.010
Chest CT 0.481±0.008 0.185±0.002 0.188±0.001 0.233±0.059 0.344±0.037 0.522±0.008 0.258±0.052 0.309±0.018 0.535±0.005 0.528±0.004
Abdomen CT 0.222±0.038 0.109±0.002 0.104±0.007 0.159±0.027 0.167±0.020 0.248±0.020 0.191±0.021 0.195±0.039 0.227±0.010 0.285±0.030
Table B.2: Bag-level AUPRC on test set for whole-scan classification. Takeaway: Normal Guidance is competitive with state-of-the-art MIL alternatives here, while improving localization (Tab. B.1). Corresponding AUROC results are in Tab. 2.
Divergence Squared error Forward KL Reverse KL
Localization 0.702±0.011 0.744±0.010 0.747±0.012
Semi-Synthetic 0.616±0.005 0.756±0.007 0.755±0.009 0.775±0.010 0.755±0.010 0.766±0.019 0.770±0.025 0.766±0.015 0.814±0.002
Classification 0.910±0.010 0.911±0.008 0.905±0.011
Head CT 0.859±0.013 0.903±0.011 0.903±0.007 0.910±0.011 0.911±0.008 0.912±0.012 0.913±0.012 0.912±0.012 0.913±0.010
Test AUPRC
Method Max pooling Mean pooling ABMIL Smooth Operator Normal Guidance (ours) TransMIL Smooth Operator Multi-Head Normal Guidance (ours) Best-in-Class Ceiling (see Sec. 4.4)
Chest CT 0.440±0.002 0.455±0.005 0.451±0.005 0.465±0.008 0.477±0.015 0.447±0.013 0.461±0.009 0.451±0.010 0.468±0.004
0.77 0.76 0.75 0.74 0.73 0.72 0.71
Abdomen CT 0.133±0.016 0.117±0.010 0.122±0.014 0.139±0.014 0.117±0.008 0.139±0.025 0.135±0.007 0.139±0.012 0.252±0.081
ABMIL with Normal Guidance 10 1
100 Regularization strength
101
Figure B.1: Sensitivity analysis for the AUPRC metric on the test set. Left: Effect of divergence D on localization and bag-level classification, with λ = 100 . Right: Effect of regularization strength λ on localization (D: Forward KL). All results from the Head CT task, using ABMIL with Normal Guidance.
14
C
Attention Visualizations
Below, you can see visuals of the learned attention weights for key MIL methods on the Head, Chest, and Abdomen CT tasks.
1 0.12
5
4
10 15 20 Slice index j
0.08 0.04 0.02 0
0.08
0.15 0.10 0.05
5
10 15 20 Slice index j
0
1
5
4
10 15 20 Slice index j
25
0.10 0.05 0.00
25
0
5
10 15 20 Slice index j
8 11 15 18 22 25 29 33
0.25
0
1
25
4
0.30
0.02
0.07 0.06 0.05 0.04 0.03 0.02 0.01 0.00
1
8 12 16 20 24 28 32 36
0.04
7 10 13 16 19 22 25 28
0.15
4
0.06
0.00
0.20
0.06
0.00
0.20
7 10 13 16 19 22 25 28 Ground truth TransMIL Smooth Operator Normal Guidance
0.10
1 0.10
0.00
25
7 10 13 16 19 22 25 28
Attention weight aij
0.05
4
0.25
10
4
20 Slice index j
0.20 0.15 0.10 0.05 0.00
30
0
10
1
8 12 16 20 24 28 32 36
4
20 Slice index j
30
8 11 15 18 22 25 29 33
0.30 0.25 Attention weight aij
Attention weight aij
0.10
0
1
Attention weight aij
Attention weight aij
0.15
0.00
Attention weight aij
7 10 13 16 19 22 25 28 Ground truth ABMIL Smooth Operator Normal Guidance
Attention weight aij
4
Attention weight aij
1 0.20
0
10
20 Slice index j
0.20 0.15 0.10 0.05 0.00
30
0
10
20 Slice index j
30
Figure C.1: Learned attention weights for 4 example CT scans (columns) from the Head CT dataset. Top row: methods based on ABMIL. Bottom row: transformer-based MIL methods.
0.01 0
50
100 150 Slice index j
0.015 0.010 0.005 0.000
200
1 22 44 66 88 109 131 153 175 197 Ground truth TransMIL Smooth Operator Normal Guidance
0.03 0.02 0.01 0.00
0
50
100 150 Slice index j
0
50
100 150 Slice index j
200
0.015 0.010 0.005 0.000
200
1 24 48 72 95 119 143 166 190 214
50
100 150 Slice index j
0.010 0.005 0
50
100 150 Slice index j
0.02 0.01 0.00
200
200
50
100 150 200 Slice index j
250
0.020
0.015 0.010 0.005 0.000
0
1 29 58 87 116 145 174 203 232 261
0.020
0.015
0.000
0
0.03
1 26 51 77 102 128 153 179 204 230
0.020 Attention weight aij
Attention weight aij
0.04
1 29 58 87 116 145 174 203 232 261
0.020 Attention weight aij
0.02
0.00
1 26 51 77 102 128 153 179 204 230
Attention weight aij
0.03
Attention weight aij
Attention weight aij
0.04
1 24 48 72 95 119 143 166 190 214 0.020
Attention weight aij
Ground truth ABMIL Smooth Operator Normal Guidance
Attention weight aij
1 22 44 66 88 109 131 153 175 197
0
50
100 150 Slice index j
200
0.015 0.010 0.005 0.000
0
50
100 150 200 Slice index j
250
Figure C.2: Learned attention weights for 4 example CT scans (columns) from the Chest CT dataset. Top row: methods based on ABMIL. Bottom row: transformer-based MIL methods.
15
0.010 0.005 0.000
0
50
100 150 Slice index j
0.025
0.08
0.020 0.015 0.010 0.005 0.000
200
1 25 49 73 97 121 145 169 193 217 0.030 Attention weight aij
0.025 0.020 0.015 0.010 0.005 0.000
0
50
100 150 Slice index j
200
50
100 Slice index j
0.06 0.04 0.02 0.00
150
1 20 39 59 78 98 117 137 156 176
Attention weight aij
Ground truth TransMIL Smooth Operator Normal Guidance
0
0.10
0.025
0.08
0.015 0.010 0.005 0.000
0
50
100 Slice index j
0
50
100 Slice index j
150
0.07 0.06 0.05 0.04 0.03 0.02 0.01 0.00
1 22 43 64 85 106 127 148 169 190
0.030 0.020
1 47 93 139 185 231 277 323 369 416
Attention weight aij
0.015
0.10 Attention weight aij
0.020
1 22 43 64 85 106 127 148 169 190
0.030
Attention weight aij
Attention weight aij
0.025
1 20 39 59 78 98 117 137 156 176
Attention weight aij
Ground truth ABMIL Smooth Operator Normal Guidance
0.06 0.04 0.02 0.00
150
0
50
100 Slice index j
150
0
100
200 300 Slice index j
400
1 47 93 139 185 231 277 323 369 416
Attention weight aij
1 25 49 73 97 121 145 169 193 217 0.030
0.07 0.06 0.05 0.04 0.03 0.02 0.01 0.00
0
100
200 300 Slice index j
400
Figure C.3: Learned attention weights for 4 example CT scans (columns) from the Abdomen CT dataset. Top row: methods based on ABMIL. Bottom row: transformer-based MIL methods.
D
Further Details about Baseline Models and Experiments
D.1
Compute Resources
All experiments use either 4 Intel Xeon 6342 CPUs (2.80 GHz) or 4 Intel Xeon 6226R CPUs (2.90 GHz) and either 1 NVIDIA RTX 6000 (24 GB), 1 NVIDIA RTX A6000 (48 GB), or 1 NVIDIA A100 GPU (40 GB). D.2
Localization Experiment Details
For most MIL methods, the attention score aij used to evaluate localization is straightforward. For max pooling MIL only, we need to construct a post-hoc attention score for each instance. Because we use element-wise max pooling MIL, we can construct attention weights αij = PM 1 ′ k=1 1[argmaxj ′ (hij k ) = j] where 1[·] denotes the indicator function. M D.3
Model Details
All MIL models use a frozen ViT-B/16 encoder (M = 768); only the pooling head and linear classifier are trained. Smooth Operator (S M). We use an approximation of S M following Castro-Macı́as et al. [2024], defined by T = 10 iterations of g (t+1) = (1−α) hi +α A g (t) starting from g (0) = hi . Here A is the normalized adjacency matrix of the chain graph representing the slice ordering of the 3D scan. The mixing scalar is parameterized as α = σ(α̃) with α̃ ∈ R learnable and initialized so αt=0 = 0.5. For SmAP, a single S M module is applied to the encoder’s output embeddings prior to ABMIL. For SmTP, a separate S M module (each with its own learnable α) is inserted after each MHSA layer. TransMIL. We follow Shao et al. [2021]: a learnable class token is prepended to hi , followed by two multi-head self-attention blocks separated by a Pyramidal Position Encoding Generator (PPEG) module. The class-token output of the second block is the bag representation. Unlike the reference implementation, which approximates self-attention with Nyström attention, we use full scaled dotproduct self-attention (torch.nn.MultiheadAttention). Both blocks use H = 8 heads at model dim M = 768 (per-head dim 96) with no dropout. PPEG is three parallel depthwise 1D convolutions with kernel sizes {3, 5, 7}, stride 1, and SAME padding, summed and added to the non-class-token positions. Reported slice-level attention ai,j is the class-token row of the second block’s attention map, averaged over heads, with the class-token-to-class-token entry removed. 16
E
Semi-Synthetic Data and its Bayes Estimator for Localization
For our data-generating process, we set M = 768 and R = 12. We sample Si uniformly between 20 and 60 for any bag regardless of its label. We draw yi , the bag-level binary label, from a Bernoulli distribution with probability 0.5. For negative bags, we draw all Si feature vectors hij ∈ RM from the same simple distribution: each entry is drawn iid from a Gaussian with mean 0 and variance 1. For positive bags, we draw the starting index ui of the modified R adjacent instances from a uniform distribution over the valid range p(ui |Si , yi =1) = UnifPMF(ui |{1, . . . , Si −R+1}).
(5)
Next, given the ui for a positive bag we draw only the first feature (of M ) of every instance in the chosen block of size R from a Gaussian with shifted-mean ∆ = 0.5. Otherwise, all features are drawn from the same zero-mean, unit variance Gaussian that defines the negative bags: Si Y M Y NormPDF(hijk |0.5, 1), if j ∈ [ui , ui +R−1] and k = 1 p(hi |ui , Si , yi =1) = (6) NormPDF(hijk |0, 1), otherwise. j=1 k=1
Using the above as ingredients, we can compute the marginal likelihood of an observed feature vector alone via the sum rule, summing over all possible starting indices p(hi |Si , yi =1) =
Si X −R+1
p(hi |ui = u, Si , yi =1)p(ui = u|Si , yi =1).
(7)
u=1
The three terms above allow computation of the Bayes estimator for the instance-level label via the right-hand-side of Eq. (4). Relevant set of ui values for index j. In the main paper, we define the set Uij as the set of possible ui values that would make index j fall in the chosen latent segment of R consecutive instances. We can define this set more formally as Uij = {ui ∈ Z : max(1, j − R + 1) ≤ ui ≤ min(Si − R + 1, j)} We know the size of the set is always at most R: |Uij | ≤ R. It will only be smaller than R for indices j within R steps of either edge.
F
Label Guidance
An assumption in MIL is that instance-level labels are unknown during training. However, if instance-level labels {yi,1 , . . . , yi,Si } are known, they can be used to define a discrete reference distribution ri to regularize the learned attention weights ai . For example, for a negative bag we could define a uniform distribution over all instances p(j|yi =0) = UnifPMF(j|{1, . . . , Si }) and for a positive bag we could define a uniform distribution over all positive instances p(j|yi =1) = UnifPMF(j|{k ∈ {1, . . . , Si } : yi,k =1}). These attention weights would achieve a perfect score for instance-level AUROC and AUPRC. Recent work has used instance-level semantic labels to guide attention [Liu et al., 2024]. Other works have used instance-level semantic labels to train classifiers to supervise attention [Huang et al., 2023].
G
Additional Related Work
Inductive biases in MIL. Other work in MIL has pursued inductive bias in the dependencies across instances. To account for 2D spatial relations among instances in MIL for whole slide imaging (WSI), Peled et al. [2026]’s probabilistic spatial attention (PSA) method modifies self-attention with a learnable distance-decayed prior. This makes nearby instances (patches) likely to attend to each other in the Si × Si pairwise self-attention matrix. Unlike our approach with a regularization term that is agnostic to any pooling strategy, their approach directly modifies the construction of attention values and relies on multi-head self-attention as a specific attention mechanism. PSA also does not 17
necessarily encourage unimodality, which seems to be a beneficial inductive bias in 3D CT datasets we tackle here. Gaussian and GMM ideas in MIL. Distantly related work [Zhang and Feng, 2021] outside of MIL has used Gaussian mixture ideas to determine how one token should pay attention to other tokens in a text sequence for NLP translation tasks, hoping to encourage concentrated attention. Our work avoids the expense of a distinct GMM for every instance in the sequence, instead forming just one Gaussian-like discrete distribution per head. Our work is also distinct in its purpose to improve localization for MIL in 3D CT settings. Our work is somewhat related to other efforts in MIL to incorporate probabilistic modeling ideas, such as work using Gaussian Processes to predict instance-level and bag-level labels given permutation-invariant per-instance embeddings [Wu et al., 2021]. That work uses Gaussian machinery to model a latent response function which is then transformed into a binary label; they do not regularize attention as we do toward helpful inductive biases such as unimodality and smoothness. G.1
Discussion of Past work on Upper Bounds using Instance-Level Labels
Correia et al. [2020] measure the performance of a permutation-invariant MIL on in-the-wild speech in Youtube videos to detect Parkinson’s or depression. They develop a “fully supervised” upper bound for MIL via a classifier that consumes one instance at a time, somehow comparing this classifier given bags of size 1 to MIL classifiers given bags of larger size in the same figure. Guillaumin et al. [2010] consider a similar instance-level upper bound in a metric-based MIL approach for multi-label MIL problems. Instead, our work develops separate best-in-class ceilings for both instance-level and bag-level tasks, uses a local neighborhood for the 1D CNN instance-level classifier instead of a “one instance in and one label out” design, and always evaluates competitor methods on bags of the same size (much larger than one). G.2
Alternatives to Weak Supervision
An alternative to weakly-supervised MIL with zero annotation costs is fully-unsupervised anomaly detection. Unsupervised methods for detecting brain lesions from 3D MRI scans have been proposed using diffusion models [Behrendt et al., 2024] or self-supervised learning where the pretext task is predicting the location of a given patch within the 2D axial slice [Baker and Brockmeier, 2026]. These approaches generally rely on an expensive bespoke-trained representation rather than frozen encoders that are more easily ported across medical image modalities, and also are likely less accurate at scan-level prediction than the deep MIL studied here due to the lack of supervised labels informing model training.
18