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: High-resolution image classification is a hard problem for DNN due to resource constraints. The authors propose a multi-scale hard-attention architecture, TNet, for high-resolution information, by only visiting the most informative image regions, to achieve better accuracy vs complexity tradeoff. The architecture achieves a better tradeoff than other baselines on both ImageNet and fMoW datasets.