Variational Structured Attention Networks for Dense Pixel-Wise Prediction

1 INTRODUCTION . Over the past decade , convolutional neural networks ( CNNs ) have become the privileged methodology to address computer vision tasks requiring dense pixel-wise prediction , such as semantic segmentation ( Chen et al. , 2016b ; Fu et al. , 2019 ) , monocular depth prediction ( Liu et al. , 2015 ; Roy & Todorovic , 2016 ) , contour detection ( Xu et al. , 2017a ) and normal surface computation ( Eigen et al. , 2014 ) . Recent studies provided clear evidence that attention mechanisms ( Mnih et al. , 2014 ) within deep networks are undoubtedly a crucial factor in improving the performance ( Chen et al. , 2016b ; Xu et al. , 2017a ; Fu et al. , 2019 ; Zhan et al. , 2018 ) . In particular , previous works demonstrated that deeply learned attentions acting as soft weights to interact with different deep features at each channel ( Zhong et al. , 2020 ; Zhang et al. , 2018 ; Song et al. , 2020 ) and at each pixel location ( Li et al. , 2020a ; Johnston & Carneiro , 2020 ; Tay et al. , 2019 ) permits to improve the pixel-wise prediction accuracy ( see Fig.1.a and Fig.1.b ) . Recently , Fu et al . ( 2019 ) proposed the Dual Attention Network ( DANet ) , embedding in a fully convolutional network ( FCN ) two complementary attention modules , specifically conceived to model separately the semantic dependencies associated to the spatial and to the channel dimensions ( Fig.1.c ) . Concurrently , other approaches have considered the use of structured attention models integrated within a graph network framework ( Zhang et al. , 2020 ; Chen et al. , 2019 ; Xu et al. , 2017a ) , showing the empirical advantage of adopting a graphical model to effectively capture the structured information present in the hidden layers of the neural network and thus enabling the learning of better deep feature representations . Notably , Xu et al . ( 2017a ) first introduced attention-gated conditional random fields ( AG-CRFs ) , a convolutional neural network implementing a probabilistic graphical model that considers attention variables as gates ( Minka & Winn , 2009 ) in order to learn improved deep features and effectively fuse multi-scale information . However , their structured attention model is only learned at the spatial-wise level , while channel-wise dependencies are not considered . This paper advances the state of the art in dense pixel-wise prediction by proposing a novel approach to learn more effective deep representations by integrating a structured attention model which jointly account for spatial- and channel-level dependencies using an attention tensor ( Fig.1.d ) within a CRF framework . More precisely , inspired from Xu et al . ( 2017a ) we model the attention as gates . Crucially , we address the question on how to enforce structure within these latent gates , in order to jointly model spatial- and channel-level dependencies while learning deep features . To do so , we hypothesize that the attention tensor is nothing but the sum of T rank-1 tensors , each of them being the tensor product of a spatial attention map and a channel attention vector . This attention tensor is used as a structured latent attention gate , enhancing the feature maps . We cast the inference problem into a maximum-likelihood estimation formulation that is made computationally tractable thanks to a variational approximation . Furthermore , we implement the maximum likelihood update rules within a neural network , so that they can be jointly learned with the preferred CNN front-end . We called our approach based on structured attention and variational inference VarIational STructured Attention Networks or VISTA-Net . We evaluate our method on multiple pixel-wise prediction problems , i.e . monocular depth estimation , semantic segmentation and surface normale prediction , considering six publicly available datasets , i.e . NYUD-V2 ( Silberman et al. , 2012 ) , KITTI ( Geiger et al. , 2013 ) , Pascal-Context ( Mottaghi et al. , 2014 ) , Pascal VOC2012 ( Everingham et al. , 2010 ) , Cityscape ( Cordts et al. , 2016 ) and ScanNet ( Dai et al. , 2017 ) . Our results demonstrate that VISTANet is able to learn rich deep representations thanks to the proposed structured attention and our probabilistic formulation , outperforming state-of-the-art methods . Related Work . Several works have considered integrating attention models within deep architectures to improve performance in several tasks such as image categorization ( Xiao et al. , 2015 ) , speech recognition ( Chorowski et al. , 2015 ) and machine translation ( Vaswani et al. , 2017 ; Kim et al. , 2017 ; Luong et al. , 2015 ) . Focusing on pixel-wise prediction , Chen et al . ( 2016b ) first described an attention model to combine multi-scale features learned by a FCN for semantic segmentation . Zhang et al . ( 2018 ) designed EncNet , a network equipped with a channel attention mechanism to model global context . Zhao et al . ( 2018 ) proposed to account for pixel-wise dependencies introducing relative position information in spatial dimension within the convolutional layers . Huang et al . ( 2019b ) described CCNet , a deep architecture that embeds a criss-cross attention module with the idea of modeling contextual dependencies using sparsely-connected graphs , such as to achieve higher computational efficiency . Fu et al . ( 2019 ) proposed to model semantic dependencies associated with spatial and channel dimensions by using two separate attention modules . Zhong et al . ( 2020 ) introduced a squeeze-and-attention network ( SANet ) specialized to pixel-wise prediction that takes into account spatial and channel inter-dependencies in an efficient way . Attention was first adopted within a CRF framework by Xu et al . ( 2017a ) , which introduced gates to control the message passing between latent variables and showed that this strategy is effective for contour detection . Our work significantly departs from these previous approaches , as we introduce a novel structured attention mechanism , jointly handling spatial- and channel-level dependencies within a probabilistic framework . Notably , we also prove that our model can be successfully employed in case of several challenging dense pixel-level prediction tasks . Our work is also closely related to previous studies on dual graph convolutional network ( Zhang et al. , 2019c ) and dynamic graph message passing networks ( Zhang et al. , 2020 ) , which have been successfully used for pixellevel prediction tasks . However , while they also resort on message passing for learning refined deep feature representations , they lack a probabilistic formulation . Finally , previous studies ( Xu et al. , 2017c ; Arnab et al. , 2016 ; Chen et al. , 2019 ) described CRF-based models for pixel-wise estima- tion , e.g . to learn and optimally fuse deep representations at multiple scales . However , they did not employ structured attention gates . 2 VARIATIONAL STRUCTURED ATTENTION NETWORKS : VISTA-NET . As previously discussed , we aim to enhance the learned representation by structuring the attention within a probabilistic formulation . One the one side , inducing structure in the attention mechanisms has been proven to be successful ( Fu et al. , 2019 ; Zhong et al. , 2020 ) . On the other side , probabilistic formulations combined with deep architectures are interesting for pixel-level prediction tasks ( Xu et al. , 2017b ) . Up to our knowledge , we are the first to bring together recent advances in pixel-wise prediction by formulating a novel structured attention mechanism within a probabilistic CRF-like inference framework . Inspired by Fu et al . ( 2019 ) , where two spatial- and a channel-wise fullrank tensors are computed , we opt to infer different spatial and channel attention variables.Very differently from Fu et al . ( 2019 ) , we propose to structure a generic attention tensor a of dimension W ×H × C ( widht , height , channels ) , as the sum of T rank-1 tensors : a = T∑ t=1 mt ⊗ vt , mt ∈ R1×W×H , vt ∈ RC×1×1 , ( 1 ) meaning that mt can be understood as an image of W ×H pixels and vt as a vector of dimension C , and ⊗ denotes the tensor product , in the case above leading to a 3-way tensor of dimensions W × H × C. Each of the tensor products within the sum yields a tensor of rank-1 , consequently limiting the rank of a to be at maximum T . Equation ( 1 ) is the algebraic expression of the proposed structured attention mechanism , and is the methodological foundation of VISTA-Net . Moreover , we inspire from the CRF formulation with gating variables proposed in ( Xu et al. , 2017a ) , and derive a new energy function and variational approximation to enable efficient learning and inference procedures . Additionally , this formulation allows us to consider the CRF kernels as latent variables and infer them from the data , together with the structured attention variables mt and vt. We believe learning the kernels is important because it allows the CRF to weight the information flow depending on the content of the rather than keeping the same weights for all images . We assume a generic CNN front-end providing a set of S multi-scale feature maps F = { fs } Ss=1 . To ease notation , we assume that each feature map has P pixels and C channels , but in practice these dimensions depend on the scale s. For each scale , we also consider the set of hidden variables zs corresponding to fs , and Z = { zs } Ss=1 . These hidden variables correspond to refined convolutional futures that incorporate information and attention from other feature maps , so as to better represent the key information for the pixel-level task at hand . Intuitively , the structured attention tensor should help refining the hidden variables to allow better performance at various pixel-level prediction tasks . As in ( Xu et al. , 2017a ) , for every pair of emitting e and receiving r scales , we consider a dedicated attention tensor ae , r . Very importantly , in our case this attention tensor is structured following ( 1 ) , and so we have a set of hidden spatial attention maps M = { mte , r } S , S , T e , r , t=1 and hidden channel attention vectors V = { vte , r } S , S , T e , r , t=1 . More precisely , m t e , r ∈ { 0 , 1 } P and vte , r ∈ { 0 , 1 } C are a binary spatial map and a stochastic channel-wise vector , hence ∑C c=1 v t , c e , r = 1 . In this way , we reduce ambiguity and ease the learning . This also means that the model is conceived to pay attention to only T channels of the feature map . While this could seem limiting at first glance we remark that : ( i ) the model learns which are the optimal T channels among the possible C that have to be used to refine the hidden variables and ( ii ) the posterior distribution of mt boils down to a convex combination of all channels , as it will appear clear when discussing the inference procedure . 2.1 ENERGY FUNCTION AND VARIATIONAL APPROXIMATION . Our model consists on three different latent variables : the hidden features Z , and the hidden attention maps M and vectors V. In addition , we also consider inferring the CRF kernels , denoted by K from the data . More precisely , the energy function associated to the proposed models writes : − E ( Z , M , V , K , F , Θ ) = ∑ s ∑ p , c φz ( z p , c r , f p , c r ) + ∑ e , r ∑ p , c , p′ , c′ ∑ t mt , pe , rv t , c e , rψ ( z p , c r , z p′ , c′ e , k e , p′ , c′ r , p , c ) + φk ( f p , c r , f p′ , c′ e , k e , p′ , c′ r , p , c ) , ( 2 ) where φz , φk and ψ are potentials to be defined and ke , p ′ , c′ r , p , c denotes the kernel value weighting the information flow from the ( p′ , c′ ) -th value of the feature map of scale e to the ( p , c ) -th value of the feature map of scale r. Since the exact posterior distribution is not computationally tractable , we opt to approximate it with the following family of separable distributions : p ( Z , M , V , K|F , Θ ) ≈ q ( Z , M , V , K ) = qz ( Z ) qm ( M ) qv ( V ) qk ( K ) . ( 3 ) In that case , the optimal solution for each of the factors of the distribution is to take the expectation w.r.t . to all the others , for instance : qz ( Z ) ∝ exp ( − Eqm ( M ) qv ( V ) qk ( K ) { E ( Z , M , V , K , F , Θ ) } ) . ( 4 ) It can be shown that the optimal variational factors write : qz ( z p , c r ) ∝ exp ( φz ( z p , c r , f p , c r ) + ∑ e 6=r ∑ t m̄t , pe , rv̄ t , c e , r ∑ p′ , c′ Eqzqk { ψ ( zp , cr , zp ′ , c′ e , k e , p′ , c′ r , p , c ) } ) , qm ( m t , p e , r ) ∝ exp ( mt , pe , r ∑ c v̄t , ce , r ∑ p′ , c′ Eqz , qk { ψ ( zp , cs , z p′ , c′ s′ , k e , p′ , c′ r , p , c ) } ) , qv ( v t , c e , r ) ∝ exp ( vt , ce , r ∑ p m̄t , pe , r ∑ p′ , c′ Eqz , qk { ψ ( zp , cs , z p′ , c′ s′ , k e , p′ , c′ r , p , c ) } ) , qk ( k e , p′ , c′ r , p , c ) ∝ exp ( φk ( f p , c r , f p′ , c′ e , k e , p′ , c′ r , p , c ) + ∑ t m̄t , pe , rv̄ t , c e , rEqz { ψ ( zp , cs , z p′ , c′ s′ , k e , p′ , c′ r , p , c ) } ) , ( 5 ) where m̄t , pe , r = Eqm { mt , pe , r } denotes the the posterior mean , and analogously for v̄t , ce , r . This result also implies that thanks to the variational approximation in ( 3 ) , the posterior distributions factorise in each of the variables above , e.g . qz ( Z ) = ∏S , P , C r , p , c=1 qz ( z p , c r ) . The relation between the various hidden variables as for their inference is shown in Figure 2 ( left ) . In addition , we also show the information flow between the hidden variables using arrows . Finally , on Figure 2 ( right ) we show the relation between the channel-wise and spatial attention variables and how the final structured attention tensor is computed .

Summary: This paper proposes the VarIational STructured Attention networks (VISTA-Net), which improves pervious SOTA models for dense pixel-wise prediction tasks. The proposed VISTA-Net is featured by two aspects: 1) A new structured attention is proposed, which is able to jointly model spatial-level and channel-level dependencies; 2) It incorporates the proposed structured attention with a CRF-like inference framework, which allows the probabilistic inference. Experimental studies are conducted on monocular depth estimation and semantic image segmentation, showing improved performances of VISTA-Net consistently.