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: This paper presents a learning framework for a hard attention mechanism. The glimpses captured by the attention mechanism are guided by the goal of minimizing output uncertainty for a downstream task such as classification. The authors pose this problem in a probabilistic framework which is based on Bayesian optimal experimental design (BOED). They devise a tractable approximation to the entropy over images and glimpse sequences and search for the glimpse sequences which minimize the output entropy of a recurrent classification model.