[論文レビュー] Iterative Patch Selection for High-Resolution Image Recognition

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

Advanced Neural Network Applications被引用数 4

ひとこと要約

本稿では、クロスアテンショントランスフォーマーを用いて顕著なパッチを段階的に選択することで、入力サイズに依存しないGPUメモリ使用量を実現する、メモリ効率の良い高解像度画像認識手法であるイテレーティブパッチ選択（IPS）を提案する。IPSにより、16GB以上の画像（例：25万パッチ）をバッチサイズ16で学習可能であり、5GBのVRAMで運用可能であり、ベースラインと比較して顕著にメモリ消費量を削減しながら高い精度を維持する。

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.

研究の動機と目的

• 標準GPU上で高解像度画像の深層ニューラルネットワーク学習を実行する際に生じる過大なメモリ消費を克服すること。
• メモリ使用量を入力サイズから分離し、厳しいハードウェア制約下でも任意に大きな画像を処理可能にする。
• 多様なドメインや画像サイズにおいて、最小限のアクセラレータメモリを用いても高い精度を維持する手法を開発すること。
• 大規模な医療および実世界の画像データセットにおいて、効率的なファインチューニングおよび自己教師付き事前学習を可能にすること。

提案手法

• IPSは、高解像度画像をパッチに分割し、共有エンコーダーを用いて各パッチを独立して埋め込み処理する。
• 勾配なしモードでクロスアテンショントランスフォーマーを用いてパッチをスコア付けし、アテンションスコアを基準として顕著なパッチを特定する。
• イテレーティブで自己逐次的なバッファ更新を採用し、上位M個の顕著なパッチを維持する。メモリ管理のため、パッチをI個のチャンクに分割して処理する。
• 選択されたパッチは、2番目のクロスアテンショントランスフォーマーを用いて集約され、分類用のグローバルな画像表現が形成される。
• 分類ヘッドは、選択されたパッチおよびアテンション重みを介して勾配が戻らないように、パッチ選択および集約モジュールと同時に学習される。
• 本手法は複数インスタンス学習に根ざしており、パッチレベルでの確率推定が可能な弱教師付きマルチラベル分類を可能にする。

実験結果

リサーチクエスチョン

• RQ1 パッチベースの手法が、入力サイズにかかわらず一定のGPUメモリ使用量を維持しながら、高解像度画像で高い精度を達成できるか？
• RQ2 クロスアテンションを用いたイテレーティブパッチ選択は、標準的なCNNおよびトランスフォーマーと比較して、大規模な入力におけるメモリ効率および性能面で優れているか？
• RQ3 IPSが、5GB未満のVRAMで、16ギガピクセル以上のウェブスライド画像のファインチューニングをどの程度可能にするか？
• RQ4 IPSは、CAMELYON16のような大規模データセットにおいて、マルチタスクおよび自己教師付き学習を低メモリオーバーヘッドでサポートできるか？
• RQ5 ハイパーパrameterであるM（選択パッチ数）およびI（1イテレーションあたりのバッチサイズ）が、計算効率およびモデル性能に与える影響はいかほどか？

主な発見

• IPSにより、25万パッチ（16ギガピクセル以上）のウェブスライド画像のファインチューニングが、バッチサイズ16で5GBのGPU VRAMのみで実現可能である。
• 1k×1kから10k×10kにスケーリングされたメガピクセルMNIST画像において、ピークメモリ使用量は1.7GBで一定を保ち、同等のCNNは2k×2k解像度で24.6GBを消費した。
• 交通標識認識、合成メガピクセルMNISTにおけるマルチタスク分類、CAMELYON16ベンチマークにおける自己教師付き事前学習の両方で高い精度を達成した。
• 3つのベンチマークすべてにおいて、1台のA100 GPUで効率的に学習およびファインチューニングが完了し、合計実行時間は、交通標識で45分、MNISTで5時間、CAMELYON16で2時間であった。
• クロスアテンション機構のアテンションスコアにより、弱教師付きのパッチレベル分類確率推定が可能となり、色分けされたアテンションマップとして可視化された。
• アブレーションスタディにより、イテレーティブパッチ選択とクロスアテンション集約の組み合わせが、低メモリ使用量と高い性能を維持する上で不可欠であることが確認された。