Near-Optimal Glimpse Sequences for Training Hard Attention Neural Networks

1 INTRODUCTION . Attention can be defined as the “ allocation of limited cognitive processing resources ” ( Anderson , 2005 ) . In humans the density of photoreceptors varies across the retina . It is much greater in the centre ( Bear et al. , 2007 ) and covers an approximately 210 degree field of view ( Traquair , 1949 ) . This means that the visual system is a limited resource with respect to observing the environment and that it must be allocated , or controlled , by some attention mechanism . We refer to this kind of controlled allocation of limited sensor resources as “ hard ” attention . This is in contrast with “ soft ” attention , the controlled application of limited computational resources to full sensory input . Hard attention can solve certain tasks using orders of magnitude less sensor bandwidth and computation than the alternatives ( Katharopoulos & Fleuret , 2019 ; Rensink , 2000 ) . It therefore may enable the use of modern approaches to computer vision in low-power settings such as mobile devices . This paper focuses on the application of hard attention in image classification . Our model of attention ( shown in Fig . 1 ) is as follows : a recurrent neural network ( RNN ) is given T steps to classify some unchanging input image . Before each step , the RNN outputs the coordinates of a pixel in the image . A patch of the image centered around this pixel is then fed into the RNN . We call this image patch a glimpse , and the coordinates a glimpse location . As such , the RNN controls its input by selecting each glimpse location , and this decision can be based on previous glimpses . After T steps , the RNN ’ s hidden state is mapped to a classification output . As with most artificial hard attention mechanisms ( Mnih et al. , 2014 ; Ba et al. , 2014 ) , this output is not differentiable with respect to the sequence of glimpse locations selected . This makes training with standard gradient backpropagation impossible , and so high variance gradient estimators such as REINFORCE ( Williams , 1992 ) are commonly used instead ( Mnih et al. , 2014 ; Ba et al. , 2014 ) . The resulting noisy gradient estimates make training difficult , especially for large T . In order to improve hard attention training , we take inspiration from neuroscience literature which suggests that visual attention is directed so as to maximally reduce entropy in an agent ’ s world model ( Bruce & Tsotsos , 2009 ; Itti & Baldi , 2009 ; Schwartenbeck et al. , 2013 ; Feldman & Friston , 2010 ) . There is a corresponding mathematical formulation of such an objective , namely Bayesian optimal experimental design ( BOED ) ( Chaloner & Verdinelli , 1995 ) . BOED tackles the problem of designing an experiment to maximally reduce uncertainty in some unknown variable . When classifying an image with hard visual attention , the ‘ experiment ’ is the process of taking a glimpse ; the ‘ design ’ is the glimpse location ; and the unknown variable is the class label . In general , BOED is applicable only when a probabilistic model of the experiment exists . This could be , for example , a prior distribution over the class label and a generative model for the observed image patch conditioned on the class label and glimpse location . We leverage generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) to provide such a model . We use methodology from BOED to introduce the following training procedure for hard attention networks , which we call partial supervision by near-optimal glimpse sequences ( PS-NOGS ) . 1 . We assume that we are given an image classification task and a corresponding labelled dataset . Then , for a subset of the training images , we determine an approximately optimal ( in the BOED sense ) glimpse location for a hard attention network to attend to at each time step . We refer to the resulting sequences of glimpse locations as near-optimal glimpse sequences . Section 4 describes our novel method to generate them . 2 . We use these near-optimal glimpse sequences as an additional supervision signal for training a hard attention network . Section 5 introduces our novel training objective for this . We empirically investigate the performance of PS-NOGS and find that it leads to faster training than our baselines , and qualitatively different behaviour with competitive accuracy . We validate the use of BOED to generate glimpse sequences through comparisons with supervision both by hand-crafted glimpse sequences , and by glimpse sequences sampled from a trained hard attention network . 2 HARD ATTENTION . Given an image , I , we consider the task of inferring its label , θ . We use an architecture based on that of Mnih et al . ( 2014 ) , shown in Fig . 1 . It runs for a fixed number of steps , T . At each step t , the RNN samples a glimpse location , lt , from a distribution conditioned on previous glimpses via the RNN ’ s hidden state . A glimpse , in the form of a contiguous square of pixels , is extracted from the image at this location . We denote this yt = ffovea ( I , lt ) . An embedding of yt and lt is then input to the RNN . After T glimpses , the network outputs a classification distribution qφ ( θ|y1 : T , l1 : T ) , where φ are the learnable network parameters . Mnih et al . ( 2014 ) use glimpses consisting of three image patches at different resolutions , but the architectures are otherwise identical . As it directly processes only a fraction of an image , this architecture is suited to low-power scenarios such as use on mobile devices . During optimisation , gradients can not be computed by simple backpropagation since ffovea is nondifferentiable . An alternative , taken by Mnih et al . ( 2014 ) and others in the literature ( Ba et al. , 2014 ; Sermanet et al. , 2014 ) , is to obtain high-variance gradient estimates using REINFORCE ( Williams , 1992 ) . Although these are unbiased , their high variance has made scaling beyond simple problems such as digit classification ( Netzer et al. , 2011 ) challenging . Section 7 describes alternatives ( Ba et al. , 2015 ; Lawson et al. , 2018 ) to training with REINFORCE , but similar problems with scalability exist . This has led many studies to focus on easing the learning task by altering the architecture : e.g. , to process a downsampled image before selecting glimpse locations ( Ba et al. , 2014 ; Sermanet et al. , 2014 ; Katharopoulos & Fleuret , 2019 ) . We summarise these innovations in Section 7 but they tend to be less suitable for low-power computation . We therefore believe that improved training of the architecture in Fig . 1 is an important research problem , and it is the focus of this paper . 3 BAYESIAN OPTIMAL EXPERIMENTAL DESIGN . Designing an experiment to be maximally informative is a fundamental problem that applies as much to tuning the parameters of a political survey ( Warwick & Lininger , 1975 ) as to deciding where to direct attention to answer a query . BOED ( Chaloner & Verdinelli , 1995 ) provides a unifying framework for this by allowing a formal comparison of possible experiments under problem-specific prior knowledge . Consider selecting the design , l , of an experiment to infer some unknown parameter , θ . For example , θ may be the median lethal dose of a drug , and l the doses of this drug given to various groups of rats ( Chaloner & Verdinelli , 1995 ) . Alternatively , as we consider in this paper , θ is the class label of an image and l determines which part of the image we observe . The experiment results in a measurement of y ∼ p ( y|l , θ ) . Following the previous examples , y could be the number of rats which die in each group or the observed pixel values . Given a prior distribution over θ and knowledge of p ( y|l , θ ) , we can use the measurement to infer a posterior distribution over θ using Bayes ’ rule : p ( θ|y , l ) = p ( y|l , θ ) p ( θ ) ∫ p ( y|l , θ ) p ( θ ) dθ . The aim of our experiment is to infer θ , and so a well designed experiment will reduce the uncertainty about θ by as much as possible . The uncertainty after the experiment can be quantified by the Shannon entropy in the posterior , H [ p ( θ|y , l ) ] = Ep ( θ|y , l ) [ − log p ( θ|y , l ) ] . ( 1 ) To maximally reduce the uncertainty , we wish to select l to minimise this posterior entropy . However , the design of the experiment must be chosen before y is measured and so we can not evaluate the posterior entropy exactly . Instead , we minimise an expectation of it over p ( y|l ) = Ep ( θ ) [ p ( y|l , θ ) ] , the marginal distribution of y . This is the expected posterior entropy , or EPE . EPE ( l ) = Ep ( y|l ) [ H [ p ( θ|y , l ) ] ] . ( 2 ) Above , we considered the case of selecting a one-off design for an experiment , such as taking a single glimpse . For the case where a sequence of glimpses can be taken , we need sequential experimental design . In this scenario , the choice of design lt can be informed by the designs and outcomes of previous experiments , l1 : t−1 and y1 : t−1 . The marginal distribution over outcomes is therefore p ( yt|l1 : t , y1 : t−1 ) rather than p ( yt|lt ) . Similarly , the posterior after observing yt is p ( θ|l1 : t , y1 : t ) . Therefore , in the sequential case which we consider throughout the rest of the paper , we greedily minimise the following form of the EPE on each iteration : EPEy1 : t−1 , l1 : t−1 ( lt ) = Ep ( yt|y1 : t−1 , l1 : t ) [ H [ p ( θ|y1 : t , l1 : t ) ] ] . ( 3 ) To summarise , sequential BOED involves , at each time t , selecting lt = arg minlt EPEy1 : t−1 , l1 : t−1 ( lt ) and then performing the experiment with design lt to observe yt . 4 GENERATING NEAR-OPTIMAL GLIMPSE SEQUENCES . Role of BOED pipeline To reiterate the outline of our method , we first annotate a portion of the training data with glimpse sequences , and then in the second stage use these to speed up the training of a hard attention mechanism . This section details our BOED pipeline for the first stage . EPE estimator BOED requires a probabilistic model of the measurements and parameters we wish to infer . That is , we need to define p ( θ , y1 : t|l1 : t ) for any l1 : t. To do so in the visual attention setting , we first define p ( θ , I ) to be the intractable joint distribution over labels and images from which our training and test data originate . To be consistent with our definition in Section 2 of y as a deterministic function of I and l , we then define p ( yi|I , li ) to be a Dirac-delta distribution on ffovea ( I , li ) . The joint distribution is then p ( θ , y1 : t|l1 : t ) = ∫ p ( θ , I ) t∏ i=1 p ( yi|I , li ) dI . ( 4 ) Given this joint distribution , EPEy1 : t−1 , l1 : t−1 ( lt ) is well defined but intractable in general . We therefore consider how to approximate it . To simplify our method for doing so , we first rearrange the expression given in Eq . ( 3 ) so that the expectation is over I rather than yt . Taking advantage of the fact that yi is a deterministic function of I and li allows it to be rewritten as follows ( proof in the appendix ) . Defining ffovea ( I , l1 : t ) = { ffovea ( I , l1 ) , . . . , ffovea ( I , lt ) } , EPEy1 : t−1 , l1 : t−1 ( lt ) = Ep ( I|y1 : t−1 , l1 : t−1 ) [ H [ p ( θ|ffovea ( I , l1 : t ) , l1 : t ) ] ] . ( 5 ) Given this form of the expected posterior entropy , we can approximate it if we can leverage the dataset to obtain : • a learned attentional variational posterior , gAVP ( θ|y1 : t , l1 : t ) ≈ p ( θ|y1 : t , l1 : t ) , • and stochastic image completion distribution rimg ( I|y1 : t−1 , l1 : t−1 ) ≈ p ( I|y1 : t−1 , l1 : t−1 ) . We expand on the form of each of these approximations later in this section . First , combining them with Eq . ( 5 ) and using a Monte Carlo estimate of the expectation yields our estimator for the EPE : EPEy1 : t−1 , l1 : t−1 ( lt ) ≈ 1 N N∑ n=1 H [ gAVP ( θ|ffovea ( I ( n ) , l1 : t ) , l1 : t ) ] ( 6 ) with I ( 1 ) , . . . , I ( N ) ∼ rimg ( I|y1 : t−1 , l1 : t−1 ) . Overview of BOED pipeline We select lt with a grid search . That is , denoting the set of allowed values of lt as L , we compute our approximation of EPEy1 : t−1 , l1 : t−1 ( lt ) for all lt ∈ L. We then select the value of lt for which this is least . To do so , our full BOED pipeline is as follows . 1 . Sample I ( 1 ) . . . , I ( N ) ∼ rimg ( I|y1 : t−1 , l1 : t−1 ) . 2 . For each lt ∈ L , approximate the expected posterior entropy with Eq . ( 6 ) . 3 . Select the value of lt for which this approximation is least . Repeating these steps for t = 1 , . . . , T yields a near-optimal glimpse sequence l1 : T for image I . Figure 2 shows an example of this process . We must do this for all images in some subset of a dataset to be able to partially supervise hard attention training as described in Section 5 . We now describe the form of gAVP ( the attentional variational posterior ) and rimg ( stochastic image completion ) . Attentional variational posterior In this section we introduce our novel approach for efficiently approximating the intractable posterior p ( θ|y1 : t , l1 : t ) . We train a convolutional neural network ( CNN ) to map from a sequence of glimpses , y1 : t , and their locations , l1 : t , to gAVP ( θ|y1 : t , l1 : t ) , an approximation of this posterior . We call this the attentional variational posterior CNN ( AVP-CNN ) . To allow a single CNN to cope with varying y1 : t , l1 : t , and even varying t , we embed its input as shown in Fig . 3 . Essentially , l1 : t is used to create an image-sized mask which is 1 for observed pixels and 0 for unobserved pixels . Elementwise multiplication of this mask with the input image sets unobserved pixels to zero . The mask is then concatenated as an additional channel . This embedding naturally maintains spatial information while enforcing an invariance to permutations of the glimpse sequence . We use a Densenet-121 ( Huang et al. , 2017 ) CNN architecture ( pretrained on ImageNet ( Deng et al. , 2009 ) ) to map from this embedding to a vector of probabilities representing gAVP . We train the network to minimise the KL divergence between its output and p ( θ|y1 : t , l1 : t ) . That is , DKL ( p ( θ|y1 : t , l1 : t ) ||gAVP ( θ|y1 : t , l1 : t ) ) . To ensure that gAVP is close for all t , l1 : t and y1 : t , the loss used is an expectation of this KL divergence over p ( y1 : t|l1 : t ) u ( t , l1 : t ) . We factorise u ( t , l1 : t ) as u ( t ) ∏t i=1 u ( li ) where , so that all times and glimpse locations are weighted equally in the loss , u ( t ) is a uniform distribution over 1 , . . . , T and u ( li ) is a uniform distribution over all image locations . Denoting the network parameters λ , the gradient of this loss is ∂ ∂λ Lλ = Ep ( θ , y1 : t|l1 : t ) u ( t , l1 : t ) [ − ∂ ∂λ log gλAVP ( θ|y1 : t , l1 : t ) ] . ( 7 ) This gradient is the same as that of a cross-entropy loss on data sampled from p ( θ , y1 : t|l1 : t ) u ( t , l1 : t ) , and can be approximated by a Monte Carlo estimate . Our approximation of the EPE in Eq . ( 6 ) involves the entropy of gAVP . Since gAVP is a categorical distribution , this is simply computed analytically . This amortised approximation of the posterior entropy is inspired by Foster et al . ( 2019 ) , but has two important differences to their estimator : • Foster et al . learn a mapping from yt to g ( θ|y1 : t , l1 : t ) , sharing information between “ nearby ” samples of yt to reduce the computational cost of the experimental design . Our AVPCNN takes this amortization further by learning a single mapping from t , l1 : t and y1 : t to gAVP ( θ|y1 : t , l1 : t ) , which yields significant further efficiency gains in our setting . • Whereas we approximate H [ p ] with H [ gAVP ] = EgAVP [ − log gAVP ] , Foster et al . use Ep [ − log g ] . This provides an upper bound on H [ p ] but is not applicable in our case as we can not sample from p ( θ|y1 : t , l1 : t ) . Both approximations are exact when gAVP = p. Stochastic image completion We considered numerous ways to form rimg ( I|y1 : t−1 , l1 : t−1 ) including inpainting ( Pathak et al. , 2016 ; Isola et al. , 2017 ) and Markov chain Monte Carlo in a generative model . Future research in generative modelling may provide alternatives to this component of our method but , for now , we choose to represent rimg using a technique we developed based on image retrieval ( Jégou et al. , 2010 ) . Of the methods we considered , this gave the best trade-off between speed and sample quality . It involves creating an empirical image distribution with 1.5 million images for each experiment using GANs with publicly available pre-trained weights ( StyleGAN ( Karras et al. , 2018 ) for CelebA-HQ and FineGAN ( Singh et al. , 2019 ) for Caltech-UCSD Birds ) . We note that the use of pre-trained models makes test leakage possible but verify in Appendix C.4 that this is unlikely to impact our results . During sampling , the database is searched for images that ‘ match ’ the previous glimpses ( y1 : t−1 and l1 : t−1 ) . How well these glimpses match a database image , I ′ , is measured by the squared distance in pixel space at glimpse locations : ∑t−1 i=1 ‖yi − ffovea ( I ′ , li ) ‖ 2 2 . This distance defines a probability distribution over the images in the database . To reduce computation , we first compare approximations of the observed parts of each image using principal component analysis ( Jolliffe , 2011 ) , and compute exact distances only when these are close . The overall procedure to sample from rimg corresponds to importance sampling ( Arulampalam et al. , 2002 ) in a model where p ( yt|I , lt ) is relaxed from a Dirac-delta distribution to a Gaussian . See the appendix for details .

Summary: The paper trains hard attention for image classification. The network is partially supervised by attention locations proposed to maximally reduce the entropy of the image label distribution. To propose these locations, the method needs an already trained image classifier conditioned on glimpses and their locations. Additionally, the method needs a generator of images, conditioned on the glimpses and their locations. This generator is approximated by searching a set of 1.5 million pre-generated images for close matches.