Is Attention Better Than Matrix Decomposition?

1 INTRODUCTION . Since self-attention and transformer ( Vaswani et al. , 2017 ) showed significant advantages over recurrent neural networks and convolutional neural networks in capturing long-distance dependencies , attention has been widely adopted by computer vision ( Wang et al. , 2018 ; Zhang et al. , 2019a ) and natural language processing ( Devlin et al. , 2019 ) for global information mining . However , is hand-crafted attention irreplaceable when modeling the global context ? This paper focuses on a new approach to design global context modules . The key idea is , if we formulate the inductive bias like the global context into an objective function , the optimization algorithm to minimize the objective function can construct a computational graph , i.e. , the architecture we need in the networks . We particularize this idea by developing a counterpart for the most representative global context module , self-attention . Considering extracting global information in the networks as finding a dictionary and the corresponding codes to capture the inherent correlation , we model the context discovery as low-rank completion of the input tensor and solve it via matrix decomposition . This paper then proposes a global correlation block , Hamburger , by employing matrix decomposition to factorize the learned representation into sub-matrices so as to recover the clean low-rank signal subspace . The iterative optimization algorithm to solve matrix decomposition defines the central computational graph , i.e. , Hamburger ’ s architecture . Our work takes advantage of the matrix decomposition models as the foundation of Hamburger , including Vector Quantization ( VQ ) ( Gray & Neuhoff , 1998 ) , Concept Decomposition ( CD ) ( Dhillon & Modha , 2001 ) , and Non-negative Matrix Factorization ( NMF ) ( Lee & Seung , 1999 ) . Additionally , instead of directly applying Back-Propagation Through Time ( BPTT ) algorithm ( Werbos et al. , 1990 ) to differentiate the iterative optimization , we adopt a truncated BPTT algorithm , i.e. , one-step gradient , to back-propagate the gradient effectively . We illustrate the advantages of Hamburger in ∗Equal first authorship †Corresponding author the fundamental vision tasks where global information has been proven crucial , including semantic segmentation and image generation . The experiments prove that optimization-designed Hamburger can perform competitively with state-of-the-art attention models when avoiding the unstable gradient back-propagated through the iterative computational graph of MD . Hamburger sets new state-ofthe-art records on the PASCAL VOC dataset ( Everingham et al. , 2010 ) and PASCAL Context dataset ( Mottaghi et al. , 2014 ) for semantic segmentation and surpasses existing attention modules for GANs in the large scale image generation on ImageNet ( Deng et al. , 2009 ) . The contributions of this paper are listed as follows : • We show a white-box approach to design global information blocks , i.e. , by turning the optimization algorithm that minimizes an objective function , in which modeling the global correlation is formulated as a low-rank completion problem , into the architecture . • We propose Hamburger , a light yet powerful global context module with O ( n ) complexity , surpassing various attention modules on semantic segmentation and image generation . • We figure out that the main obstacle of applying MD in the networks is the unstable backward gradient through its iterative optimization algorithm . As a pragmatic solution , the proposed one-step gradient facilitates the training of Hamburger with MDs . 2 METHODOLOGY . 2.1 WARM UP . Since matrix decomposition is pivotal to the proposed Hamburger , we first review the idea of matrix decomposition . A common view is that matrix decomposition factorizes the observed matrix into a product of several sub-matrices , e.g. , Singular Value Decomposition . However , a more illuminating perspective is that , by assuming the generation process , matrix decomposition acts as the inverse of the generation , disassembling the atoms that make up the complex data . From the reconstruction of the original matrices , matrix decomposition recovers the latent structure of observed data . Suppose that the given data are arranged as the columns of a large matrix X = [ x1 , · · · , xn ] ∈ Rd×n . A general assumption is that there is a low-dimensional subspace , or a union of multiple subspaces hidden in X . That is , there exists a dictionary matrix D = [ d1 , · · · , dr ] ∈ Rd×r and corresponding codes C = [ c1 , · · · , cn ] ∈ Rr×n that X can be expressed as generation←−−−−−−−−− X = X̄ + E = DC + E , −−−−−−−−−→ decomposition ( 1 ) where X̄ ∈ Rd×n is the output low-rank reconstruction , and E ∈ Rd×n is the noise matrix to be discarded . Here we assume that the recovered matrix X̄ has the low-rank property , such that rank ( X̄ ) ≤ min ( rank ( D ) , rank ( C ) ) ≤ r min ( d , n ) . ( 2 ) Different MDs can be derived by assuming structures to matrices D , C , and E ( Kolda & Bader , 2009 ; Udell et al. , 2016 ) . MD is usually formulated as an objective with various constraints and then solved by optimization algorithms , with classic applications to image denoising ( Wright et al. , 2009 ; Lu et al. , 2014 ) , inpainting ( Mairal et al. , 2010 ) , and feature extraction ( Zhang et al. , 2012 ) . 2.2 PROPOSED METHOD . We focus on building global context modules for the networks without painstaking hand-crafted design . Before starting our discussion , we review the representative hand-designed context block self-attention pithily . The attention mechanism aims at finding a group of concepts for further conscious reasoning from massive unconscious context ( Xu et al. , 2015 ; Bengio , 2017 ; Goyal et al. , 2019 ) . As a representative , self-attention ( Vaswani et al. , 2017 ) is proposed for learning long-range dependencies in machine translation , Attention ( Q , K , V ) = softmax ( QK > √ d ) V , ( 3 ) where Q , K , V ∈ Rn×d are features projected by linear transformations from the input . Selfattention extracts global information via attending all tokens at a time rather than the typical one-byone processing of recurrent neural networks . Though self-attention and its variants achieved great success , researchers are confronted with ( 1 ) developing new global context modules based on self-attention , typically via hand-crafted engineering , and ( 2 ) explaining why current attention models work . This paper bypasses both issues and finds a method to easily design global context modules via a well-defined white-box toolkit . We try to formulate the human inductive bias , like the global context , as an objective function and use the optimization algorithm to solve such a problem to design the module ’ s architecture . The optimization algorithm creates a computational graph , takes some input , and finally outputs the solution . We apply the computational graph of optimization algorithms for the central part of our context module . Based on this approach , we need to model the networks ’ global information issue as an optimization problem . Take the convolutional neural networks ( CNN ) as an example for further discussion . The networks output a tensor X ∈ RC×H×W after we feed into an image . Since the tensor can be seen as a set of HW C-dimensional hyper-pixels , we unfold the tensor into a matrix X ∈ RC×HW . When the module learns the long-range dependencies or the global context , the hidden assumption is that the hyper-pixels are inherently correlated . For the sake of simplicity , we assume that hyper-pixels are linearly dependent , which means that each hyper-pixel in X can be expressed as the linear combination of bases whose elements are typically much less than HW . In the ideal situation , the global information hidden in X can be low-rank . However , due to vanilla CNN ’ s poor ability to model the global context ( Wang et al. , 2018 ; Zhang et al. , 2019a ) , the learned X is usually corrupted with redundant information or incompleteness . The above analysis suggests a potential method to model the global context , i.e. , by completing the low-rank part X̄ in the unfolded matrix X and discarding the noise part E , using the classic matrix decomposition models described in Eq . ( 1 ) , which filters out the redundancy and incompleteness at the same time . We thus model learning the global context as a low-rank completion problem with matrix decomposition as its solution . Using the notion of Sec . 2.1 , the general objective function of matrix decomposition is min D , C L ( X , DC ) +R1 ( D ) +R2 ( C ) ( 4 ) where L is the reconstruction loss , R1 andR2 are regularization terms for the dictionary D and the codes C. Denote the optimization algorithm to minimize Eq . ( 4 ) asM.M is the core architecture we deploy in our global context module . To help readers further understand this modeling , We also provide a more intuitive illustration in Appendix G. In the later sections , we introduce our global context block , Hamburger , and then discuss detailed MD models and optimization algorithms forM . Finally , we handle the gradient issue for back-propagation through matrix decomposition . 2.2.1 HAMBURGER . Hamburger consists of one slice of “ ham ” ( matrix decomposition ) and two slices of “ bread ” ( linear transformation ) . As the name implies , Hamburger first maps the input Z ∈ Rdz×n into feature space with a linear transformation Wl ∈ Rd×dz , namely “ lower bread ” , then uses matrix decomposition M to solve a low-rank signal subspace , corresponding to the “ ham ” , and finally transforms extracted signals into the output with another linear transformation Wu ∈ Rdz×d , called “ upper bread ” , H ( Z ) = WuM ( WlZ ) , ( 5 ) where M is matrix decomposition to recover the clear latent structure , functioning as a global nonlinearity . Detailed architectures ofM , i.e. , optimization algorithms to factorize X , are discussed in Sec . 2.2.2 . Fig . 1 describes the architecture of Hamburger , where it collaborates with the networks via Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) , a skip connection , and finally outputs Y , Y = Z + BN ( H ( Z ) ) . ( 6 ) 2.2.2 HAMS . This section describes the structure of “ ham ” , i.e. , M in Eq . ( 5 ) . As discussed in the previous section , by formulating the global information discovery as an optimization problem of MD , algorithms to solve MD naturally composeM.M takes the output of “ lower bread ” as its input and computes a low-rank reconstruction as its output , denoted as X and X̄ , respectively . M ( X ) = X̄ = DC . ( 7 ) We investigate two MD models forM , Vector Quantization ( VQ ) , and Non-negative Matrix Factorization ( NMF ) to solve D and C and reconstruct X̄ , while leaving Concept Decomposition ( CD ) to Appendix B . The selected MD models are introduced briefly because we endeavor to illustrate the importance of the low-rank inductive bias and the optimization-driven designing method for global context modules rather than any specific MD models . It is preferred to abstract the MD part as a whole , i.e. , M in the context of this paper , and focus on how Hamburger can show the superiority in its entirety . Vector Quantization Vector Quantization ( VQ ) ( Gray & Neuhoff , 1998 ) , a classic data compression algorithm , can be formulated as an optimization problem in term of matrix decomposition : min D , C ‖X −DC‖F s.t . ci ∈ { e1 , e2 , · · · , er } , ( 8 ) where ei is the canonical basis vector , ei = [ 0 , · · · , 1 , · · · , 0 ] > ith . The solution to minimize the objective in Eq . ( 8 ) is K-means ( Gray & Neuhoff , 1998 ) . However , to ensure that VQ is differentiable , we replace the hard arg min and Euclidean distance with softmax and cosine similarity , leading to Alg . 1 , where cosine ( D , X ) is a similarity matrix whose entries satisfy cosine ( D , X ) ij = d > i xj ‖d‖‖x‖ , and softmax is applied column-wise and T is the temperature . Further we can obtain a hard assignment by a one-hot vector when T → 0 . Algorithm 1 Ham : Soft VQ Input X . Initialize D , C. for k from 1 to K do C ← softmax ( 1T cosine ( D , X ) ) D ←XC > diag ( C1n ) −1 end for Output X̄ = DC . Algorithm 2 Ham : NMF with MU Input X . Initialize non-negative D , C for k from 1 to K do Cij ← Cij ( D > X ) ij ( D > DC ) ij Dij ←Dij ( XC > ) ij ( DCC > ) ij end for Output X̄ = DC . Non-negative Matrix Factorization If we impose non-negative constraints on the dictionary D and the codes C , it leads to Non-negative Matrix Factorization ( NMF ) ( Lee & Seung , 1999 ) : min D , C ‖X −DC‖F s.t.Dij ≥ 0 , Cjk ≥ 0 . ( 9 ) To satisfy the non-negative constraints , we add a ReLU non-linearity before putting X into NMF . We apply the Multiplicative Update ( MU ) rules ( Lee & Seung , 2001 ) in Alg . 2 to solve NMF , which guarantees the convergence . As white-box global context modules , VQ , CD , and NMF are straightforward and light , showing remarkable efficiency . They are formulated into optimization algorithms that mainly consist of matrix multiplications with the complexity O ( ndr ) , much cheaper than complexity O ( n2d ) in self-attention as r n. All three MDs are memory-friendly since they avoid generating a large n× n matrix as an intermediate variable , like the product of Q and K of self-attention in Eq . ( 3 ) . In the later section , our experiments prove MDs are at least on par with self-attention , though the architectures ofM are created by optimization and look different from classic dot product self-attention .

Summary: This paper proposes to use matrix decomposition to construct low-rank representations to find the long-distance correlations in context, which is demonstrated more effective than popular self-attention mechanism. Combining linear transformation and matrix decomposition (core part), authors design Hamburger block to model global dependencies from input as residual output. The authors propose differentiable modified Vector Quantization and Non-negative Matrix Factorization to perform matrix decomposition. They propose one-step gradient, an approximation of Back-Propagation Through Time (BPTT) algorithm, to back-propagate gradient of matrix decomposition. They conduct experiments on semantic segmentation and image generation to demonstrate the superiority of their methods regarding modelling global dependencies and computational cost.