[Paper Review] Iterative Patch Selection for High-Resolution Image Recognition

Benjamin Bergner, Christoph Lippert+1 more|arXiv (Cornell University)|Oct 24, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Iterative Patch Selection (IPS), a memory-efficient method for high-resolution image recognition that decouples GPU memory usage from input size by iteratively selecting salient patches via a cross-attention transformer. IPS enables training on gigapixel images (e.g., 250k patches) with only 5 GB VRAM at batch size 16, maintaining high accuracy while drastically reducing memory consumption compared to baselines.

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.

Motivation & Objective

• To address the challenge of training deep neural networks on high-resolution images due to prohibitive memory consumption on standard GPUs.
• To decouple memory usage from input size, enabling processing of arbitrarily large images under tight hardware constraints.
• To develop a method that maintains high accuracy while using minimal accelerator memory across diverse domains and image sizes.
• To enable efficient fine-tuning and self-supervised pre-training on large-scale medical and real-world image datasets.

Proposed method

• IPS processes high-resolution images by dividing them into patches and embedding each patch independently using a shared encoder.
• A cross-attention transformer scores patches in no-gradient mode to identify the most salient ones, using attention scores as selection criteria.
• The method employs an iterative, autoregressive buffer update to maintain the top-M salient patches, processing patches in chunks of size I to manage memory.
• The selected patches are aggregated via a second cross-attention transformer to form a global image representation for classification.
• The classification head is trained jointly with the patch selection and aggregation modules, with gradients stopped from flowing back through the selected patches and attention weights.
• The approach is grounded in Multiple Instance Learning and allows for weakly-supervised multi-label classification with patch-level probability estimation.

Experimental results

Research questions

• RQ1 Can a patch-based method achieve high accuracy on high-resolution images while maintaining constant GPU memory usage regardless of input size?
• RQ2 How does iterative patch selection via cross-attention compare to standard CNNs and transformers in terms of memory efficiency and performance on large inputs?
• RQ3 To what extent can IPS enable fine-tuning on whole-slide images (e.g., 16+ gigapixels) with minimal VRAM?
• RQ4 Can IPS support multi-task and self-supervised learning on large-scale datasets like CAMELYON16 with low memory overhead?
• RQ5 How do hyperparameters M (number of selected patches) and I (batch size per iteration) affect computational efficiency and model performance?

Key findings

• IPS enables fine-tuning on whole-slide images with up to 250,000 patches (>16 gigapixels) using only 5 GB of GPU VRAM at a batch size of 16.
• For megapixel MNIST images scaled from 1k×1k to 10k×10k, peak memory usage remained constant at 1.7 GB, while a comparable CNN consumed 24.6 GB at 2k×2k resolution.
• The method achieved high accuracy on traffic sign recognition, multi-task classification on synthetic megapixel MNIST, and self-supervised pre-training on the CAMELYON16 benchmark.
• Training and fine-tuning on all three benchmarks completed efficiently on a single A100 GPU, with total runtimes of 45 minutes (traffic signs), 5 hours (MNIST), and 2 hours (CAMELYON16).
• The attention scores from the cross-attention mechanism enabled weakly-supervised patch-level class probability estimation, visualized as color-coded attention maps.
• Ablation studies confirmed that the combination of iterative patch selection and cross-attention aggregation is key to maintaining low memory usage and high performance.