Hard-Attention for Scalable Image Classification

1 Introduction . In image classification , deep neural networks ( DNNs ) are typically designed and optimized for a specific input resolution , e.g . 224× 224 px . Using modern DNNs on images of higher resolution ( as happens e.g. , in satellite or medical imaging ) is a non-trivial problem due to the subtlety of scaling model architectures [ 57 ] , and rapid increase in computational and memory requirements . A linear increase in the spatial dimensions of the input , results in a quadratic increase in computational complexity and memory , and can easily lead to resource bottlenecks . This can be mitigated with careful engineering , e.g. , streaming [ 46 ] or gradient checkpointing [ 42 ] . However , such solutions are content-agnostic , and don ’ t take advantage of the fact that discriminative information may be sparse and distributed across various image scales , deeming processing of the whole input unnecessary . Our goal is to leverage high-resolution information , while dispensing with the unsustainable quadratic complexity to input scale . To this end , we propose Traversal Network ( TNet ) , a multi-scale hardattention architecture , which traverses image scale-space in a top-down fashion , visiting only the most informative image regions along the way . TNet is recursive , and can be applied to inputs of virtually any resolution ; an outline of its processing flow is presented in Fig . 1 ( a ) . Our method draws its intuition from the way humans use saccades to explore the visual world . TNet offers an adjustable trade-off between accuracy and complexity , by changing the number of attended image regions . This way , complexity increases linearly with the number of attended locations , irrespective of the input resolution . Also , hard-attention explicitly reveals the image regions that our model values the most , providing a certain degree of interpretability ( Fig . 1 ( c ) ) . Importantly , interpretability comes without any extra cost beyond inference , in contrast to popular attribution methods , which require at least an additional backward pass [ 53 ] , or numerous forward passes [ 1 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Attention may also reduce data acquisition cost [ 60 ] , by allowing only a fraction of the high-resolution content to be acquired . TNet is trained end-to-end by employing a modified version of REINFORCE rule [ 68 ] , while using only classification labels . Our architecture is modular , and most classification models could be adopted as its backbone for feature extraction . This way , we can directly take advantage of various performance benefits offered by existing optimized deep models . Hard-attention is the mechanism that allows TNet to dispense with quadratic complexity to input scale , and as a result , we evaluate our model against strong hard-attention baselines on ImageNet [ 18 ] . A summary of our results is depicted in Fig . 1 ( b ) , where we see that TNet offers a better trade-off between accuracy and complexity measured in FLOPs ( similar behavior is observed with actual timings and memory ) . We extend our experiments to fMoW dataset , which consists of high-resolution satellite images [ 12 ] . We process images up to 896 × 896 px , getting up to 2.5x faster processing compared to baselines operating on the same resolution , while achieving higher accuracy as well . We find improvements in accuracy surprising , because TNet is processing only part of the input , in contrast to fully convolutional baselines . We primarily attribute this behavior to a novel regularization method , which encourages classification based on individual attended locations . We verify its efficacy through an ablation study . 2 Related work . Attention . Attention has a long history in the artificial neural networks literature [ 31 ] , and in the modern era of deep learning it has been used very successfully in various problems [ 36 , 14 , 4 , 70 , 22 , 67 , 55 ] . Two main forms of attention are : soft attention which processes everything but weights various regions differently ; and hard attention which selects only a fraction of the data for processing . Hard-attention models address various use-cases , and can be motivated by interpretability [ 18 ] , reduction of high-resolution data acquisition cost [ 60 ] , or computational efficiency [ 33 ] . Our goal is to offer a single model that takes advantage of all these benefits . Our model is conceptually similar to glimpse-based models [ 2 , 43 , 50 , 54 , 19 , 3 ] . An important difference is that we don ’ t restrict our attention policy to 2D space , but we consider scale dimension as well . Also , our model parallelizes feature extraction at each processing level , instead of being fully sequential . Furthermore , we don ’ t use recurrent neural networks ( RNNs ) to combine features from different locations , but instead , we simply average them ( see Section 3 ) . This way , gradients flow directly to extracted features , without the need to backpropagate through RNN steps . This is a simpler strategy compared to LSTM in order to avoid vanishing gradients . Recent work explores soft attention mechanisms based on transformers , which originate from the natural language processing community [ 63 ] . Transformers have already been used extensively in machine vision [ 45 , 74 , 16 , 8 , 48 ] , and research interest in related directions has increased [ 6 , 11 , 5 ] . Multi-scale representations . We identify four broad categories of multi-scale processing methods . ( 1 ) Image pyramid methods extract multi-scale features by processing multi-scale inputs [ 17 , 47 , 34 , 44 ] . Our model belongs to this category , and due to its recursive nature , it can extract features from an arbitrary number of pyramid levels ( see Section 3 ) . ( 2 ) Encoding schemes take advantage of the inherently hierarchical nature of deep neural nets , and reuse features from different layers , since they contain information of different scale [ 24 , 41 , 10 ] . ( 3 ) Encoding-Decoding schemes follow up the feed-forward processing ( encoding ) with a decoder , that gradually recovers the spatial resolution of early feature maps , by combining coarse and fine features [ 52 , 40 ] . ( 4 ) Spatial modules are incorporated into the forward pass , to alter feature extraction between layers [ 71 , 9 , 65 ] . Computational efficiency . There are multiple ways to adjust the computational cost of deep neural networks . We organize them into four categories . ( 1 ) Compression methods aim to remove redundancy from already trained models [ 38 , 26 , 72 ] . ( 2 ) Lightweight design strategies are used to replace network components with computationally lighter counterparts [ 32 , 30 , 51 , 27 , 66 ] . ( 3 ) Partial computation methods selectively utilize parts of a network , creating paths of computation with different costs [ 37 , 20 , 73 , 29 , 69 ] . ( 4 ) Attention methods selectively process parts of the input , based on their importance for the task at hand [ 49 , 39 , 33 , 56 ] . This is the strategy we follow in our architecture . 3 Architecture . 3.1 Processing flow . We present our architecture by walking through the example in Figure 2 , where we process an image with original resolution of 896× 896 px ( 1© in the top left corner ) . In the first level , we downscale the image to 224× 224 px and pass it through the feature extraction module , in order to produce a feature vector V1 that contains a coarse description of the original image . To proceed to the next level , we feed an intermediate feature map , F1 , from the feature extraction module to the location module , which considers a number of candidate locations described by F1 , and predicts their importance ( in this particular example , the candidate locations form a 2× 2 regular grid ( 2© ) , and the location module yields 4 predictions ) . We express region importance as attendance probability , which parametrizes a categorical distribution used for sampling without replacement ; in our current example , we sample 2 locations ( 3© ) . In the 2nd processing level , we crop the selected regions from the full-resolution image , resize them to 224× 224 px , and feed them to the feature extraction module to obtain the corresponding feature vectors ( here V21 and V23 ) . The original input resolution allows us to move to a 3rd processing level , where we feed F21 and F23 to the location module , leading to 2 Categorical distributions . We sample 1 location from each one of them to get V35 and V39 . Features extracted from all levels are passed through the positional encoding module , which injects information about the spatial position and scale of the image regions the features describe . The resulting vectors , { V ′∗ } , are averaged ( 4© ) into a single comprehensive representation , V ′ , that is fed to the classification module for the final prediction . 3.2 Modules . The feature extraction module receives an image of fixed size as input , and outputs a feature vector V and an intermediate spatial representation F . The input size , a hyperparameter we call base resolution , defines the minimum amount of information that can be processed . Hence , it constrains the minimum cost that our model has to pay in computational and memory resources . In our experiments , feature extraction modules are implemented using CNNs . The location module predicts K = n2 probabilities of a Categorical distribution over the locations within a given n× n grid of candidate image regions . It receives a feature map of size n× n× c as input , where each 1× 1× c vector describes the corresponding location of the grid . The feature map is passed through a series of 1× 1 convolutions ( contextual information is infused as well , e.g. , via squeeze and excitation [ 28 ] ) , yielding K logits , which are transformed to relative region importance via a softmax layer . The positional encoding module receives a feature vector f and a positional encoding vector p , and combines them ( e.g. , through a fully connected layer ) to an output feature vector f ′ . We use a variant of the fixed positional encodings based on sine and cosine functions introduced in [ 63 ] . Instead of a single dimension of time , we have three : two spatial dimensions and scale . The classification module projects the final feature vector ( e.g. , via a linear layer ) to classification logits . We provide the exact module architectures in Appendix 2.1 , along with justification of our design choices . 4 Training . 4.1 Learning rule . Our model is not end-to-end differentiable because of location sampling . We address this problem using a variant of the REINFORCE [ 68 ] learning rule : LF = 1 N ·M N ·M∑ i=1 [ ∂ log p ( yi|li , xi , w ) ∂w + λf ( Ri − b ) ∂ log p ( li|xi , w ) ∂w ] ( 1 ) where xi is the i-th image , yi is its label , and w are the parameters of our model . p ( li|xi , w ) is the probability that the sequence of locations li is attended for image xi , and p ( yi|li , xi , w ) is the probability of predicting the correct label after attending to li . N ·M is the total number of examples used for each update . The size of our original batch B is N , and we derive ( 1 ) using a Monte Carlo estimator with M samples to approximate the expectation∑ li p ( l i|xi , w ) [ ∂ log p ( yi|li , xi , w ) ∂w + log p ( yi|l i , xi , w ) ∂ log p ( li|xi , w ) ∂w ] for each image xi in B . To reduce the variance of the estimator , we replace log p ( yi|li , xi , w ) with a discrete indicator function Ri , which is equal to 1 for correct predictions and 0 otherwise [ 2 ] . To the same end , we use baseline b , which corresponds to the exponential moving average of the mean reward Ri ∀i , and is updated after processing each training batch [ 70 ] : bn = 0.9 · bn−1 + 0.1 · 1 NM NM∑ i=1 Rni ( 2 ) where Rni is the reward for the i-th image in the n-th batch . For simplicity , we drop the subscript of bn in ( 1 ) . λf is a weighting hyperparameter . We provide a detailed derivation of our learning rule in Appendix 1.1 . The first term of LF is used to update the parameters in order to maximize the probability of the correct label . The second term is used to update the location selection process , according to the utility of the attended sequence li in the prediction of the correct label .

Summary: The paper reduces the computational cost of processing large images by leveraging a multi-scale hard attention mechanism. Starting from a low resolution version of an image, it progressively identifies and zooms in the most discriminative areas of the image. By ignoring most of the irrelevant content of the image, it processes images faster than strong traditional approaches, while achieving higher accuracy than these baselines thanks to a novel regularization mechanism.