Rebuttal For NKLh

We thank the reviewer for the thoughtful evaluation. We address the points raised under Weaknesses, Questions, and Limitations below:

W1

We thank the reviewer for scrutinizing our data splits and apologize for omitting details in the main text. We selected 50 CAM16 WSIs to emulate the extreme low-data regimes typical of rare cancers (≈70 WSIs/type [1]), a scenario closely aligned with few-shot learning. CAM16’s official test set is used in full (l. 525–526), matching standard protocols. As noted (l. 145–148), patch-level outputs are critical for pathologists, so patch metrics remain central. For CAM17, only 500/1000 slides have slide-level labels, just 100 carry pixel annotations (one discarded for quality; l. 508, 531), limiting patch evaluation (l. 530–531). Despite this, our method still outperforms prior approaches (Tabs. 1, 2). We have now run four-fold experiments on complete CAM16 and CAM17 splits (omitting CAM17 patch metrics) to confirm our initial splits and validate our conclusions. While CAM16 evaluation follows its standard test set, we acknowledge K-fold cross-validation could further bolster robustness. For CAM17, our random stratified sampling across Macro/Micro/ITC classes reduces variance relative to naïve splits [2], though it does not replace full cross-validation. We will clarify these choices and incorporate full-dataset results in the final manuscript.

W2

We thank the reviewer for the careful comparison. On CAM16, the lower baseline scores reflect the extreme low-data split; transformer-based MIL methods are inherently data-hungry. For CAM17, few papers report slide-level results: [3] obtains AUC = 0.863/ACC = 0.815 for DSMIL on 500 WSIs, while our closest configuration (DINO-Domain) yields 0.833/0.800 on 100 WSIs. Their encoder was SimCLR-trained on CAM16 patches, ours DINO-trained on in-domain patches; the larger sample size partly offsets their out-of-distribution (OOD) features. For TCGA, any apparent discrepancy is purely presentational: our table shows baseline AUCs slightly above [5] and slightly below the others. Such small variations are expected given differences in splits, background thresholds, JPEG quality, backbone pre-training (SimCLR vs DINO), and unspecified hyperparameters—there is currently no community standard. We therefore maintain that our evaluation is transparent, rigorous, and fully consistent with prior work.

Reported metrics for TCGA in different papers:

W3

We thank you for raising this point, which Reviewer tjrh also mentioned, external experiments to evaluate generalization. We think out-of-distribution generalization experiments are out of our work’s context; see our response to reviewer tjrh Weakness 4 for further details. Nevertheless, we appreciate more discussion on different aspects of our method’s improvement in generalization. Variance is not uniformly low: Tables 1 and 3 show that Max-pooling consistently exhibits higher variance than MaxSoft across all datasets—especially CAM16 and CAM17—and across all encoders. This aligns with our theory: Max-pooling is LSE-pooling with β = ∞, lacks differentiability, and consequently widens the generalization gap (Theorems 1 & 4).

W4

We appreciate the reviewer noting the similarity to Straight-Through Gumbel Softmax [9]. Historically, STEs replaced non-differentiable binary activations with the identity in the backward pass; the more general Backward-Pass Differentiable Approximation (BPDA) was formalized in [11], and—as STE is a special case of BPDA—we cited BPDA at l. 61. Functionally, Straight-Through Gumbel Softmax forwards a categorical sample but back-propagates the softmax gradient, whereas MaxSoft forwards a hard max and back-propagates the Log-Sum-Exp gradient, a choice informed by our theoretical bounds on LSE temperature to tighten the generalization gap. Furthermore, STE-based methods target discrete latent variables, while our work handles continuous patch scores and slide-level outputs, so the correspondence is not straightforward. We will cite both Straight-Through Gumbel Softmax and BPDA in the final manuscript and have included an STE-variant ablation in our response to Q3.

W5

We thank the reviewer for the critical perspective; it sharpens our manuscript. Our main point may not have been explicit enough. In the WSI-MIL pipeline, the final stage trains a slide-level model on frozen features, so the effective sample count is the number of WSIs, not their many patches. This is analogous to training ViT-S/16 on ImageNet-1K: the dataset is 1.28 M images, not $(224 \div 16) \times (224 \div 16) \times 1.28 \approx 250$M, which is the number of patch tokens within them. Slide-level sample scarcity is well documented in the literature [7, 14] and motivates all augmentation baselines we report. The referenced work [13] addresses patch-level classification with patch labels, not slide-level MIL. It thus lies outside the scope of our study and does not speak to our claim. We will further clarify this distinction in the final revision.

W6

We apologize for the confusion from last-minute edits: the encoder for the augmentation experiments was inadvertently unstated, and the baselines from Tables 1 and 3 were omitted in Tables 2 and 4. Nevertheless, DINO-Domain results confirm our claim at l. 206: comparing Tables 1→2 and 3→4, MaxSoft outperforms all augmentation baselines. On CAM16 it achieves AUC 0.934 / ACC 0.863 / Patch AUC 0.919 (vs. lower scores for Remix, RankMix, AugDiff, SSRDL); on CAM17 it reaches Slide AUC 0.983 / ACC 0.867 / Patch AUC 0.839; for TCGA it records Slide AUC 0.940 / ACC 0.885; and on SICAP Slide AUC 0.850 / ACC 0.808. We will explicitly specify the encoder and include the missing baseline rows in the final manuscript.

Q1

We chose m=3 (Table 5; App J) using Random Rotation, Random Gaussian Blur, and Random Color Jitter. A single augmentation collapses to PerSlide sampling and degrades performance (Tables 2, 4; l. 276–277), whereas PerPatch sampling yields diverse per-patch variations and drives the gains we observe. We dropped Random H&E-DAB Jitter (low density/coverage; no AUC‐variance or ECE benefit), Random Elastic Deformation, and Random Affine Transformation (limited real-world relevance). In practice, augmentations with higher density/coverage and lower FID (i.e., closer to the original data manifold) perform best. Our trio was selected on these qualitative grounds rather than via an exhaustive sweep over mm; we can provide a performance-vs-mm curve if desired.

Q2

Yes. As noted in l. 267, we tune β on the validation split; we will state this more explicitly in the final manuscript.

Q3

Thank you for mentioning this. We will make sure to incorporate it in the related work. Also, to provide a more comprehensive analysis on methods that are relatively similar to our method, Noisy-AND and Straight Through Gumbel Softmax, we decided to provide experiments on them following the setup that we already tested in our manuscript on CAM17 but due to character limit we cannot provide now and will provide in discussion if asked.

Q4

We thank the reviewer for suggesting Adaptive Noisy-AND pooling. We’ll cite it in Related Work and have run comparative experiments—including Adaptive Noisy-AND and the STE variant of Straight-Through Gumbel Softmax—under our CAM17 setup. Due to space constraints, we’ve omitted the detailed results here but can provide them upon request.

[1] Ding T., Multimodal Pathology Foundation Model [2] Haugh M., Further Variance Reduction Methods (note) [3] Fourkioti O., CAMIL: Context-Aware MIL [4] Chen T., Simple Contrastive Learning Framework [5] Shao Z., TransMIL [6] Tang W., Feature Re-Embedding [7] Zhang H., DTFD-MIL [8] Dual-stream MIL Network (Self-supervised) [9] Jang E., Gumbel-Softmax Reparameterization [10] Bengio Y., Gradients Through Stochastic Neurons [11] Athalye A., Obfuscated Gradients [12] Yin P., Straight-Through Estimator [13] Aben N., Large-Scale Pathology Foundation Models [14] Liu P., Pseudo-Bag Mixup [15] Kraus O., Deep MIL for Microscopy

We thank the reviewer for the constructive evaluation. We address the points raised under Weaknesses, Questions, and Limitations below:

W1

We thank the reviewer for noting that on diffuse-signal tasks like TCGA-Lung (≈80 % slide-level signal), MaxSoft’s localization-oriented pooling yields smaller gains (Tables 1 & 3; [2–6]). Such tasks are inherently easier—strong baselines already perform well—so the narrower margin reflects task simplicity, not a weakness of our method. Nonetheless, even modest improvements are valuable in high-stakes diagnostics. We will incorporate this balanced discussion and clarify our focus on slide-level classification rather than molecular-signature prediction.

W2&Q2

There is no ω discussed in the entire paper, and we think you mean β so we answer assuming you meant β. We agree that this is part of our method, what we did, and how we experimented. Although we argue that at the moment there is no solution in the deep learning community that does not require selection of hyperparameters, and we do it in a principled way by choosing it with our independent validation set, which we do not train the model on. Below is a thorough explanation of how we chose β’s range based on our theoretical analysis,s which also matches with our experiments and can be considered as automatic.

1. What the theory already gives us

• Upper anchor (pivot to max‑like behavior). Corollary 1 shows that the soft weights stop being almost‑uniform and start collapsing on the maximum once

where $N$ is the number of patches and $\varphi=q_{N/2}-q_{1}$ is an (unknown) inter‑quantile gap. In practice we drop $\varphi$ (or set $\varphi\approx1$) and simply treat

as the sensitivity ceiling.

• Lower anchor (guaranteed true‑positive rate). Theorem 6 proves that, with confidence $1-\alpha$, any desired TP rate is achieved provided

where $N_c$ is the (random) number of tumor patches, $V$ their sub‑Gaussian variance and $\gamma\!\in\!(\tfrac12,1)$ the tail exponent. Taking the “worst reasonable” constants $C_\alpha=1,\;V=1,\;\gamma=\tfrac12$ gives a safe lower bound

Between these two anchors every β is Pareto‑optimal: smaller values tighten the generalization bound (Theorem 4, small‑β analysis), while larger values buy sensitivity at the cost of variance.

2. A suggestion for β picker

1. Estimate $N$ and $N_c$. $N$ is known at WSI load‑time; $N_c$ can be (i) the empirical tumor fraction on the training split, or (ii) a running Bayesian estimate updated after each epoch.
2. Compute the bracket $[\beta_{\text{low}},\beta_{\text{high}}]$ exactly as above.
3. Select β on‑the‑fly
   • During training we start at the geometric mean $\beta_{0}=\sqrt{\beta_{\text{low}}\beta_{\text{high}}}$ to balance variance and sensitivity.
   • Every $k$ epochs we re‑evaluate validation AUC and, if the β‑distribution‑adjusted confidence interval still overlaps the desired TP threshold, we decay β by √2; otherwise we increase it by √2, but never step outside the bracket.

• Because both anchors are monotone in $N$ and $N_c$, this scheduler stays differentiable and data‑driven; no manual grid search is required.

3. Assumptions made explicit

• Constants $C_\alpha=V=1,\;\gamma=\tfrac12$ correspond to unit‑variance sub‑Gaussian logits with mid‑heavy tails—exactly the regime sampled in Appendix A.4.
• If you have tighter priors (e.g. heavier tails ⇒ larger $\gamma$), simply plug them into (*); the code path is identical.

In short: yes, the paper’s theorems give two numerical guideposts; turning them into an automatic scheduler is straightforward, and we now include that reference implementation in the supplementary.

W3

We derive generalization guarantees via uniform‐convergence bounds: a drop in training loss implies a drop in test loss regardless of optimizer. Empirically, training loss decreases and MaxSoft generalizes best across datasets. As noted (l. 61), our update follows BPDA [7], whose cure for forward–backward mismatch is extra epochs; accordingly, we train all methods for 500 epochs (MIL training is fast—≈5 min on CAM17; Table 1). A formal optimization analysis, akin to work on STEs in binary nets [8], remains future work.

W4

We thank the reviewer for stressing rigorous evaluation. As detailed in Sec. 5.2, we use a standard split with an independent test set, and all proofs (Thm. 1; App. A.1, Lem. 3, 5; Thm. 6) assume i.i.d. sampling. Out-of-distribution (OOD) generalization is beyond our primary scope; nonetheless, we report the only cross-dataset check our data allow—training on CAM17 and evaluating on CAM16—with results in the table below. This aligns with well-known findings that higher in-distribution accuracy tends to boost OOD performance [9–12].

W5

We thank the reviewer for calling attention to our PerSlide vs. PerPatch distinction. In our setup, PerSlide applies the same augmentation to every patch in a WSI, whereas PerPatch samples independently per patch. For example, ABMIL’s “random rotation and mirror every patch” is actually PerSlide, since all tiles use the same transform. The works it builds on [13] provide no augmentation code or detail, and other related papers [14, 15] only report per-slide mean–std normalization. To our knowledge, no prior method explicitly implements PerPatch as we define it. We will clarify these definitions in the final manuscript.

Q1

We thank the reviewer for the suggestion. As images cannot be included in the rebuttal, we reproduce Fig. 3 as a table spanning all four datasets, thereby adding the requested third dataset and illustrating β’s effect on both diffuse- and localized-signal tasks.

L1

We thank the reviewer for this insightful point. Our theory indeed prescribes an optimal β that adapts to each WSI’s tumor fraction—ideally varying at every weight update—but we instead pick a fixed β via a rule‐of‐thumb derived from overall dataset statistics. This is suboptimal given the wide slide‐to‐slide variation in cancer burden (e.g., Fig. 6’s nearly fully malignant CAM17 example). To address this, we propose a simple heuristic that estimates tumor fraction online from the model’s own (initially noisy) predictions and adjusts β according to our theoretical equations. Though early‐epoch instability and alternative strategies remain possible, this per-slide β selection is a concrete MaxSoft limitation we’ll explicitly discuss in the final manuscript.

[1] Liu D., Dual-Attention MIL Framework [2] Tang W., MIL with Masked Hard Instance Mining [3] Shao Z., TransMIL [4] Zhang H., DTFD-MIL [5] Bontempo G., DAS-MIL [6] Qu L., DGMIL [7] Athalye A., Obfuscated Gradients [8] Yin P., Straight-Through Estimator [9] Recht B., Do ImageNet Classifiers Generalize? [10] Taori R., Robustness to Natural Distribution Shifts [11] Miller J., Accuracy on the Line [12] Hendrycks D., The Many Faces of Robustness [13] Juyal D., SC-MIL: Supervised Contrastive MIL [14] Qiu P., SC-MIL: Sparsely Coding MIL [15] Yang Z., SCMIL: Sparse Context-aware MIL

Thanks for your response. I will keep my accept rating.

Rebuttal For KwFs

We thank the reviewer for the constructive feedback. We appreciate your recognition of our theory-driven methodology, the intuitiveness of PerPatch, and the breadth of our empirical study. We also value the suggestion to apply PerPatch within alternative MIL pipelines and regard this as a promising future direction. We note that our theoretical section introduces new theorems for the underexplored Log-Sum-Exp (LSE) pooling, which we believe will benefit the broader ML community. We respond to each point below and will incorporate the feedback in the revision.

Owing to space constraints, we use “transformer-based” to encompass both transformer- and attention-based approaches. We likewise omit rarely cited metrics and standard deviations, consistent with prevailing practice in the literature.

Initially, we observe that the Questions section repeats the Weaknesses points exactly. We were not sure that this was intentional by the reviewer. We are open to discuss any additional questions to improve our work and really appreciate it.

W1&Q1&L1: Motivation underdeveloped.

We thank the reviewer for raising this concern and apologize if our exposition made the motivation appear incomplete. Below, we restate, without introducing new claims, the body of evidence already present in the paper.

Well-documented data scarcity in MIL. Numerous MIL studies are driven by the scarcity of labeled slides; in particular, [1] and [2] are motivated enuclearntirely by sample scarcity and both explicitly observe that transformer-based MIL is most vulnerable. More broadly, vision transformers are repeatedly characterized as data-hungry because of their weak inductive bias [3–5].

ViT as an MIL model. Removing positional embeddings converts a ViT into a set-based MIL architecture. The MoCoV3 ablation in [6] trains such a “ViT-MIL” on the full 1.28 M-image ImageNet-1K yet still observes the same data-hungriness issue, showing that the limitation is architectural rather than dataset-specific.

Our empirical evidence. Figure 9 demonstrates that test accuracy for ABMIL, DSMIL, Snuffy, and Prov-GigaPath rises monotonically with slide count, closely tracking the slope in Fig. 3 of [3]. This pattern replicates across all four datasets we evaluate (CAM16, CAM17, SICAP, and TCGA).

The reviewer proposes several alternative explanations, each of which we tested:

Encoder type. Tables 1 & 3 evaluate four encoders: DINO-Natural (ImageNet-1K), DINO-Domain (pretrained on the training split patches, e.g., CAM16), UNI (100 M pathology patches), and Prov-GigaPath (1.3 B pathology patches). Across every dataset, MaxSoft pooling surpasses transformer pooling, ruling out encoder capacity or domain mismatch.

Data quality. Three of the four datasets were curated under stringent, pathologist-led QC protocols [6-9]; visual inspection confirms their clarity. Consistent with the broader CV literature, however, quantity—not quality alone—enhances transformer hunger (e.g., LAION-2B in [9]), so the gap remains.

JPEG fidelity. Digital-pathology pipelines typically store tiled patches as JPEG Q = 75. We therefore generated three new variants—Q = 50, Q = 100, and Q = 75 + Gaussian blur—to test whether artefacts influence sample efficiency. The accompanying table (Table S4) shows an unchanged ordering, demonstrating that neither lowering nor raising quality affects the conclusion.

Training protocol & domain alignment. All models are trained from scratch via standard i.i.d. sampling and evaluated on held-out test splits drawn from the same dataset; DINO-Domain is trained exclusively on in-domain patches (Section 5, Appendix D, and line 163). The transformer gap persists under these strictly matched conditions.

Pretraining scale. UNI and Prov-GigaPath are foundation encoders trained on 100 M and 1.3 B patches, respectively. We finetuned the Prov-GigaPath MIL head (pretrained on ≈200,000 slides)—originally trained on **≈**200 K slides—on SICAP-MIL, yet it still underperforms MaxSoft, indicating that present-day slide counts remain insufficient for transformer MIL on WSIs.

Collectively, these observations validate our original claim: under current pathology data regimes, transformer-based MIL is data-inefficient, whereas MaxSoft delivers robust gains under the same conditions.

W2&Q2: Data-efficiency untested.

Figure 9 directly addresses data efficiency. Under limited training samples, at least one simple pooling (Max, Mean, LSE, or MaxSoft) surpasses transformer-based MIL; specifically, Mean pooling leads at 50 slides, while MaxSoft dominates from ≥100 slides.

W3&Q3: Missing non-transformer baselines.

We appreciate the request to suggest graph MIL models; however, our study deliberately excluded them because their current formulations introduce well-documented limitations that render them unreliable for WSI classification. As noted in the cited survey, two persistent challenges are “graph-structure definition” and “dependence on segmentation output.” Most pipelines define each nucleus as a node and connect only spatially adjacent pairs, which raises three issues:

1. reliance on a separate nuclei-segmentation model discards information whenever segmentation errs;
2. the distance threshold for edge creation lacks a biologically motivated rationale; and
3. nucleus-centric graphs omit non-nuclear cues—e.g., tissue morphology—that pathologists routinely use.

Methods that first encode patches with ViTs and then build a graph [10] are essentially transformer-based models in disguise. Theoretical analyses show that a ViT is a message-passing GNN on a fully connected token graph, while a Message Passing Neural Network (MPNN) with a virtual node can approximate self-attention; thus, such hybrids are functionally equivalent to transformer MIL [11–13].

Nonetheless, to address your concern we added GTP [13]—one of the strongest graph baselines—to the Additional Baselines table. Its results confirm our main conclusions. We also note a practical drawback: GTP constructs an explicit adjacency matrix, incurring O(N²) memory per slide and frequently triggering CUDA out-of-memory errors (cf. GitHub issue #25, opened 20 Mar 2025, unresolved to date).

W3: Missing recent MIL baselines.

We thank the reviewer for highlighting additional MIL baselines. Our study already benchmarks ABMIL and DSMIL—two widely cited transformer-based MIL methods—and Snuffy, a current SOTA whose Tables 1–2 show clear gains over TransMIL (ref. 8). Of the suggested non-graph papers (refs. 3, 4, 7, 8), ref. 3 is unrelated to MIL, WSI, or pathology; refs. 4 and 7 propose prototype-based morphology models and ref. 4 explicitly state it is not MIL, with ref. 7 reporting very weak CAM16 results at 20×—the magnification we use (l. 525); and ref. 8 is another transformer-style MIL architecture. The “Additional Baselines” table already corroborates our findings. As we were not sure which one of the above methods picked the reviewer interests due to explained reasons, if the reviewer deems any further method essential, we would appreciate if they explain why it was necassary, and also we will gladly include it.

Additional Baselines: CAM16 and CAM17, TCGA, and SICAP

Finally, we noticed that reference [1] is cited but not discussed in your review; please let us know if there are specific points from that work you would like us to address.

[1] Zhang H., DTFD-MIL

[2] Liu P., Pseudo-Bag Mixup

[3] Dosovitskiy A., An Image Is Worth 16×16 Words

[4] Beyer L., Better Plain ViT Baselines

[5] Tolstikhin I., MLP-Mixer

[5] Chen X., Empirical Study of Self-Supervised ViTs

[6] Bejnordi B., DL for Lymph Node Metastases Detection

[7] Bandi P., Camelyon17 Challenge

[8] Cooper L., Pancancer Insights from TCGA

[9] Shao Z., TransMIL

[10] Bontempo G., DAS-MIL

[11] Joshi C., Transformers Are Graph Neural Networks

[12] Cai C., On the Connection Between MPNN and Graph Transformer

[13] Zheng Y., Graph-Transformer for WSI Classification