Title: ‌Navigating the MIL Trade-Off: Flexible Pooling for Whole Slide Image Classification
Venue: NeurIPS 2025
Accepted: yes
Keywords: Whole Slide Image

Self-Supervised Learning

Multiple Instance Learning

Data Augmentation
Abstract: Multiple Instance Learning (MIL) is a standard weakly supervised approach for Whole Slide Image (WSI) classification, where performance hinges on both feature representation and MIL pooling strategies. Recent research has predominantly focused on Transformer-based architectures adapted for WSIs. However, we argue that this trend faces a fundamental limitation: data scarcity. In typical settings, Transformer models yield only marginal gains without access to large-scale datasets—resources that are virtually inaccessible to all but a few well-funded research labs. Motivated by this, we revisit simple, non-attention MIL with unsupervised slide features and analyze temperature-$\beta$-controlled log-sum-exp (LSE) pooling. For slides partitioned into $N$ patches, we theoretically show that LSE has a smooth transition at a critical $\beta_{\mathrm{crit}}=\mathcal{O}(\log N)$ threshold, interpolating between mean-like aggregation (stable, better generalization but less sensitive) and max-like aggregation (more sensitive but looser generalization bounds). Grounded in this analysis, we introduce Maxsoft—a novel MIL pooling function that enables flexible control over this trade-off, allowing adaptation to specific tasks and datasets. To further tackle real-world deployment challenges such as specimen heterogeneity, we propose PerPatch augmentation—a simple yet effective technique that enhances model robustness. Empirically, Maxsoft achieves state-of-the-art performance in low-data regimes across four major benchmarks (CAMELYON16, CAMELYON17, TCGA-Lung, and SICAP-MIL), often matching or surpassing large-scale foundation models. When combined with PerPatch augmentation, this performance is further improved through increased robustness. Code is available at \href{https://github.com/jafarinia/maxsoft}{\texttt{https://github.com/jafarinia/maxsoft}}


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies pooling operations for weakly supervised, bag‐level classification of whole‐slide images (WSIs). It first derives theoretical bounds showing how the Log-Sum-Exp (LSE) pooling’s temperature $\\beta$ trades off generalisation error bounds. Building on this, the authors propose MaxSoft pooling, which applies a hard max during the forward pass but uses the differentiable LSE gradient (with tuned $\\beta$) in the backward pass. They introduce PerPatch augmentation, which precomputes m augmented variants per patch and samples them independently to increase bag diversity without incurring runtime costs. Experiments on four pathology benchmarks (CAMELYON16/17, TCGA-Lung, SICAP-MIL) and five classical MIL datasets compare MaxSoft and PerPatch against standard pooling and recent attention-based MIL methods.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper presents theoretical derivations concerning the effect of the $\\beta$ parameter on generalisation error (Theorem 1) and true positive/false negative rates (Theorem 2). While I appreciate the effort to provide theoretical grounding for the proposed method, a detailed verification of the correctness of these proofs is unfortunately beyond my expertise. Therefore, I have proceeded with the assumption that the proofs as presented in the main paper and the Appendix are mathematically correct.\n\n### **Strengths:**\n- Simplicity of MaxSoft: The forward-max/backward-LSE trick is easy to implement and introduces minimal overhead over attention-based pooling methods or other methods that require extra parameters.\n\n- Although beyond my area of expertise, the theoretical analysis seems to support the idea of combining max-pooling and mean-pooling.\n\n- PerPatch precomputes m variants per patch, allowing for the pre-computation of their respective features using the frozen encoder.\n\n- The experiments use a lot of different backbones for extracting features, and the authors also run ablation studies for the parameter $\\beta$.\n\n### **Weaknesses:**\nMy biggest concern is the experimental setup, both dataset utilisation and baseline comparisons. \n1. For CAMELYON16/17 datasets, the use of only 50 and 99 WSIs, respectively, is highly restrictive and significantly limits the generalizability of the findings. The paper should:\n\t- Use the full dataset and report metrics on the full dataset. The numbers do not need to be SOTA as the proposed approach is simpler than other approaches.\n\t- If authors decide to use a smaller subset, then it should be explicitly stated that the proposed approach works better only in \"low-data regime\" and also justified (for example, due to the use of data augmentation). For the experiments, the authors should divide the train set into K equal parts and report the mean (with std) accuracy for those parts. Or randomly sample n WSIs, with replacement, but sample them a lot of times so as to cover the whole dataset (bootstrap).\n\n2. Furthermore, the reported performance of established baselines such as ABMIL and DS-MIL across all datasets appears consistently lower than typically expected in the literature. For CAMELYON16/17, it could be due to the use of a small subset of the training data (point 1. above) but the numbers are low for TCGA-Lung as well. For example, the accuracy and AUC for ABMIL and DSMIL are higher (compared to many but not all in the paper) even when using ResNet50 pre-trained on ImageNet as a backbone [1-3].\n\n3. While the paper theoretically discusses the generalisation gap, the experimental evaluation demonstrates only low variance in performance metrics across internal splits. An empirical validation of a smaller generalisation gap typically requires testing on an external (out-of-distribution) dataset.\n\n4. The proposed MaxSoft operation is very similar to a Straight Through Gumbel Softmax. Proper attribution to Straight Through Estimators (STEs) is missing.\n\n5. The authors' claim (lines 40-46) that Transformer-based architectures offer little benefit due to small WSI datasets overlooks a critical aspect of WSI data scale. While WSI counts may be limited, individual gigapixel WSIs yield an enormous number of patches, which is the true unit of training for patch-based Transformers. Therefore, the total patch count available, even from moderately sized WSI datasets, is often substantial and perhaps sufficient to meet the data demands of Transformer models. This is supported by recent literature [4] (section II D. *The effect of training data size*).\n\n6. In line 206, the authors claim that existing methods might even \"degrade performance\" but I could not find anything in the referenced papers to support this claim.\n\n\n**[1]** Shao, Zhuchen, et al. \"Transmil: Transformer based correlated multiple instance learning for whole slide image classification.\" Advances in neural information processing systems 34 (2021): 2136-2147.\n\n**[2]** Tang, Wenhao, et al. \"Feature re-embedding: Towards foundation model-level performance in computational pathology.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.\n\n**[3]** Zhang, Hongrun, et al. \"Dtfd-mil: Double-tier feature distillation multiple instance learning for histopathology whole slide image classification.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n**[4]** Aben, Nanne, et al. \"Towards large-scale training of pathology foundation models.\" arXiv preprint arXiv:2404.15217 (2024).\n\nQUESTIONS:\n- How many augmentations (m) per patch did you use? How to pick the number m for a dataset? I suspect that one gets diminishing (or even negative) returns after a specific value.\n\n- It seems like $\\beta$ needs to be tuned to a dataset. Did you tune the values of $\\beta$ using the validation split?\n\n- Adaptive Noisy-AND pooling [5] also tries to balance Max and Mean pooling. It may be worth mentioning in the related work.\n\n- Why does related work include a section on foundation models? Is it only because the paper uses these models as frozen encoders?\n\n**[5]** Kraus, Oren Z., Jimmy Lei Ba, and Brendan J. Frey. \"Classifying and segmenting microscopy images with deep multiple instance learning.\" Bioinformatics 32.12 (2016): i52-i59.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- This paper argues that attention and transformer-based MIL models offer limited benefits without massive data and computational resources. Instead, the paper returns to classical pooling strategies like Log-Sum-Exp (LSE). It theoretically analyze LSE pooling, showing a trade-off between generalization and sensitivity, depending on the temperature parameter which makes it choose between a mean and max pooling function.\n\n- The paper proposes MaxSoft, a new pooling method that uses the max operation in the forward pass but smooth gradients (via LSE) in the backward pass for stability.\n\n- It also introduces PerPatch augmentation, a simple strategy to improve training robustness by applying random augmentations independently to each patch in a WSI.\n\n- It show strong empirical results, where MaxSoft and PerPatch outperform both classical and Transformer-based methods, often matching the performance of massive foundation models with far lower compute.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The theoretical analysis explaining when and why different pooling strategies work, and showing generalization vs. sensitivity trade-offs with max vs mean is a useful framework which leads to intuitive conclusions on utility of each pooling type. Theorem 2 showing generalization vs TPR tension is something practitioners have observed with MIL models and the theoretical analysis helps structure the argument well.\n- MaxSoft is an interesting idea where max pooling is used for forward pass but a smooth approximation is used for the backward pass using a moderate $\\beta$ in LSE pooling. It's simple to implement and results in no extra parameters or complexity.\n- The authors make a strong and honest case against over-reliance on attention/transformer models in small datasets, which is often glossed over in literature.\n- The paper has extensive experimental results. Experiments cover multiple datasets (CAMELYON16/17, TCGA-Lung, SICAP-MIL) and compare with powerful baselines (Prov-GigaPath, SSRDL, ABMIL, DSMIL, Snuffy, R2T-MIL).\n\n**Weaknesses**\n- Generalization-sensitivity tradeoffs predominantly exist in problems where the signal in the WSI is very localized. This is true for Camelyon and the SICAP-MIL dataset, but other problems like NSCLC subtyping or prediction of molecular signatures from WSI where signal is more diffused, these pooling type tradeoffs changes significantly. This is seen even in the results in TCGA-Lung vs other datasets. This solution offers less value in these cases.  \n- MaxSoft requires manual selection of 𝜔 (tuned via validation), and performance is sensitive to this choice.\n- During backward pass, gradients come from a different function (LSE with moderate 𝜔) than the forward pass (hard max). This breaks consistency since the gradients are not true gradients of the original function. It also leads to training vs inference mismatch, it can hurt generalization. \n- Assessing true generalization without an independent heldout test set is hard and results to show how this method translates to better generalization and lower FNR for Camelyon problems should ideally deploy on a new test set outside of the competition.\n- Many MIL implementations apply per-patch augmentations randomly in the data loader, so it happens by default. Specific papers and review papers and their code [2] have mentioned applying per-patch augmentations, so this doesn't seem to be a novel contribution.\n\n[1] Gadermayr, Michael, and Maximilian Tschuchnig. \"Multiple instance learning for digital pathology: A review of the state-of-the-art, limitations & future potential.\" Computerized Medical Imaging and Graphics 112 (2024): 102337.\n\n[2] Juyal, Dinkar, et al. \"SC-MIL: Supervised Contrastive Multiple Instance Learning for Imbalanced Classification in Pathology.\" Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2024.\n\nQUESTIONS:\n- Can we have the third dataset in Figure 3 as well. The effect of $/beta$ on the a problem with diffused signal is important to note so readers are aware of the kind of problems this is most suitable for.\n- Given how critical ω is to MaxSoft’s success, have you considered a way to automatically choose it based on bag size, dataset statistics, or validation confidence?\n\nLIMITATIONS:\nThe paper does mention limitations briefly, but it's not very specific to their analysis or method, but applies more generally for most MIL methods.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses the problem of Multiple Instance Learning (MIL) in the context of Whole Slide Image (WSI) classification. The authors argue that the main limitation of transformer-based MIL methods lies in their dependency on large-scale training data: \"Transformer models yield only marginal gains without access to large-scale datasets.\" To mitigate this issue, the authors propose a pooling strategy inspired by both max and mean pooling. Specifically, they introduce MaxSoft, a simple, differentiable alternative to attention-based aggregation mechanisms. Additionally, they propose a novel augmentation scheme, PerPatch, which enhances patch diversity during training and improves robustness.\n\nSTRENGTHS AND WEAKNESSES:\nWeaknesses\n1) The central motivation—that transformer-based MIL methods underperform in low-data regimes—is underdeveloped and insufficiently supported. While data scarcity is a known challenge in digital pathology, it is unclear whether the limitations observed by the authors are intrinsic to transformer architectures or confounded by other variables (the type of encoder used, the quality and domain alignment of pretraining data). \n\n2) The paper would benefit from a more rigorous empirical characterization of why and when transformer-based methods fail in these settings. This is especially important given that the authors’ main contribution is motivated by this claim. Moreover, while the proposed MaxSoft pooling is elegant and easy to implement, it is not obvious why this simpler mechanism should not suffer from the same data-efficiency limitations as attention mechanisms (no experiment is provided about data scarcity). \n\n3) Another omission is the lack of discussion of non-transformer MIL alternatives, such as graph-based models [1], which are relevant and increasingly adopted in histopathology. Including such methods in the benchmark—or at least discussing them—would strengthen the experimental comparison. Some recent MIL baselines not mentioned in the paper that should be considered for a more comprehensive evaluation include [2–8].\n\nStrengths\n1) MaxSoft is a lightweight and theoretically well-motivated pooling strategy. It is simple to implement and does not require many parameters.\n2)The PerPatch augmentation scheme is intuitive and may be beneficial across various MIL pipelines, not only the proposed one.\n3)The experimental evaluation is extensive, spanning multiple WSI datasets, with thorough ablations and implementation details included in the supplementary material.\n\nSuggested Missing References\n[1] Graph Neural Networks in Histopathology: Emerging Trends and Future Directions\n[2] H2-MIL: Exploring Hierarchical Representation with Heterogeneous Multiple Instance Learning for Whole Slide Image Analysis\n[3] A Trainable Optimal Transport Embedding for Feature Aggregation and Its Relationship to Attention\n[4] Morphological Prototyping for Unsupervised Slide Representation Learning in Computational Pathology\n[5] DAS-MIL: Distilling Across Scales for MIL Classification of Histological WSIs\n[6] A Graph-Transformer for Whole Slide Image Classification\n[7] Prototypical Multiple Instance Learning for Predicting Lymph Node Metastasis of Breast Cancer from Whole-Slide Pathological Images\n[8] TransMIL: Transformer-Based Correlated Multiple Instance Learning for Whole Slide Image Classification\n\nQUESTIONS:\n1) The central motivation—that transformer-based MIL methods underperform in low-data regimes—is underdeveloped and insufficiently supported. While data scarcity is a known challenge in digital pathology, it is unclear whether the limitations observed by the authors are intrinsic to transformer architectures or confounded by other variables (the type of encoder used, the quality and domain alignment of pretraining data). \n\n2) The paper would benefit from a more rigorous empirical characterization of why and when transformer-based methods fail in these settings. This is especially important given that the authors’ main contribution is motivated by this claim. Moreover, while the proposed MaxSoft pooling is elegant and easy to implement, it is not obvious why this simpler mechanism should not suffer from the same data-efficiency limitations as attention mechanisms (no experiment is provided about data scarcity). \n\n3) Another omission is the lack of discussion of non-transformer MIL alternatives, such as graph-based models [1], which are relevant and increasingly adopted in histopathology. Including such methods in the benchmark—or at least discussing them—would strengthen the experimental comparison. Some recent MIL baselines not mentioned in the paper that should be considered for a more comprehensive evaluation include [2–8].\n\nLIMITATIONS:\nSeveral competitor are missing and despite the strong theoretical content of the supplementary material the main motivation is still not adequately supported.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}