[论文解读] Processing Megapixel Images with Deep Attention-Sampling Models

Angelos Katharopoulos, François Fleuret|arXiv (Cornell University)|May 3, 2019

Medical Imaging Techniques and Applications被引用 22

一句话总结

本文提出一种端到端可微的注意力采样模型，通过在低分辨率图像视图上计算的注意力分布对图像的少量像素区域进行采样，仅处理百万像素图像的一小部分。该方法在计算量和显存使用减少10倍的同时，实现了与全图注意力模型相当的精度，且通过标准SGD进行训练，梯度估计方差最小化。

ABSTRACT

Existing deep architectures cannot operate on very large signals such as megapixel images due to computational and memory constraints. To tackle this limitation, we propose a fully differentiable end-to-end trainable model that samples and processes only a fraction of the full resolution input image. The locations to process are sampled from an attention distribution computed from a low resolution view of the input. We refer to our method as attention sampling and it can process images of several megapixels with a standard single GPU setup. We show that sampling from the attention distribution results in an unbiased estimator of the full model with minimal variance, and we derive an unbiased estimator of the gradient that we use to train our model end-to-end with a normal SGD procedure. This new method is evaluated on three classification tasks, where we show that it allows to reduce computation and memory footprint by an order of magnitude for the same accuracy as classical architectures. We also show the consistency of the sampling that indeed focuses on informative parts of the input images.

研究动机与目标

• 解决现有CNN在处理百万像素图像时面临的计算与显存限制。
• 实现在高分辨率图像上端到端训练模型，而无需每patch的标注。
• 在保持与全图处理相当的分类精度的同时，降低计算成本与显存使用。
• 开发一种可微采样机制，使梯度能够通过采样过程反向传播。

提出的方法

• 模型从输入图像的低分辨率版本中计算注意力分布，以识别信息丰富的区域。
• 基于该注意力分布，从完整图像中采样少量patch（例如5个）。
• 通过注意力权重对采样patch的特征进行加权平均聚合。
• 推导出无偏梯度估计器，实现使用标准SGD的端到端训练，避免使用强化学习或变分推断。
• 证明基于注意力的采样可产生全模型输出的最小方差估计器。
• 实现中采用ResNet-based特征提取器和四层卷积网络作为注意力头。

实验结果

研究问题

• RQ1 能否设计一种可微采样机制，在无需每patch标注的情况下高效处理百万像素图像？
• RQ2 基于注意力分布的采样是否相比均匀采样或启发式采样具有更低方差的估计器？
• RQ3 该模型能否在显存和计算量降低一个数量级的同时，实现与全注意力模型相当的精度？
• RQ4 在采样过程中，能否在不依赖强化学习的前提下，使用标准SGD实现端到端训练？

主要发现

• 所提出的注意力采样模型在仅使用5个patch的情况下，测试误差与处理全部192个patch的Deep MIL模型相当。
• 与Deep MIL相比，该模型将显存使用减少最多30倍，推理时间减少最多25倍。
• 基于注意力分布的采样可产生全模型输出的最小方差估计器，具有统计最优性。
• 即使没有每patch标签，注意力分布也能持续聚焦于信息丰富的区域，如交通标志。
• 在下采样图像上，该模型的泛化能力优于标准CNN，后者会因丢失关键细节（如限速数字）而表现下降。
• 该方法使单张GPU即可实现百万像素图像的训练与推理，克服了以往高分辨率处理的限制。