ADMIL: Attention-Distilled Multiple Instance Learning for Selective Foundation Model Inference in Pathology

Duncan Stothers; Ren-Chin Wu; William Lotter

Source-linked AI summary

ADMIL: Attention-Distilled Multiple Instance Learning for Selective Foundation Model Inference in Pathology

Duncan Stothers, Ren-Chin Wu, William Lotter

arXiv:2608.22066v1Aug 22, 2026cs.CVcs.AIcs.LG

TL;DR

Exhaustive foundation-model encoding of every foreground tile is costly even when ABMIL attention concentrates on a small informative subset. ADMIL distills teacher attention into PriorNet for top-K tile selection, then applies a selected-bag ABMIL student; across three benchmarks, it preserves headline performance while avoiding most expensive embeddings, with limitations concerning teacher reliability, operating-point selection, validation scope, and training overhead.

Problem
Exhaustive foundation-model inference over thousands of foreground tiles can dominate latency, compute costs, and throughput despite concentrated task-relevant ABMIL attention.
Method
ADMIL trains KL-distilled PriorNet to rank foreground tiles from raw images, selects the top-K, and uses their foundation-model embeddings in a separate selected-bag ABMIL predictor.
Results
ADMIL matches full-teacher mean headline performance using 4 BRACS tiles, 8 PANDA tiles, and 128 CAMELYON16 tiles, with foundation-model calls reduced by over 98% in all three datasets.
Takeaways & Limitations
Teacher attention can be distilled early enough to reduce expensive tile inference while preserving slide-level performance across three distinct pathology tasks.
Takeaways & Limitations
ADMIL depends on stable, relevant teacher attention, requires validation-based K selection for deployment, and needs broader dataset and foundation-model validation while adding training overhead.

Abstract

from arXiv · show

Attention-based multiple instance learning (ABMIL) using pathology foundation model embeddings is effective for slide-level tasks, but exhaustive inference requires applying a large image encoder to every foreground tile despite the subsequent attention distribution often concentrating over a small subset of informative regions. We introduce ADMIL (Attention-Distilled Multiple Instance Learning), a selective-compute framework that distills an ABMIL teacher's attention into a lightweight tile-selection model, PriorNet. Using an EfficientNet architecture, PriorNet learns the teacher attention distribution from raw tile pixels with KL divergence; at inference, it scores the foreground pool, selects the top-K tiles, and invokes the expensive foundation model only on that subset before a selected-bag ABMIL student predicts the slide label. Across BRACS, PANDA, and CAMELYON16, ADMIL matches full-teacher headline performance at K=4, 8, and 128 tiles, respectively, avoiding >98% of foundation model (Virchow2) tile embeddings and model inference FLOPs. Random and teacher-attention oracle controls show that this result depends on task-relevant selection rather than tile-count reduction alone. Quantitative and qualitative analyses suggest that PriorNet recovers the teacher's tile ordering with high fidelity while focusing on task-relevant morphological regions. ADMIL shows that nearly all expensive tile encodings can be removed without sacrificing slide-level performance, providing a potential path for more efficient deployment in clinical settings where latency and compute costs are key considerations.

Introduction

ADMIL addresses the cost of exhaustive foundation-model inference by distilling a teacher ABMIL model’s task-specific attention into PriorNet, which selects informative tiles before expensive encoding. The framework preserves high-capacity representations where they matter while targeting redundant foundation-model evaluations.

Thousands to tens of thousands of foreground tiles can make exhaustive foundation-model inference dominate pathology-slide latency, compute costs, and throughput.
ABMIL attention often concentrates on a small subset of morphologically informative regions, motivating cheap preselection before foundation-model inference.
ADMIL distills a trained ABMIL teacher’s tile attention into PriorNet using KL divergence, then trains a selected-bag ABMIL model on the highest-scoring tiles.
PriorNet gates access to the expensive foundation model rather than replacing its encoder or the ABMIL architecture, using task-specific teacher behavior instead of random subsampling.
98.5%–99.8% fewer foundation-model tile embeddings and 98.3%–99.6% lower model-inference FLOPs are reported while matching full-teacher headline performance across BRACS, PANDA, and CAMELYON16.
PriorNet selection shows strong concordance with teacher- and student-selected tiles in quantitative and qualitative evaluations.

Related Work

Related work establishes weakly supervised ABMIL and pathology foundation models as effective components of WSI analysis, while recent efficiency methods reduce processing through spatial selection, distillation, or cheaper encoders. ADMIL instead targets the upstream cost of producing foundation-model features for every candidate tile.

Weakly supervised WSI classification treats each slide as a tile bag with slide-level supervision, and ABMIL aggregates tile features through learned attention-weighted averaging.
Pathology foundation models provide strong general-purpose tile representations and improve downstream performance, but increase the cost of exhaustive tile encoding.
ZoomMIL and HDMIL reduce gigapixel pathology computation through coarse-to-fine processing or low-resolution attention distillation for higher-resolution selection.
Recent methods also distill large-model features or use cheaper WSI-level models for selection, whereas ADMIL focuses on efficient inference within foundation-model ABMIL pipelines.

Methods: ADMIL

ADMIL represents slides as foreground tile bags, uses a frozen full-bag foundation-model ABMIL teacher to provide attention supervision, and trains PriorNet to rank tiles for selective encoding. A separate top-K ABMIL student then predicts from only the selected foundation-model embeddings.

Problem Setting and Full-Bag Teacher: Each slide is represented as a foreground tile bag X = {x_i}^N, where the number of tiles varies by slide.
Problem Setting and Full-Bag Teacher: The frozen foundation-model teacher maps each tile to an embedding and exports its attention distribution as tile-level supervision for PriorNet.
Attention Distillation with KL Divergence: PriorNet produces one scalar score per foreground tile and normalizes scores across the slide before ranking candidates.
Attention Distillation with KL Divergence: PriorNet is trained with a KL distillation objective, using attention-guided sampling that retains highly attended tiles and samples additional tiles from the remaining foreground pool.
Top-K Selection and Selected-Bag ABMIL: ADMIL orders tiles by PriorNet score and selects the top-K candidates for foundation-model embedding at inference.
Top-K Selection and Selected-Bag ABMIL: A separate top-K ABMIL is trained on the selected-bag condition using the original slide-level task loss.

Methods: Evaluation Approach

The evaluation compares ADMIL with full-teacher, random-selection, and teacher-oracle baselines while measuring predictive performance, efficiency, and attention-selection fidelity. It also tests whether results are robust to the PriorNet supervision fraction.

Baselines and controls: ADMIL is evaluated against the full foundation-model ABMIL teacher, random top-K selection, and a teacher top-K oracle.The oracle estimates performance when teacher evidence is recovered perfectly but is not deployable because it requires full-bag inference.
Attention fidelity: Teacher-support retrieval measures the fraction of teacher top-K tiles recovered among PriorNet’s top-T tiles, with cells masked when T < K.Correlation and retrieval are computed from full-foreground teacher and PriorNet attention distributions, distinct from selected-bag student attention.
Robustness analysis: The supervision-budget analysis compares 5%, 10%, and 20% of foreground candidates while holding other procedures fixed to assess robustness to Z.The 10% setting is the default for headline results; alternative fractions test dependence on the supervision fraction rather than define another headline configuration.
Efficiency: Inference efficiency is reported as foundation-model embeddings avoided and FLOPs saved at the descriptive critical budget Kcrit.The FLOPs calculation includes PriorNet, foundation-model, and ABMIL inference operations.

Methods: Experimental Setup

Experiments evaluate ADMIL on three public WSI benchmarks using fixed test partitions, dataset-specific tasks and metrics, and three-seed training procedures. Virchow2 supplies tile embeddings, while PriorNet and ABMIL models operate on foreground-only tiles.

Datasets and tasks: The study evaluates CAMELYON16 metastasis detection, PANDA six-class ISUP grading, and BRACS using public WSI benchmarks.CAMELYON16 averages 12,391 foreground tiles per WSI, while PANDA averages 531.
Data splitting: Three global runs use seeds 0, 1, and 2, with fixed test partitions and label-stratified 70/30 training-validation splits.PANDA is split at the patient level because patients can contribute multiple slides.
Metrics: Performance is reported with dataset-specific metrics: AUROC and AUPRC for CAMELYON16, QWK and within-1 accuracy for PANDA, and balanced accuracy and macro AUROC for BRACS.Reported metrics average test-set performance across three independent runs, with uncertainty estimated using 1,000 hierarchical bootstrap replicates.
Model architectures: Virchow2 is the foundation-model tile encoder, full-bag teachers and selected-bag students use non-gated ABMIL, and PriorNet uses an EfficientNet-B0 backbone.All models process foreground-only 224 × 224 tiles.
Training: PriorNet is trained for 20 epochs with AdamW and KL divergence, independently of deployment budget K.Checkpoint selection uses the validation split.

Results

ADMIL reaches full-teacher performance at small tile budgets across BRACS, PANDA, and CAMELYON16 while avoiding nearly all foundation-model tile embeddings. PriorNet also recovers teacher attention ordering and produces qualitatively aligned, sparser student attention, with performance robust to supervision-budget variation.

Headline performance: K = 4, 8, and 128 first reach the teacher threshold for BRACS balanced accuracy, PANDA QWK, and CAMELYON16 AUROC, respectively.These budgets correspond to 99.8%, 98.5%, and 99.0% fewer foundation-model tile embeddings, respectively.
Inference efficiency: 99.8%, 98.5%, and 99.0% fewer foundation-model tile embeddings are used at the critical budgets for BRACS, PANDA, and CAMELYON16, respectively.The reduction translates to similar levels of model inference FLOPs saved given PriorNet’s lightweight architecture.
Selection controls: PriorNet performance consistently approaches the teacher-selection oracle and substantially exceeds random sampling, especially at low K.This comparison supports task-relevant selection rather than tile-count reduction alone.
Attention fidelity: PriorNet’s top selected tile matches the top teacher tile 14%, 9%, and 26% of the time for CAMELYON16, BRACS, and PANDA, respectively.Retrieval recall increases as PriorNet retrieval depth grows, indicating useful approximation of the teacher’s task-specific evidence ranking.
Qualitative attention: The selected-bag student largely aligns with teacher attention in examples while producing a sparser attention map.In CAMELYON16, both models emphasize tumor-cell tiles; in PANDA, overlapping high-attention tiles contain high-grade prostate carcinoma morphology.
Robustness: 5%, 10%, and 20% supervision-budget curves follow similar trajectories with overlapping confidence intervals across datasets and metrics.The authors retain 10% as the default and interpret the alternatives as evidence of robustness to the supervision fraction.

Discussion

ADMIL removes substantial redundant tile computation by using task-specific selection, while retaining full-teacher performance at very small budgets. Its effectiveness depends on learned relevance ordering, with limitations around teacher reliability, operating-point selection, validation scope, and training overhead.

Results: Four, eight, and 128 selected tiles match full-teacher headline performance on BRACS, PANDA, and CAMELYON16, respectively.These correspond to more than 98% fewer foundation model calls across all three datasets.
Selection mechanism: Random top-K students perform worse in the low-budget regime, whereas teacher-attention oracle and PriorNet selection retain more full-bag performance.This supports task-relevant selection rather than tile-count reduction alone as the explanation for ADMIL’s low-budget results.
Efficiency: 98.3%–99.6% lower model inference FLOPs and 98.5%–99.8% fewer foundation model tile embeddings result from PriorNet screening before Virchow2 inference.Virchow2 requires 340.132 GFLOPs per tile versus 0.769 GFLOPs for PriorNet, so scanning the foreground pool adds only about 0.226% of full-teacher model-forward FLOPs.
Attention fidelity: Spearman correlations exceeding Pearson correlations indicate that preserving teacher-relevant tile ordering matters more than exact attention-probability calibration.Correlation is higher for PANDA and CAMELYON16 than BRACS, where heterogeneous morphologies and a larger physical field of view may make isolated-tile attention harder to predict.
Attention fidelity: BRACS still reaches full-teacher balanced accuracy with four selected tiles despite imperfect global correlation.This suggests effective decision-focused retrieval does not require perfect attention agreement.
Sensitivity: 5%, 10%, and 20% PriorNet supervision budgets retain similar qualitative curves and largely overlapping uncertainty, without a consistent ordering.The experiments indicate that performance is not tied to one exact supervision budget.
Positioning: ADMIL is tailored to efficient inference in foundation model-based MIL pipelines while preserving their representations and prediction architecture.Its task-specific teacher supervision distinguishes it from task-agnostic compression and heuristic or random selection methods.
Limitations: ADMIL depends on stable, relevant teacher attention, requires validation-based K selection for deployment, and remains limited by three benchmark datasets, one foundation model, and added training overhead.The reported Kcrit is a post-hoc test-curve summary rather than a deployable validation-selected operating point.

Conclusion

ADMIL turns a trained ABMIL teacher’s attention into a selective-compute policy using PriorNet and a selected-bag ABMIL student. Across three tasks, it preserves teacher performance while avoiding most foundation model embeddings, with supporting evidence from attention agreement, retrieval, sensitivity, and morphology analyses.

Conclusion: ADMIL converts a trained ABMIL teacher’s attention into a selective-compute policy with a KL-trained PriorNet and selected-bag ABMIL student.The framework distills teacher attention before selecting tiles for expensive foundation model processing.
Conclusion: Across three distinct tasks, ADMIL preserves teacher performance while avoiding most foundation model embeddings.The conclusion is supported jointly by quantitative attention agreement, teacher-support retrieval, supervision-budget sensitivity, and morphology-consistent examples.

Inference Efficiency Assessment Details

Inference efficiency is summarized using the critical tile budget needed to reach teacher performance, alongside avoided foundation-model embeddings and saved model-inference FLOPs. These FLOP reductions exclude I/O and preprocessing, and Kcrit is a post-hoc descriptive statistic rather than a validation-selected deployment hyperparameter.

Efficiency is assessed by the percentage of foundation-model tile embeddings avoided and FLOPs saved.FLOP accounting includes PriorNet scoring over the full foreground pool and the teacher/student MIL heads.
FLOP reductions describe model-forward computation rather than wall-clock latency or total computation.I/O and preprocessing are excluded from the calculation.
Kcrit identifies the minimum tile budget at which ADMIL reaches teacher performance for a dataset and task.It is defined from the complete test performance curve.
Kcrit is a post-hoc descriptive summary computed from complete test performance curves, while checkpoint selection uses validation data only.It should not be interpreted as a validation-selected deployment hyperparameter.