Iterative Patch Selection for High-Resolution Image Recognition
Authors: Benjamin Bergner
Christoph Lippert
Aravindh Mahendran
Venue: ICLR 2023 poster (2023)
Keywords: high-resolution images
memory-efficient deep learning
multiple instance learning
transformer
image recognition
computer vision

Abstract: High-resolution images are prevalent in various applications, such as autonomous driving and computer-aided diagnosis. However, training neural networks on such images is computationally challenging and easily leads to out-of-memory errors even on modern GPUs. We propose a simple method, Iterative Patch Selection (IPS), which decouples the memory usage from the input size and thus enables the processing of arbitrarily large images under tight hardware constraints. IPS achieves this by selecting only the most salient patches, which are then aggregated into a global representation for image recognition. For both patch selection and aggregation, a cross-attention based transformer is introduced, which exhibits a close connection to Multiple Instance Learning. Our method demonstrates strong performance and has wide applicability across different domains, training regimes and image sizes while using minimal accelerator memory. For example, we are able to finetune our model on whole-slide images consisting of up to 250k patches (>16 gigapixels) with only 5 GB of GPU VRAM at a batch size of 16.

Published as a conference paper at ICLR 2023
ITERATIVE PATCH SELECTION FOR HIGH-RESOLUTION IMAGE RECOGNITION
Benjamin Bergner', Christoph L Aravindh Mahendran*
Hasso Plattner Institute for Digital Engineering, University of Potsdam Haaso Plattner Institute for Digital Health the Icahn School of Medicine Mount Sinai Google Research, Brain Team {fir (firstname. 2 w aravindhmmgpogge..cmm
ABSTRACT
High-resolution images are prevalent in various applications, such as autonomous driving and computer-aided diagnosis. However, training neural networks on such images is computationally challenging and easily leads out-of-memory errors even on modern GPUs. We propose a simple method, Iterative Patch Selection (IPS), which decouples the memory usage from the input size and thus enables the processing of arbitrarily large images under tight hardware constraints. IPS achieves this by selecting only the most salient patches, which are then aggregated into a global representation for image recognition. For both patch selection and aggregation, a cross-attention based transformer is introduced, which exhibits a close connection to Multiple Instance Learning. Our method demonstrates strong performance and has wide applicability across different domains, training regimes and image sizes while using minimal accelerator memory. For example, we are able to finetune our model on whole-slide images consisting of up to 250k patches (> 16 gigapixels) with only GB of GPU VRAM at batch size of 16.
INTRODUCTION
Image recognition has made great strides in cent years, spawning landmark architectures such as AlexNet (Krizhevsky et al. 2012) or ResNet (He et al. 2016). These networks are typically designed and optimized for datasets like ImageNet (Rus- sakovsky et al. 2013). which consist of natural im- ages well below one megapixel In contrast, real- world applications often rely on high-resolution im- ages that reveal detailed information about an ob- ject of interest. For example, in self-driving cars, megapixel images are beneficial to recognize distant traffic signs far in advance and react in time (Sahin) 2019). In medical imaging, a pathology diagnosis system has process gigapixel microscope slides recognize cancer cells, as illustrated in Fig.
Training neural networks high-resolution images is challenging and can lead to out-of-memory errors even on dedicated nigh-performance hardware. Al- though downsizing the image can fix this problem, details critical for recogmition may be lost the pro- cess (Sabottke & Spieler 2020 Katharopoulos & Fleuret] 2019). Reducing the batch size is another common approach to decrease memory usage, but it does not scale to arbitrarily large inputs and may lead to instabilities in networks involving batch normalization (Lian & Liu 2019). On the other hand, distributed learning across multiple devices increases resources but is more costly and incurs higher energy consumption (Strubell et al. 2019).
Figure 1: IPS selects salient patches for high- resolution image recognition. In the top right insert, green contours represent ground truth cancerous cells and white overlays indicate high scoring patches from our model. Exam- ple patches are shown in the bottom right.
For instance, 256 image corresponds only 0.06 megapixels.

---

Published as a conference paper at ICLR 2023
We propose Iterative Patch Selection (IPS), simple patch-based approach that decouples the con- sumed memory from the imput size and thus enables the efficient processing of high-resolution im- ages without rumning out of memory. IPS works in two steps: First, the most salient patches of an image are identified in no-gradient mode. Then, only selected patches are aggregated to train the network. We find that the attention scores of a cross-attention based transformer link both of these steps, and have close connection to Multiple Instance Learning (MIL).
In the experiments, we demonstrate strong performance across three very different domains and training regimes: traffic sign recognition on megapixel images, multi-task classification on syn- thetic megapixel MNIST digits, and using self-supervised pre-training together with our method for memory meemor--effiiient learning on the gigapixel CAMELYONIC 1666 benchmark. Furthermore, our method exhibits significantly lower memory consumption compared to various baselines. For example, when scaling megapixel MNIST images from 1k to 10k pixels per side at batch size of 16, we can keep peak memory usage at a constant 1.7 GB while maintaining high accuracy, in contrast to a comparable CNN, which already consumes 24.6 GB at resolution of 2kx In an ablation study, we further analyze and provide insights into the key factors driving computational efficiency in IPS. Finally, we visualize exemplary attention distributions and present approach to obtain patch-level class probabilities in a weakly-supervise multi-label classification setting.
2
METHODS
We regard an image as a set of N patches. Each patch is embedded independently by a shared encoder network. resulting in D-dimensional representations, X R.NDD, Given the embeddings, we select the most salient patches and aggregate the information across these patches for the classifi- cation task. Thus method, illustrated in Fig. [ consists of two consecutively executed modules: an iterative patch selection module that selects a fixed number of patches and a transformer- patch aggregation module that combines patch embeddings compute a global image embedding that is passed on to a classification head. Crucially, the patch aggregation module consists of a cross- attention layer, that is used by the patch selection module in no no-gradient mode to score patches. We discuss these in detail next and provide code at htttp:://gitai htt.e//ww...coobbbberneeee/////p
2.1 ITERATIVE PATCH SELECTION
Given infinite memory, one could use an attention module to score each patch and select the top M patches for aggregation. However, due to limited GPU memory, one cannot compute and store all patch embeddings in memory at the same time. We instead propose to iterate over patches, I at a time, and autoregressively maintain set of top M patch embeddings. In other words, say P%r is buffer of M patch embeddings at time step and pf+1 are the next patch embeddings in the autoregressive update step. We run the following for T iteeations:
p++1 ToP-M(PKs U P |att),
(X,.... XM} is the initial buffer of embeddings X ((I....M) and ++1 E M+1 are attention scores of considered patches at iteration t 1. based on which the selec- tion in Top-M is made. These attention scores are obtained from the cross-attention transformer as described in Sect. 2.2 The output of IPS is a set of M patches corresponding to Pf- Note that both patch embedding and patch selection are executed in no-gradient and evaluation mode. The former entails that no gradients are computed and stored, which renders IPS runtime and memory-effiient. The latter ensures deterministic patch selection behavior when using BatchNorm and Dropout.
where T -
Data loading We introduce three data loading strategies that trade off memory and runtime effi- ciency during IPS. In eager loading, a batch of images loaded onto the GPU and IPS is applied to each image in parallel-this is the fastest variant but requires storing multiple images at once. In eager sequential loading, individual images are loaded onto the GPU and thus patches are selected for one image at time until batch of M patches per image is selected for training. This enables the processing of different sequence lengths without padding and reduces memory usage at the cost of higher runtime. In contrast, lazy loading loads batch of images onto CPU memory. Then, only patches and corresponding embeddings pertinent to the current iteration are stored on the GPU-this decouples GPU memory usage from the image size, again at the cost of a higher runtime.
2

---

Published as a conference paper at ICLR 2023
Iterative Patch Selection (No-Gradient Mode)
2
Transformer
Scores
Class Head
Top-M
Top-M
Attention Scorer
Attention Scorer
Attention
Add Norm
Scorer
Shared weights
MLP
Patch
Embeddings
Add & Norm
Memory
Ps
Pir
Pi
Multi-Head
Cross-Attentinn
Input Image
Embed
K
iteration
Query token
Selected Patches by
Pb
P7
P2
Pf
IPS (re-embedded))
Figure 2: Left: Illustration of Iterative Patch Selection running for 3 iterations with N = 16. M = and / 4. Right: Cross-attention transformer module that aggregates patches selected by IPS.
2.2 TRANSFORMER- BASED PATCH AGGREGATION
After selecting M patches in the IPS stage, these patches are embedded again in gradient and training mode, and the patch aggregation module can aggregate the resulting embeddings X Mx using convex combination:
M
am Xm
m=1
Each attention score am is the result of function of the corresponding patch embedding X M) Softmax (Ji()) (x;;)) M)) and can be learned by neural network f parame- terized by B. We observe that the weighted average in Eq. 2 constitutes a cross-attention layer by defining f follows:
QK"
8 (X
VD
where queries 8 q and keys K= W linear transformations of learnalee query and embedding X respectively, with weights W4 W E IR Dx and Dj being the projected dimension. Furthermore, instead of aggregating patch embeddings directly, we aggregate a linear projection called values in typical transformer notation: V =x; W", with W R Dx and D, being the value dimension. In practice, we use multi-head cross- artss-attention and place it in standard transformer (Vaswani et al. 2017). optionally adding DD sinusooia position encoding (details in Appendix B). Note that there is a close connection between the patch selection and patch aggregation modules. They both share the cross-attention layer. Whereas IPS runs in no-gradient mode and cannot train its parameters, by sharing them with the patch aggregation module, one still learn select patches relevant for the downstream task.
token Φº E
3 RELATED WORK
The setting described in Sect constiitess a MIL problem (Maron & Lozano- Pérez) 1997). where an image is seen as a bag that contains multiple patch-level features called instances. The goal is then to predict the label of unseen bags given their instances. However, training labels are only available for bags, not for instances. The model must thus learn to attend to patches that best represent the image for given task. Eq. provides such mechanism and is known as MIL pooling function. In particular, our cross-attens cross-attention layer follows the weighted collective assumption stating that all instances contribute the bag representation, albeit to varying degrees according to the attention scores (Foulds Frank 2010 (mmorss) 2013). A MIL approach similar to ours is DeepMIL (Use al. 2018), which gated attention mechanism compute scores for the aggregation of N
3

---

Published as a conference paper at ICLR 2023
patches: (X,) w(tann(X,,, O sigm(X, V))", with U, V E RDDx and w R being parameters, and D,, hidden nodes. In contrast, our multi-head cross-attention layer provides more capacity by computing multiple intermediate bag representations. Furthermore. it is MIL pooling function that can be naturally integrated into a standard and field-proven transformer module.
Other MIL methods common for binary classification of large images classify each patch and use max- max-ppooing or a variant thereof on the class scores for patch selection and training (Campanella ct al., 2018 Zhou et al., 2018, Lerousseau et al. 2020). Our method is more general and can be used for multi- ttask, label, class} classification while being able to attend to multiple patches. In early experiments, we found that max-pooling in the baseline leads to overfitting. Instead, we employ TopMIL (based on [Campanella et al. (2019)) as a baseline, which selects the top M patches in a single iteration according to the maximum class logit of each patch. Selected patches are then re- embedded, classified, class logits are averaged, and an activation function (e.g., softmax) is applied.
Another line of work seeks to reduce memory usage for the processing of high-resolution images by first identifying regions of interest in low resolution and then sampling patches from these regions in high resolution (Katharopoulos & Fleuret, 2019 Cordonnier et al., 2021). This is also related to early attention approaches that rapidly process the entire image based on simple features (Culhane & Tsotsos| 1992 Koch & Ullman| 1987). as well as to human attention where saceadic eye movements may be informed by peripheral vision Burt 1988). However, especially with gigapixel images, low- resolution images may either still be too large to encode, or too small to identify informative regions. In contrast, IPS extracts deep features from high-resolution patches to estimate attention scores.
For very large images, patch- wise self-supervised learning is a promising way to reduce memory usage (Dehaene et al. 2020 Chen al. 2022 Li al. 2021). After pre-training, low-dimensional features are extracted and used to train a downstream network. Although the input size decreases considerably in this way, may still require too much memory for very long sequences. Pre-training can be used orthogonally to IPS to process sequences of any length, as demonstrated in Sect. 4.2
Our method is related to Sparse Mixture of Experts, a conditional computation approach in which specialized subnetworks (experts) are responsible for the processing of specific input patterns (Ja-) cobs al., 199T). Specifically, [Shazeer et al. (2017) and Riquelme et al. (202T) use a gating func- tion to decide which expert is assigned to which part of the input. This gating function is non- differentiable but shares its output with a differentiable aggregation operation that combines the expert's 's outputs. On a high level this is similar to how we share the cross-attention module between no no-gradient mode IPS and the gradient mode patch aggregator. Attention Sampling (Katharopoulos & Fleuret 2019) also reuses aggregation weights for patch sampling. IPS uses Top-M instead.
4 EXPERIMENTS
We evaluate the performance and efficiency of our method on three challenging datasets from a vari- ety of domains and training regimes: Multi-class recognition of distant traffic signs in megapixel im- ages, weakly-supervised classification in gigapixel whole-slide images (WSI) using self-supervised representations, and multi-task learning of inter-patch relations on synthetic megapixel MNIST benchmark. All training and baseline hyperparameters are provided in Appendix D
Baselines To show that IPS selects salient patches, two other patch selection methods, Random Patch Selection (RPS) and Differential Patch Selection (DPS) (Cordonnier et al. (021). used for comparison, both employing the same transformer as IPS. Next, to verify that the cross-altention transformer is strong MIL pooling operator, it is compared DeepMIL (Ilse et al., 2018 and TopMIL as introduced in Sect. 3, Finally, we compare to standard CNN applied on the original image and smaller sized version to assess how resolution affects performance.
Metrics Each of the main experiments is performed five times with random parameter initializa- tions, and the average classification performance and standard deviation are reported on the test set using either accuracy or AUC score. We also report computational efficiency metrics, including maximum GPU memory usage (VRAM) and training runtime for a batch size of 16 and various values of M. The runtime is calculated for single forward and backward pass and is averaged over all iterations of one training epoch excluding the first and last iterations. Both metrics are calculated on single NVIDIA A 100 GPU in all experiments.
4

---

Table Published as a conference paper at ICLR 2023
4.1 TRAFFIC SIGNS RECOGNITION
We first evaluate our method on the Swedish traffic signs dataset, which consists of 747 training and 684 test images with 1.3 megapixel resolution, as in [Ktharoppollos & Fleuret (2019). Each image shows speed limit sign of 50, 70, 80 km/h or no sign. This problem re- quires high-resolution images to dis- tinguish visually similar traffic signs, some of which appear very small due to their distance from the cam- era (Fig. [5] left). We resize each image to 1200 1600 and extract 192 non-overlapping patches of size 100x 100. Due to the small num- ber of data points, a ResNet-18 with ImageNet-Il Ik weights is used as en- coder for all methods and then fine- tuned. For IPS, we use eager loading and set I = 32.
Table I: Traffic sign recognition: Our IPS transformer out- performs all baselines using 1.2 GB of VRAM while being faster than the CNN. DPS uses = 768 patches (see Ap- pendix D.5). VRAM is reported in gigabytes and time in milliseconds.
M
Acc. [%]
VRAM
Time
IPS Transformer
192
98.4 +0.5
14.2
323
(ours)
10
98.6
#0.5
1.3
126
2
98.5
+0.8
1.2
114
RPS Transformer
10
50.7 #2.6
1.3
31
2
28.3 #3.2
0.9
26
DPS Transformer
10
94.0 1:.2
3.5
337
2
97.2 40.7
2.5
296
TopMIL
10
97.7 +0.5
4.5
109
TopMIL
2
98.2 40.5
4.5
96
DeepMIL
192
96.2 +1.6
14.2
323
DeepMIL+
192
97.7 11.3
14.3
323
CNN
84.1 #4.7
13.4
231
CNN 0.3x
76.3
10.9
1.4
83
Table D shows that the average accu- racy of the IPS transformer is consis- tently high regardless of the value of M. For M (2, 10). memory usage is more than 10x less than for the CNN, and for M = 2, training time is about half that of the CNN. As expected, RPS does not select salient patches due the low signal-to-noise ratio in the data. DPS localizes the salient patches but performs worse than IPS in all metrics and is sensitive to the choice of M. The MIL. baselines perform well in general, and TopMIL comes closest to IPS in accuracy. However, TopMIL requires knowledge about the number of informative patches, since the logits are already for . = 10, the performance is slightly degraded. In addition, TopMIL. is faster but less memory- efficient than IPS because in TopMIL all patches are scored at once. DeepMIL, on the other hand, lacks patch selection mechanism and hence includes all patches in the aggregation, which increases training time and VRAM. For M 192, DeepMILL differs from the IPS transformer only in the ag- gregation function, yet the transformer performs 2.2 percentage points better on average. We also report results for DeepMIL+, which has similar capacity to the transformer for fairer comparison. achieves an average accuracy of 97.7%, less than 98.4% of IPS. A CNN that takes the entire image as input and has no attention mechanism performs much worse and is inefficient. Downscaling the image by a factor of 3 further degrades the CNN's performance, as details of distant traffic signs are lost.
For a fair comparison against prior work, we used the same transformer and pre-trained model where applicable. This results in better numbers than reported by DPS, which achieves 97. 97.2% in vs. 91.7% in Cordonnier et al. (2021). Nonetheless, our IPS transformer improves on that by up to 1.4 percentage points while being more memory and runtime efficient.
4.2 GIGAPIXEL WSI CLASSIHICATION
Next, we consider the CAMELYONI6 dataset (Litjens al. 2018), which consists of 270 training and 129 test WSIs of gigapixel resolution for the recognition of metastases in lymph node cells. The task formulated as a seakly-supervised binary classification problem, i.e., only image-level labels about the presence of metastases are used. The WSIs come with multiple resolutions, but we attempt to solve this task with the highest available resolution (40x magnification). Since a large portion of the images may not contain any cells, Otsu's method (Otsul 1979) is first applied to filter out the background. Only patches with at least 1% of the area being foreground are included a sensitive criterion that ensures that abnormal cells are left out. Then, non-overlapping patches of size 256 x 256 are extracted. This gives a total of 28.6 million patches (avg.: 71.8k, std.: 37.9k), with the largest WSI having 251.8k patches (about 20% of the size of ImageNet-1k).
5

---

Published as a conference paper at ICLR 2023
Table 2: CAMELYONI6: IPS achieves performance close to the strongly-supervised state-of-the-
art. t: Average number of patches is reported for these methods, NR: N and M are unknown.
M
Total (N)
Pretr.
AUC [%]
VRAM [GB]
Time [ms]
IPS Transformer (ours)
5k
all
BYOL
98.1 10.3
4.7
355
1k
all
BYOL
97.5 40.2
4.1
313
100
all
BYOL
97.3 #0.6
4.0
313
5k
70k
BYOL
95.1 #1.7
2.9
286
DIe
10k:
10k
BYOL
84.1 #1.9
4.5
26
50k:
50k
BYOL
93.8 #1.4
22.5
110
70k:
70K
BYOL
94.5-111
31.5
150
TopMIL
5k
70K
BYOL
71.8 12.0
84
1k
70K
BYOL
76.2 #2.9
19.8
76
100
70k
BYOL.
84.4 #6.9
19.8
74
DSMIL-LC [L. ct al.
2021
SA
8&1
SimCLR
91.7
CLAM (Luu et al. 2021
42k
42
CPC
93.6
CLAM-SH Wang ef al.
2022
NR
NR
SRCI.
94.2
DeepMIL Dehaene et
2020
10A:
10k
MoCo v2
98.7
Challenge Winner Wang et
2016
None
92.5
Strong superv. (Bejnordi el al.
2017
None
99.4
Due to the large volume of patches, it is impractical to learn directly from pixels. Instead, we train BYOL. (Grill al., 2020), a state-of-the-art self-supervised learning algorithm, with a ResNer-50 encoder and extract features for all patches and WSIs (see Appendix [.6] for details). Each patch is thus represented as a 2.048-dimensional vector, which is further projected down to 512 features and finally processed by the respective aggregator. Note that even when processing patches at the low-dimensional feature level, aggregating thousands of patch features can become prohibitively expensive. For example, with DeepMIL, we aggregate around 70k patch features before running out of memory on an A100 at a batch size of 16. Therefore, in each of the runs, a random sample of up to 70k patch features is drawn for the baseline MIL methods, similar to Dehaene al. (2020). In contrast, IPS can easily process all patches. Due to the varying number of patches in each WSI, batch cannot be built without expensive padding. We thus switch to eager sequential loading in IPS, i.e. slides are loaded sequentially and M patch features per image are selected and cached for subsequent mini-batch training. all IPS experiments,
5k.
We report AUC scores, as in the challenge (Bejnordi et al., 2017). in Table 2 Best performance achieved by IPS using M 5k, with an improvement of more than 3 percentage points over the runner-up baseline. We also outperform the CAMELYON16 challenge winner (Wang et al.) 2016), who leveraged patch-level labels, and approach the result of the strongly-supervised state of the art (Bejnordi et al. (017). The memory usage in IPS is 4-5 GB depending M. However, IPS has a higher runtime because all patches are taken into account and eager sequential loading is used. For comparison, (using the same random seeds) and observe slightly better performance than DeepMIL. DPS was not considered since the downscaled image exceeded our hardware limits. We also tested CNNs trained on downscaled images up to 4k pixels but they perform worse than a naive classifier.
also reduce the total number of patches before the IPS stage to 70k
Several recent works adopt pre-trained feature extractor and report results for CAMELYON16 (see bottom of Table 2 Given their respective setups, most consider all patches for aggregation but make compromises that one can avoid with IPS: A batch size of (Wang al. 2022 Li al., 2021), training with features from the third ResNet-50 block, which effectively halves the size of the embeddings (Lu et al.) 2021), or using smaller resolutions (Dehaene al. 2020 Liet al. 2021). The work by Dehaene et al. (2020) achieves a slightly better result than our IPS transformer (avg. AUC: 98.1% vs. 98.7%) by using DeepMIL on a random subset of up to 10k patches. We do not assume that this is due to our transformer, as the DeepMIL baseline reproduces their aggregation function yet performs significantly worse. Instead, it could be due a different pre-training algorithm (MoCo v2), pre- pre-processing (closed-souree U-Net to filter out the background), lower magnification level (20x), which reduces the total number of patches (avg. 9.8k vs. 71.8k in ours).

---

Published as a conference paper at ICLR 2023
Table 3: Results for megapixel MNIST. The IPS Transformer accurately solves all relational tasks with a minimal number of patches and can leverage overlapping patches with small memory usage.
Total (N)
Majority
Max
Top
Multi-label
VRAM [GB]
Time [ms]
IPS Transformer
900
3,481
99.8 10.1
99.0 10.2
95.0 =11
93.6 13.3
15.5 GB
57U ms
(ours)
100
3,481
99.7 10.1
98.8 10.2
95.9
10.8
96.1 10.7
2.2 GB
326 ms
5
3,481
97.5 #1.0
98.4 #0.1
96.7 =0.5
90.2 #1.0
1.1 GB
301 ms
900
900
99.1 10.2
93.7 10.5
91.7
11.4
82.5 13.0
15.0
314 ms
900
98.9 +0.3
94.1 10.8
92.3
#0.7
84.2 12.0
1.8 GB
114 ms
900 5
94.8 10.7
92.4 10.8
91.0
72.1 13.0
0.7 GB
89 ms
RPS Transformer
900
3,481
82.5 11.6
78.5 1:199
77.6 110
47.0 12.4
15.4 GB
314 ms
100
3,481
33.8 #1.3
34.9 #1.3
29.1 #1.5
0.9 =0.3
2.2 GB
47
DPS Transformer
100
900
97.2
93.2 10.4
92.3 1:066
81.4 L1S
GB
353
5
900
88.1 #3.1
81.7 +2.6
80.1 #3.6
48.3 +:577
2.2 GB
285
TopMIL
900
900
98.2 #0044
90.0 #1.0
62.9 #0.66
71.9 +4.00
GB
381
100
900
98.4 1:0.3
92.4 10.3
53.8 11.0
78.0
4.6 GB
105
5
900
97.2 +0.6
91.4 #0.6
71.2 =1.0
72.3 #1.7
4.6 GB
73 ms
DeepMIL
900
900
98.8 4.0.2
93.8 10.2
92.3
4.09
80.7 4:1.2
GB
316 ms
CNN
99.3 10.4
91.9 10.6
57.0 11.1
71.7 19.0
13.7 GB
209
CNN 0.3x
90.8 +0.7
76.7 +1.3
55.8 =0.4
42.3 #2.66
1.6 GB
30 ms
4.3
INTER-PATCH REASONING
A single informative patch was sufficient to solve the previous tasks. In contrast, megapixel MNIST
introduced in Katharopoulos & Fleuret (2019) requires the recognition of multiple patches and their relations. The dataset consists of 5,000 training and 1,000 test images of size 1,500 1,500. We extract patches of size 50x50 without overlap (N 900) or with 50% overlap (N 3,481). In each image, 5 MNIST digits are placed, 3 of which belong to the same class and 2 of which to other classes. In addition, 50 noisy lines are added (see Fig. S right). The task is then to predict the majority class. We found that this problem can be solved well by most baselines, so we extend the setup with three more complex tasks: max: detect the maximum digit, top: identify the topmost digit,
the presence/absence of all classes needs to be recognized. We frame this as a multi-task learning problem and use learnable query tokens in the cross-attention layer, i.e. task representations are learned. Similarly, 4 gated-attention layers are used for DeepMII. All methods also utilize multiple classification heads and add a 11) -inussisaa position encoding to the patch features. For the CNNs, a sinusoidal channel is concatenated with the input. A simplified ResNet- 18 consisting of residual blocks is used as the encoder for all methods. For IPS, we use eager loading and set I = 100.
The results in Table show that almost all methods were able to obtain high accuracies for the tasks majority and max. However, only DeepMIL. DPS and IPS were able to make appropriate use of the positional encoding to solve task top. For the hardest task, multi-label, only the IPS transformer obtained accuracies above 90% by using overlapping patches. However, for both TopMIL. and Deep-
MIL, we were not able to use overlapping patches due to out-of-memory errors (on a single A A100). Interestingly, IPS achieves high performance throughout with M = 100 using overlapping patches at a memory usage of only 2.2 GB while maintaining a competitive runtime. Furthermore, the IPS transformer also achieves high performance with M 5 patches, which is the minimum number of patches required to solve all tasks.
4.4 ABLATION STUDY
Effect of image size How does input size affect performance and efficiency in IPS? To address this question. we generate new megapixel MNIST datasets by scaling the image size from 1k to 10k pixels per side (400 to 40k patches) and attempt to solve the same tasks as in the main experiment. We also linearly scale the number of noise patterns from 33 to 333 to make the task increasingly difficult. Fig. 3 shows a comparison of memory usage, runtime and accuracy (task majority) for different data loading options of our IPS transformer (M- = 100, I 100, non-overlapping patches, default batch size of 16), as well as the baseline CNN.
7

---

Published as a conference paper at ICLR 2023
Eager loading
1.00
Poak G2C Menuur (UB)
loading
Eapor seg. Joading
Bustime [ms]]
Actumay [84)
15
CNN
CNN batch size /4
20>
80
buger
Lerr loading
Bager seq. loading
IPS Transformer
10
UNN
20
CNN
hateis
12
25
3k
38
6k
7X
86
9k:
10k
1k
3k
6k
53
7k
5k 9k.
10k
1k
26
4k 56
7k 86
10
Width [pcc
Width/HHeigtt [pre]
Wilthh/Higgh [pe]
Figure 3: Peak GPU memory (left), training runtime (center) and accuracy (right) as function of image size megapixel MNIST: With IPS and lazy loading, memory usage can be kept constant. High performance is mostly maintained compared to the baseline CNN.
250
Punk canu: Munocy (C30)
slae
200
Number of terrattons
[ms]
150
size
100
60
25
25
-0-
50
50
100
100
20-
all
Mumory
Memory
Memory
Figure 4: VRAM, training runtime and number iterations for different combinations of M and J.
The CNN exhibits high memory footprint that quickly exceeds the hardware limits even when using a reduced batch size In IPS with eager loading, the VRAM is dominated by the batch of inputs and can be reduced by sequential loading. In contrast, with lazy loading the memory usage kept constant at 1.7 GB regardless of the input size. The runtime
eager loading is less than that of the CNN and almost approaches CNN using batch size of 4. Lazy loading is faster than eager sequential loading, suggesting that the latter should only be used situations where different sequence lengths need to be processed. In terms of performance, achieves high accuracy for up to 8k pixels and decreases only slightly starting from 9k pixels likely due to the declining iigise-t-noisee ratio (at 10k pixels, only 5 out of 40k patches relevant). In comparison, the CNN performance drops rapidly starting from 3k pixels towards a naive classifier.
note that the IPS transformer
Effect of M and Fig. shows that VRAM and runtime tend to increase with M as more patches need to be processed in training mode. With increasing values of I, a higher VRAM but lower runtime can be observed, since IPS runs for fewer iterations. However, there are more subtle cases. Along the diagonal in Fig. g (left), peak VRAM is observed in the non-IPS part and thus remains constant for different values of I. For example, for M = 50, VRAM remains constant for / E {25, 50, 100). so setting I 100 results in faster runtime without compromising memory usage. In contrast, the upper left triangle the graph, the VRAM dominated by IPS in no-gradient mode. For example, = all uses 25GB of VRAM for M 10 (i.e., I 890). Interestingly, when processing all remaining patches in a single iteration (I = - M), VRAM can be reduced by increasing M up to 200, because is reduced by M.
Effect of patch size In Table 4 we investigate the effect of patch size (25-400 px per side) on the performance and efficiency of traffic sign recognition. One can notice that the accuracy is similar for all patch sizes and only slightly decreases for very small patch sizes of 25 size increases, the runtime for IPS tends decrease since fewer iterations are run, while the time to embed and aggregate selected patches (non-IPS time) increases because more pixels need to be processed in training mode.
As the patch
4.5 INTERPRETABILITY
In Fig. [] we superimpose attention images, which shows that salient patches tended to while the background is neglected. In low signal-to-noise settings, such visualizations
8

---

Published as a conference paper at ICLR 2023
Table 4: Effect of patch size traffic sign recognition and efficiency (M 2,
32, single run).
Size [px]
Acc. [%]
VRAM [GB]
IPS Time [ms]
Non-IPS Time [ms]
Total Time [ms]
96.7
0.9
417
423
50
98.7
1.0
167
174
98.8
1.2
104
10
114
200
99.1
3.3
74
15
89
400
98.4
4.2
88
49
137
Figure 5: Patch scores and crops of most salient patches. Left: Traffic signs obtain high scores. Cen- ter: Most cancer cells (black contour) are attended. Right: Patch-level class mappings in MNIST.
can support humans to localize tiny regions of interest that might otherwise be missed (e.g., Fig. However, attention scores alone may not be satisfactory when multiple objects of different classes are present. For example, attention scores in megapixel MNIST indicate relevant patches, but not to which classes these patches belong. We propose simple fix to obtain patch-level class scores for C classes. We add new classification head: Ym = gd (X where gd IR + RC with parameters
sh. For training, these scores compute an image-level prediction: ÿ = o o +===1 Ym). where an activation function. Crucially, we stop gradients from flowing back through X and am such that gd is trained as a separate read-out head jointly with the base model. We demonstrate this for megapixel MNIST (N 900), where the patch-level classifier learned from the multi-label task, using only image-level labels. In Fig. (right), visualize the inferred per-patch class label, argmax, Ym.. using colors and am using the alpha channel.
5 CONCLUSION
Reflecting on the experiments, one can observe that the IPS transformer performs well in various domains and training regimes and exhibits interesting capabilities: constant memory usage, learning with arbitrarily long sequences of varying lengths, multi-task learning, and modeling of inter-pal inter-patch relations. However, there are also limitations. For example, while the GPU memory can be kept con- stant, runtime cannot. In section 4.4 gave some intuition about what influence hyperparameters M and have in this regard. Furthermore, performance can degrade if the signal-to-noise ratio be- comes small, shown in MNIST scaling experiment-in practice, this might be mitigated by pre-training. Nonetheless, we believe that our main contribution, maintaining a low memory footprint for high-resolution image processing, can foster interesting applications that would other- wise be prohibitively expensive. Examples of such applications are learning with multiple medical scans for diagnosis/prognosis, the processing of satellite images, prediction tasks in videos, graphs, point clouds or whole-genon mhole-genome sequences. Technically, it would be interesting to extend IPS to dense prediction tasks, such as segmentation or depth estimation.

---

Published as a conference paper at ICLR 2023
ETHICS STATEMENT
Datasets CAMELYON16 a public benchmark dataset consisting of histopathology slides ob- tained from tissue of human subjects. The dataset has been approved by an ethics committee, see Litjens et al. (2018) for more information. For the Swedish traffic signs dataset, we searched the data for humans and excluded one image where face of a pedestrian was clearly visible. Megapixel MNIST builds on the widely used MNIST dataset, which contains digits that are ethically uncritical.
Applications In this paper, we exclusively conduct experiments and demonstrate potential appli- cations for human benefit, such as early traffic signs detection, which could be useful for advanced driver assistance systems, or abnormality detection in medical images, which could assist physi- cians. However, more care is needed before any of these systems can be deployed in worldworld applications including an assessment of biases and fairness and monitoring by medical profession- als where relevant. Furthermore, the attention scores currently not usable as detection method. Beyond the showcased applications, our method could be misused, just like any other visual object detection system, but now also on gigapixel images. We do not explore such applications in our paper and do not plan to do so in future work.
Computational efficiency Our method enables the training deep neural networks on arbitrary large inputs even on low-cost consumer hardware with limited VRAM, which makes Al devel- opment more accessible. All training and fine-tuning experiments on downstream tasks can be efficiently run a single A100 GPU, with the following total training runtimes: Traffic sign recog- nition: 45 min (M = 10, I = 32, 150 epochs), megapixel MNIST: 5h (M = 100, / = 100, N 3,481, 150 epochs), CAMELYONI6: 2h (M 5,000, I 5,000, N all, 50 epochs).
REPRODUCIBILITY STATEMENT
well the supplementary source code Github enable the reproduction of our method. Our metrics and baselines lean on existing work and described in Sect. 4 Necessary steps to reproduce our experiments are described in respective subsections of Sect. 4 We provide further details about all network architectures, hyperparameters for our method and the basclines, as well as pre-processing steps in Appendix D The tools used to measure the training runtime and peak GPU memory are listed in Appendix D.7
Sect. the pseudocode provided in Appendix A
ACKNOWLEDGMENTS
We would like to thank Joan Puigcerver for suggesting relevant related work, Romeo Sommerfeld for fruitful discussions on interpretability, Noel Danz for useful discussions about applications, as well as Jonathan Kreidler, Konrad Lux, Selena Braune and Florian Sold for providing and managing the compute infrastructure. This project has received funding from the German Federal Ministry for Economic Affairs and Climate Action (BMWK) in the project DAKI-FWS (reference number: 01MK21009E).
REFERENCES
Jaume Amores. Multiple instance classification: Review, taxonomy and comparative study. Artificial intelligence, 201:81-105, 2013,
Moab Arar, Ariel Shamir, and Amit Bermano. Learned queries for efficient local attention. In Pro- ceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 10841- 10852, 2022.
Babak Ehteshami Bejnordi, Mitko Veta, Paul Johannes Van Diest, Bram Van Ginneken, Nico Karssemeijer, Geert Litjens, Jeroen AWM Van Der Laak, Meyke Hermsen, Quirine F Manson, Maschenka Balkenhol, et al. Diagnostic assessment of deep learning algorithms for detection of lymph node metastases in women with breast cancer. Jama, 318(22):2199222 2210, 2017.
10

---

Published as a conference paper at ICLR 2023
Peter Burt. Attention mechanisms for vision in a dynamic world. In 9th international conference on pattem recognition, pp. 977-978. IEEE Computer Society, 1988.
Gabriele Campanella, Vitor Werneck Krauss Silva, and Thomas J Fuchs. Terabyte-scale deep multiple instance learning for classification and localization in pathology. arXiv preprint arXiv:1805.06663 2018.
Gabriele Campanella, Matthew G Hanna, Luke Geneslaw, Allen Miraforr Vitor Werneck Krauss Silva, Klaus Busam, Edi Brogi, Victor E Reuter, David S Klimstra, and Thomas Fuchs. Clinical-grade computational pathology using weakly supervised deep learning whole slide images. Nature medicine, 25(8):1301-139, 2019.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference computer vision, pp. 213-229. Springer, 2020.
Chun-Fu Richard Chen, Quanfu Fan, and Rameswar Panda. Crossvit: Cross-attention multi-scale vision transformer for image classification. Proceedings of the IEEE/CVF international con- ference on computer vision, pp. 357-366, 2021.
Richard Chen, Chengkuan Chen, Yicong Li, Tiffany Y Chen, Andrew D Trister, Rahul G Kr- ishnan, and Faisal Mahmood. Scaling vision transformers to gigapixel images via hicrarchical sulf-super learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 16144-16155 2022.
dean-Baptiss Cordonnier, Aravindh Mahendran, Alexey Dosovitskiy, Dirk Weissenborn, Jakob Uszkoreit, and Thomas Unterthiner. Differentiable patch selection for image recognition. Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2351- 2360, 2021.
Sean M Culhane and John Tsotsos. An attentional prototype for early vision. In European Conference Computer Vision, pp. 551-560. Springer, 1992.
Olivier Dehaene, Axel Camara, Olivier Moindrot, Axel de Lavergne, and Pierre Courtiol, Self- supervision closes the gap between weak and strong supervision in histology. arXiv preprint arXiv:2012.03583, 2020.
James Foulds and Eibe Frank. A review of multi-instance learning assumptions. The knowledge engineering review, 25(1):1-25, 2010.
Jean-Bastier Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach self-supervised learning. Advances in neural information processing systems, 33:21271-21284. 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Maximilian Ilse, Jakub Tomezak, and Max Welling. Attention-based deep multiple instance learn- ing. In International conference on machine learning, pp. 2127-2136. PMLR, 2018
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79 87, 1991.
Andrew Jacgle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint 1x11:22007.14955 2021.
Angelos Katharopoulos and François Fleuret. Processing megapixel images with deep attention- sampling models. In International Conference Machine Learning, pp. 3282-3291 PMLR, 2019.
11

---

Published as a conference paper at ICLR 2023
Christof Koch and Shimon Ullman. Shifts in selective visual attention: towards the underlying neural circuitry. Matters of intelligence, pp. 115-141. Springer, 1987.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Marvin Lerousseau, Maria Vakalopoulou, Marion Classe, Julien Adam, Enzo Battistella, Alexandre Carré, Théo Estienne, Théophraste Henry, Eric Deutsch, and Nikos Paragios. Weakly supervised multiple instance learning histopathological tumor segmentation. International Conference on Medical Image Computing and Computer-Assister Intervention, pp. 470-479. Springer, 2020.
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, pp.
2021.
Xiangru Lian and Ji Liu. Revisit batch normalization: New understanding and refinement via com- position optimization. In The 22nd International Conference on Artificial Intelligence and Statis- tics, pp. 3254-32555 PMLR, 2019.
Geert Litjens, Peter Bandi, Babak Ehteshami Bejnordi, Osear Geessink, Maschenka Balkenhol, Peter Bult, Altuna Halilovic. Mcyke Hermsen, Rob van de Loo, Rob Vogels, et al. 1399 h&e- stained sentinel lymph node sections of breast cancer patients: the camelyon dataset. GigaScience, 7(6):giy065, 2018.
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Ming Y Lu, Drew FK 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-550, 2021.
Oded Maron and Tomás Lozano- Pérez. A framework for multiple multiple-instance instance learning. Advances in neural information processing systems, 10, 1997.
Nobuyuki Otsu. A threshold selection method from gray-level histograms. IEEE transactions on systems, man, and cybernetics, 9(1):22-66, 1979.
Hans Pinckaers, Bram Van Ginneken, and Geert Litjens. Streaming convolutional neural networks for end-to-end learning with multi-megapixe images. IEEE transactions on pattern analysis and machine intelligence, 2020.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34.8583-8555 2021.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJC),, 115(3):211225 2015. doi: 10. 1007/1 0077/11263-011 1263-015-0816.y.
Michael Ryoo, AJ Piergiovanni, Anurag Arnab, Mostafa Dehghani, and Anelia Angelova. Token- learner: Adaptive space-time tokenization for videos. Advances in Neural Information Processing Systems,
2021.
Carl F Sabottke and Bradley M Spieler.
The effect of image resolution on deep learning in radiog- raphy. Radiology. Artificial intelligence, 2(1), 2020.
Furkan Sahin. Long-range, nigh-resolution camera optical design for assisted and autonomous driving. In Photonics, volume 6, pp. 73. MDPI, 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gad -ated cixture-of-expertt layer. arXiv preprint 170.0007000005535 2017.
12

---

Published as a conference paper at ICLR 2023
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. arXiv preprint
2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Dayong Wang, Aditya Khosla, Rishab Gargeya, Humayun Irshad, and Andrew H Beck. Deep learning for identifying metastatic breast cancer. arXiv preprint arXiv: 05711, 2016.
Xiyue Wang, Sen Yang, Jun Zhang, Minghui Wang, Jing Zhang, Wei Yang, Junzhou Huang, and Xiao Han. Transformer-basee unsupervised contrastive learning for histopathological image clas- sification. Medical Image Analysis, 81:10255y 2022.
Lei Zhou, Yu Zhao, Jie Yang. Qi Yu, and Xun Xu. Deep multiple instance learning for automatic detection of diabetic retinopathy in retinal images, IET Image Processing, 12(4):563-571, 2018.
13

---

Published as a conference paper at ICLR 2023
APPENDIX
A PSEUDOCODE
Algorithm 1: Pseudocode for IPS and Patch Aggregation
for step = 1, 2...... do
Load image patches Activate no-gradient and evaluation mode
PM Embed first patches
for = to T do
P = kmbed next patches
a'+l Compute their scores
pt+1 select top M embeddings from Pyr P end
Fetch patches according to Ph
Activate gradient and training mode
Embed and aggregate selected patches
Classify, compute loss, optimize
end
B MULTI-HEAD
AND TRANSFORMER
In practice we use multi-head cross-attention (MCA), which computes multiple intermediate repre- sentations zA, h 1...., corresponding H heads and then aggregates them:
z concat (z) ..... =*). z W
(4)
(5)
with 2 E RD,, x" E () E RD and W E H)), Note that each representation zh computed with head-specific attention scores a² For patch selection, these scores are averaged over heads to obtain a scalar score per patch. The transformer module is composed of an MCA layer, multi-layer perceptron (MLP) and Layemnrm (LN) and is formulated as:
MCA(X* + X
(6)
z' LN(2 qO)
(7)
z" -L(MMP((2)) z'),
(8)
where Pos IR MDD is an optional 1/ sinusoiiaa position encoding and the MLP consists of two layers with ReLU non-linearity.
C EXTENDED RELATED WORK
Cross-attention is an essential component of our method because it links the patch selection and patch aggregation modules. Therefore, we attention a concept in which the queries stem from another input than the one utilized by keys and values. It has previously been employed various transformer-bast based architectures and applications. In the original transformer, it is used in an autoregressive decoder for machine translation /Vaswani al. 2017). where keys and values are computed from the encoder output, whereas queries are computed from the already translated sequence. Cross applies cross-attention to efficiently fuse patch encodings corresponding to varying patch sizes (Chen et al. 2021). Cross-attention can also be applied to object detection. For example, DETR defines learnable object tokens that cross-attend to patch embeddings, where the number of tokens reflects the maximum number of objects that can be detected (Carion al. 2020). More related which involves learnable queries as a substitute for local convolutions and the more computationally expensive self-attention mechanism. Another interesting architecture is Perceiver IO, which applies
shed more light its use in the literature. Cross-
work is Query and-Attend (Arar al. 2022
14

---

Published as a conference paper at ICLR 2023
cross-attention to the encoder for dimensionality reduction and the decoder for learning task- specific outputs (Jegle et al.) 2021). TokenLearner (Ryoo et al. 2021) also reduces the number of tokens in the encoder for image and video recognition tasks but without using cross-attention module. In our work, cross-attention serves two purposes. First, we treat it as a MII. pooling function that is used in the rear part of the network combine information from multiple patches. Second, we take advantage of the attention scores to select fixed number of most salient patches, which makes memory consumption independent of input size and thus allows us to process arbitrarily large images.
Checkpointing Another memory-efficien approach the StreamingCNN (Pinckaers et al. 2020), which employs gradient checkpointing to process large images sequentially through CNN, which demonstrated on images up to 8k 8k pixels. Due to the use of checkpoints, multiple forward and backward passes required. Furthermore, backprop is performed for all tiles, which further slows down the training runtime. In IPS, all patches only run once through the network in fast no- ---ardddent mode. Only selected patches then require second forward and backward pass.
D EXPERIMENT DETAILS
D. TRAINING
All models are trained for 150 epochs (megapixel MNIST, traffic signs) or 50 epochs (CAME- LYONI6) on the respective training sets, and results are reported after the last epoch on the test set without early stopping. The batch size is 16, and AdamW with weight decay of 0.1 is used as op- timizer (Loshchilov & Hutter 2017). After a linear warm-up period 10 epochs, the learning rate is set to 0.0003 when finetuning pre-trained networks and 0.001 when training from scratch. The learning rate is then decayed by a factor of 1,000 over the course of training using a cosine schedule.
D.2 ARCHITECTURE COMPONENTS
Table provides high-level overview of the various components used by the baselines and method. Note that all methods use the same encoder for a fair comparison.
Table 5: Module pipeline of baselines and our method
#
CNN DPS Transformer
TopMIL
DeepMIL
IPS Transformer (ours)
Encoder
Scorer + Difff Top-M
Patch selection Encoder
Encoder
IPS
2
Class. head
Encoder
Gated attention
Encoder
3
Activation
Cross-attention transf.
Class. head
Class. hcad
transf.
4
Class. head
Avg. pool
Activation
Class, head
5
Activation
Activation
Activation
D.3
ENCODERS
The encoders differ depending on the dataset and are listed in Table After applying each encoder, global average pooling used to obtain a feature vector per patch. The CNN basclines also use these encoders, but the encoders are applied to the image instead of the patches.
D.4 TRANSFORMER HYPERPARAMETERS
The hyperparameters of the cross-attention transformer follow default values (Vaswani et al. 2017)
and are listed in Table 7 We slightly deviate from these settings for very small values of M. In particular, for M = 2 in traffic signs, we set Dinner = D. For both traffic signs with M = 2 and megapixel MNIST with M 5, we attention dropout to O.
15

---

Published as a conference paper at ICLR 2023
Table 6: Encoder architectures and pre-training for each dataset.
Dataset
Encoder
Pre-training
Features
Traffic signs CAMELYON16
ResNet-18
ImageNet-lk
512
ResNet-50 projector
BYOL
512
Megapixel MNIST
ResNet-18 (2 blocks)
None
128
Table Transformer hyperparameters. Dinner is the hidden layer size the feed-forward MLP of the transformer block.
Setting
Value
Heads
D
Features
D&
D/8
Do
D/8
Dinner
4D
dropout
0.1
attention dropout
0.1
D.5 BASELINE HYPERPARAMETERS
In DPS, we mostly follow the setup of Cordonnier et al. (202)). We downsample the original image by 3 to obtain low-resolution image (MNIST: 500 500. traffic signs: 400x533). As scorer, we use the first two blocks of ResNet-18 without pre-training for megapixel MNIST and with ImageNet-I 1k pre-training for traffic signs, followed by a Conv layer (kernel: 3, padding: 0) with single output channel and max pooling (kernel: 2, stride: 2). The resulting score maps have shape 30x30 for megapixel MNIST and 24 32 for traffic signs. Thus, we consider N 900 patches and N 768 patches for megapixel MNIST and traffic signs, respectively. We use 500 noise samples for perturbed top-K, employ gradient 2 norm clipping with maximum cutoff value of 0.1 and set perturbation noise MNIST and M = 5, o-decay was deactivated as it resulted in degraded performance towards the end of training.
0.05, which linearly decays to 0 over the course of training. For megapixel
Table 8: Classification heads for each dataset for DeepMIL. and TopMIL.
#
Traffic signs
Megapixel MNIST
CAMELYON16
fc (# classes)
fc (128) ReLU fc (# classes)
fo (128) + ReLU fo (64) ReLU fo (# classes)
2
3
For DeepMIL., D₁ 128 following llse et al. (2018), whereas D, 1024 for DeepMIL++ The classification heads of the MIL. baselines differ slightly between datasets and are listed Table In traffic signs recognition, single output layer is utilized following Îlse et al. (2018). More capacity added for DeepMIL+: fo (2,048) ReLU, fe (512) + ReLU, fc (# classes). For megapixel MNIST, we obtained better results by adding an additional layer, which could be due to the multi-task setup. For CAMELYONI6, 16, we followed the settings of Dehaene et al. (2020). In DPS and our method, we use the same setup for all datasets: fo (# classes). The leteh-level classification head introduced in uses the same structure as the MIL methods megapixel MNIST (Table middle column).
Sect. 4.5
D.6 PRE-PROCESSING OF CAMELYON16
Since large fraction of WSI may not contain tissue, we first read each image chunks, create histogram of pixel intensities, and apply Otsu's method to the histogram to obtain threshold
16

---

Published as a conference paper at ICLR 2023
per image. We then extract non-overlapping patches of size 256x256, and retain patches with least 1% the area being foreground. BYOL is then trained with a ResNet-50 for 500 epochs using the training settings of Grill et al. (2020) and the augmentations listed in Table 9 which mostly follow those used in Dehaene et al. (2020). each pre-training epoch, 270,000 patches approximately 1,000 patches per training slide. We train
sampled at random, corresponding on A100 GPUs with a batch size of 256 using mixed precision. After pre-training, each patch center cropped with 224 px per side and 2.048-dimensional features resulting from the ResNet- 50 encoder are extracted. The features are then normalized and stored in an HDF5 file. For the downstream task, we train a projector (fe (512), BatchNorm, ReLU), and its output is passed on to the aggregation module. For the pre-processing of the megapixel MNIST and traffic sign datasets, we adapt the code of Katharopoulos & Fleuret (201).
Table BYOL augmentations. Square brackets indicate varying probabilities each of the views.
Type
Parameters
Probability
Rotation
{0', (09900,
90" 180", 270"}
1.0
VerticalFlip
0.5
HorizontalFlip
0.5
ResizedCrop
scale (0.2, 1.0)
1.0
ColorJitter
brightness = 0.8,
0.8
contrast = 0.8,
saturation 0.8,
hue = 0.2
Grayscale GaussianBlur Solarize
0.2
kernel size = 23
[1.0, 0.1
threshold 150
(0.0, 0.2
D.7 COMPUTATIONAL EFFICIENCY
We measure the training runtime for a full forward and backward pass without data loading for batch size of 16, and report the average of all but the first and last iterations of epoch. For time recording, we utilize torch. cuda. Event, which can be used as follows to measure the elapsed time for single iteration:
import torch
start_event = torch. E
end event = torch. cuda. Event (enable timing=True)
start_event.recorz
#forward and backward pass
end_ovent.rccord
corch.cuda www...cw.......c.ii.i.....m ()
time start_cvent. time (end_event)
After running an experiment for one epoch, the following code can be used to measure the peak GPU memory usage:
stats torch. cuda memory sttttorysstttt)))))
peak") peak go mem stats "allocated bytes. pe.aaa) 1024
E IPS WITH DEEPMIL
We want explore how IPS performs with different aggregation module than ours. Like
attention module, the gated attention mechanism used in DeepMIL (llse et al. 2018) also assigns attention scores each patch. We thus train IPS jointly with DeepMIL megapixel MNIST (multi- task) and report results for each of 5 runs in Table 10 for the majority task with N = 900, = 100 and different values of M.
17

---

Published as a conference paper at ICLR 2023
Table 10: Performance of DeepMIlL when combined with IPS.
M
Accuracy per run (sorted)
5
9.1
9.6
10.0
10.0
94.3
100
10.0
98.6
98.7
98.7
99.2
200
98.5
98.7
98.9
99.0
99.1
For M = 5, one can observe that the model was able to localize the digits in only one out of five runs. However, when increasing the memory buffer to 100, only one run failed, and for M 200 all runs were successful. For IPS + transformer, all runs achieved over 90% accuracy, even for M = 5. The average accuracy of IPS DeepMIL for MM 200 is 98.8%, which on par with IPS+transformer for .M 100 (98.9%). This shows that IPS can be used with different attention-based aggregation modules, however IPS + transformer more robust for small values of M.
F ADDITIONAL VISUALIZATIONS
Epoch O
Fpoch
Fpoch
Fprch
Fpech 20
25
Epoch
Epoch
Epoch
Figure 6: Evolution of attention scores in IPS for M 10. The traffic sign (80 km/h) is localized starting from epoch 20, The entropy of the attention distribution appears to decrease during training.
80
70
80
80
50
80
8
50
80
70
Figure 7: Patches with the highest attention scores (decreasing from left-to-right then top-to-bottom) out of all images from the traffic sign recognition test set.
18

---

Published as a conference paper at ICLR 2023
Figure 8: Attention scores and crops of the most salient patches for example test images in the traffic sign recognition dataset.
8
8
6
(
Z
9
9
9
2
2
3
3
0
2
3
2
9
Figure 9: Patch-level classification scores and crops of patches that obtained highest attention scores (decreasing from left-to-right then top-to- to-ttbbbotto per image) in megapixel MNIST (N 3,481).
19

---

Published as a conference paper at ICLR 2023
Figure 10: Attention maps overlaid on the input image and crops of the most salient patches for two examples of the CAMELYON 16 dataset.
20


Reviews:
[{"review_id": "k2wK1JVtA0t", "reviewer": "Reviewer_jHHu", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors deal with the problem of training neural networks when the input images are very large and thus cannot fit entirely into GPU memory (ie, cases where even a batch size of 1 is too large). They propose an Iterative Patch Selection algorithm which selects only the most salient patches which are then aggregated into a global representation for image recognition. A transformer is used for both patch selection and aggregation. Their experiments are able to process >16 gigapixel images on a 5GB GPU.", "summary_of_the_review": "Overall I think this paper should be of interest to the ICLR community as it helps enable big scale deep learning under limited resources. The main weakness I see is that the authors need to do a better job positioning the paper with respect to early work on attention.", "strength_and_weaknesses": "Strengths:\n\nS1: The paper addresses an important problem that plagues graduate students and researchers who may not have access to the latest and most expensive clusters of GPUs, namely that of being able to do vision under limited resources. It should therefore be of interest to the ICLR community.\n\nS2: the paper is well written and is up to ICLR standards. However as I indicate below, in certain places the authors should move some of the details from the appendix into the main text.\n\n\nWeaknesses:\n\nW1: Fig 1: it is unclear how the patch scores are shown in this figure\n\nW2: under Eq.1, X_1,...,X_M is not well defined.\n\nW3: page 2, last paragraph: the relevant information is really in the appendix. I suggest the authors add more details here. This is currently too high level.\n\nW4: under Eq 3, the variables just under Eq 3 need to be discussed in more detail\n\nW5: in the paper's prior work, the authors should make a better connection to early approaches for simulating saccades in vision and attention beams. Why do humans have attention? \n\nW6: Appendix A pseudocode: this is not pseudocode, it is actual code. I suggest edits to make it more readable.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "imbWbyaso-E", "reviewer": "Reviewer_jpUm", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new method of multiple-instance learning from high resolution images, which uses a patch selection scheme that processes available patches in mini-batches, scores them, then aggregates the top-M highest-scoring patches’ embeddings to reach a bag/image-level prediction. The proposed IPS Transformer out-performs state-of-the-art methods on 3 different high resolution image benchmarks. It is shown that IPS Transformer can be less sensitive to the choice of pool size (M), use less memory, while performing similarly to existing methods.", "summary_of_the_review": "The paper presents a practical yet simple solution to MIL on high-resolution images and experimentally validates its performance on 3 benchmarks. The presentation and analysis are of high quality, warranting the paper’s acceptance.\n\nEdit after rebuttal: I did not previously have major concerns regarding this submission, and after carefully reading the other reviews as well as the authors' diligent response and action, I maintain my rating of 8 as no further concerns exist from my side.", "strength_and_weaknesses": "The manuscript is written quite well and the related work is handled particularly well. The experimental settings are clear and seem to be fair, with many related works being reproduced by the authors themselves. The IPS Transformer clearly performs well while being less sensitive to the choice of M, uses lower GPU RAM, and has competitive inference times.\n\nThough the proposed algorithm is not excessively complex, the concrete steps of how to train IPS Transformer is scattered in the text of Sec. 2. It is therefore difficult to clearly understand the steps involved and which modules are trained when. The paper may benefit from either an algorithm figure in the main paper or a pipeline/workflow diagram figure.\n\nThe experimental results and ablation studies look complete.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PDb4-GPwyv", "reviewer": "Reviewer_AtFj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The problem of solving computer vision tasks with neural networks on large images is considered. Authors propose a method that splits an image into patches and then applying greedy iterative approach (IPS iterative patch selection) to select the representative ones. Selection is performed by creating a buffer and iteratively updating it by scoring them. Final buffer is then transformed via a convex combination of embeddings. The approach is more heursitic than theoretically justified. Results are demonstrated on multiple dataset where high resolution matters, the method is compared with multiple others. ", "summary_of_the_review": "Paper tackles interesting problem of working with high resolution images. The solution is a system that selects informative patches and processes them with transformer. The novelty of the method is on the weak side and some components that can be considered interesting to the community are not well explained. Given little novelty and no significant theoretical justification of the method my recommendation is top reject the paper. \n\n\n------Post rebuttal------\nThank authors for addressing raised points. The way attention module is trained without gradients is more clear now and updates to the figure are great. Reading other reviews it seems they are more positive. On my end, paper has only empirical contribution is it is not groundbreaking, however, given other reviews I will increase my score to 6 and will not object if the paper is accepted. ", "strength_and_weaknesses": "Strength:\n\n- The problem is very relevant and considers important problem outside of the main focus. Solving CV tasks at high resolution is challenging and authors provide a working (according to results) solution. \n- Overall novelty of the paper is in 2 components: (i) iterative patch selection and (ii) patch aggregation.\n- The idea of using transformer attention for scoring that is learned without gradient is novel and seem to work fine.\n- The related work section is quite broad and gives a nice overview for the person not in the field.\n\nWeaknesses/questions:\n\n1) Significance of the method and novelty. The approach seem to be hand designed and heuristic. We don't see much of the theory to back it up. The contribution of (i) iterative patch selection can be considered as a straightforward and not novel. Patch aggregation with the idea of scoring patches is also not groundbreaking. The way patch scores are training is an in interesting point.\n2) The explanation of how attention scorer is trained is a bit vague. Paper states multiple times that there is no gradient propagation (no-gradient mode), but then it is not clear how to train the scorer. There is a second network that aggregates patches, but it seem to take patch + positional encoding meaning previous scores are not needed. If both blocks share the same attention then we can train only the second one, and use attention scores for the selector. A more detailed explanation of the novel part is needed. \n3) The implementation uses a relatively heavy network resnet50 to embed patches. This network by itself can potentially solve the problem if information from all patches is aggregated by some way of pooling (max for example). Probably learning an embedder will be useful for the work. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HMHZEWPHayg", "reviewer": "Reviewer_Ffm2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes  a simple method, Iterative Patch Selection (IPS), which can select the image patches from a high-resolution image and meet the memory constraint. Two-stage DNN is used to produce the final result, the first part of the DNN is used to select the input patches in autoregressive fashion. The selected patches are then processed by the second part of the DNN to produce the final results. The results indicate that ISP can lead to superior accuracy with minimum level of memory footprint. ", "summary_of_the_review": "Image patch selection for large-scale image is an interesting and practical problem. The paper adopts the transformer based DNN to perform both image patch selection and classification. However, I am a bit concerned about the novelty of the work and the reproducibility of the work, given that no code has been submitted.", "strength_and_weaknesses": "Strong point:\n+ Image patch searching for high-resolution image is an interesting and practical problem.\n+ The proposed solution is technically sound.\n+ The evaluation results are comprehensive.\n\nWeak points:\n- Compared with \"Differentiable Patch Selection (DPS)\" which also use transformer based image patch selection, the contribution is not that significant.\n- The no-gradient mode is confusing to me. Why can't you use the encoder transformer to compute the scores for all the patches, then use some approximation function (e.g., softmax) to estimate the top-k function. This will make the whole DNN differentiable.\n- Not sure whether the proposed architecture will bring additional hardware constraint due to the storage of the intermediate results within the DNN. The author should justify that.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster