Kronecker Attention Networks

1 INTRODUCTION . Deep learning networks with attention operators have demonstrated great capabilities of solving challenging problems in various tasks such as computer vision ( Xu et al. , 2015 ; Lu et al. , 2016 ) , natural language processing ( Bahdanau et al. , 2015 ; Vaswani et al. , 2017 ) , and network embedding ( Veličković et al. , 2017 ) . Attention operators are capable of capturing long-range relationships and brings significant performance boosts ( Li et al. , 2018 ; Malinowski et al. , 2018 ) . The application scenarios of attention operators range from 1-D data like texts to high-order and high-dimensional data such as images and videos . However , attention operators suffer from the excessive usage of computational resources when applied on high-order or high-dimensional data . The memory and computational cost increases dramatically with the increase of input orders and dimensions . This prevents attention operators from being applied in broader scenarios . To address this limitation , some studies focus on reducing spatial sizes of inputs such as down-sampling input data ( Wang et al. , 2018 ) or attending selected part of data ( Huang et al. , 2018 ) . However , such kind of methods inevitably results in information and performance loss . In this work , we propose novel and efficient attention operators , known as Kronecker attention operators ( KAOs ) , for high-order data , which avoid flattening and operate on high-order data directly . Experimental results show that KAOs are as effective as original attention operators , while dramatically reducing the amount of required computational resources . In particular , we employ KAOs to design a family of efficient modules , leading to our compact deep models known as Kronecker attention networks ( KANets ) . KANets significantly outperform prior compact models on the image classification task , with fewer parameters and less computational cost . We also investigate the above problem from a probabilistic perspective . Specifically , regular attention operators flatten the data and assume the flattened data follow multivariate normal distributions . This assumption not only results in high computational cost and memory usage , but also fails to preserve the spatial or spatial-temporal structures of data . Our KAOs , instead , use matrix-variate normal distributions to model the data , where the Kronecker covariance structure is able to capture relationships among spatial or spatial-temporal dimensions . 2 BACKGROUND AND RELATED WORK . In this section , we describe the attention and related non-local operators . 2.1 ATTENTION OPERATOR . The inputs to an attention operator include a query matrix Q = [ q1 , q2 , ⋯ , qm ] ∈ Rd×m , a key matrix K = [ k1 , k2 , ⋯ , kn ] ∈ Rd×n , and a value matrix V = [ v1 , v2 , ⋯ , vn ] ∈ Rp×n . The attention operation computes the responses of a query vector qi by attending it to all key vectors in K and uses the results to take a weighted sum over value vectors in V . The layer-wise forward-propagation operation of an attention operator can be expressed asO = V softmax ( KTQ ) . Matrix multiplication betweenKT andQ results in a coefficient matrixE =KTQ , in which each element eij is calculated by the inner product between kTi and qj . This coefficient matrix E computes similarity scores between every query vector qi , and every key vector kj and is normalized by a column-wise softmax operator to make every column sum to 1 . The outputO ∈ Rp×m is obtained by multiplying V with the normalized E. In self-attention operators ( Vaswani et al. , 2017 ) , we have Q = K = V . The computational cost in attention operator is O ( m × n × ( d + p ) ) . The memory required for storing the intermediate coefficient matrix E is O ( mn ) . If d = p and m = n , the time and space complexities become O ( m2 × d ) and O ( m2 ) , respectively . There are several other ways to compute E fromQ andK , including Gaussian function , dot product , concatenation , and embedded Gaussian function . It has been shown that dot product is the simplest but most effective one ( Wang et al. , 2018 ) . Therefore , we focus on the dot product similarity function in this work . In practice , we can first perform separate linear transformations on each input matrix , resulting in the following attention operator : O =W V V Softmax ( ( WKK ) TWQQ ) , whereW V ∈ Rp ′×p , WK ∈ Rd ′×d , andWQ ∈ Rd ′×d . For notational simplicity , we omit linear transformations in the following discussion . 2.2 NON-LOCAL OPERATOR Non-local operators ( Wang et al. , 2018 ) apply self-attention operators on higher-order data such as images and videos . Taking 2-D data as an example , the input to the non-local operator is a third-order tensor X ∈ Rh×w×c , where h , w , and c denote the height , width , and number of channels , respectively . The tensor is first converted into a matrixX ( 3 ) ∈ Rc×hw by unfolding along mode-3 ( Kolda & Bader , 2009 ) , as illus- trated in Figure 1 . Then we perform the attention operation by setting Q = K = V = X ( 3 ) . The output of the attention operator is converted back to a third-order tensor as the final output . One practical challenge of the non-local operator is its excessive usage of computational resources . If h = w , the computational cost of a 2-D non-local operator is O ( h4×c ) . The memory used to store the coefficient matrix incurs O ( h4 ) space complexity . The time and space complexities are prohibitively high for high-dimensional and high-order data . 3 KRONECKER ATTENTION NETWORKS . In this section , we describe our proposed Kronecker attention operators , which are efficient and effective attention operators on high-order data . 3.1 KRONECKER ATTENTION OPERATORS . We describe the Kronecker attention operators ( KAO ) in the context of self-attention on 2-D data , but they can be easily generalized to generic attentions . In this case , the input to the ` th layer is a third-order tensor X ( ` ) ∈ Rh×w×c . We propose to use horizontal and lateral average matrices to represent original mode-3 unfolding without much information loss . The horizontal average matrix H and the lateral average matrix L are computed as H = 1 h h ∑ i=1 X ( ` ) i∶∶ ∈ R w×c , L = 1 w w ∑ j=1 X ( ` ) ∶j∶ ∈ R h×c , ( 1 ) where X ( ` ) i∶∶ and X ( ` ) ∶j∶ are the horizontal and lateral slices ( Kolda & Bader , 2009 ) of tensor X ( ` ) , respectively . We then form a matrix C by juxtaposingHT and LT as C = [ HT , LT ] ∈ Rc× ( h+w ) . Based on the horizontal and lateral average matrices contained in C , we propose two Kronecker attention operators ( KAOs ) , i.e. , KAOKV and KAOQKV . In KAOKV as shown in Figure 2 ( b ) , we useX ( ` ) ( 3 ) as the query matrix and C as the key and value matrices as O = attn ( X ( ` ) ( 3 ) , C , C ) ∈ R c×hw . ( 2 ) Note that the number of columns in O depends on the number of query vectors . Thus , we obtain hw output vectors from the attention operation . Similar to the regular attention operator , O is folded back to a third-order tensor Y ( ` ) ∈ Rh×w×c by considering the column vectors inO as mode-3 fibers of Y ( ` ) . KAOKV uses Y ( ` ) as the output of layer ` . If h = w , the time and space complexities of KAOKV are O ( hw × c × ( h + w ) ) = O ( h3 × c ) and O ( hw × ( h + w ) ) = O ( h3 ) , respectively . Compared to the original local operator on 2-D data , KAOKV reduces time and space complexities by a factor of h. In order to reduce the time and space complexities further , we propose another operator known as KAOQKV . In KAOQKV as shown in Figure 2 ( c ) , we use C as the query , key , and value matrices as [ H̃ h̄ , L̃ ® w ] =O = attn ( C , C , C ) ∈ Rc× ( h+w ) . ( 3 ) The final output tensor Y ( ` ) ∈ Rh×w×c is obtained as Y ( ` ) ∶∶i = H̃ T i∶ | L̃Ti∶ , where H̃i∶ and L̃i∶ are the ith rows of the corresponding matrices . That is , the ith frontal slice of Y ( ` ) is obtained by computing the outer sum of the ith rows of H̃ and L̃ . If h = w , the time and space complexities of KAOQKV are O ( ( h +w ) × c × ( h +w ) ) = O ( h2 × c ) and O ( ( h +w ) × ( h +w ) ) = O ( h2 ) , respectively . Thus , the time and space complexities have been reduced by a factor of h2 as compared to the original local operator , and by a factor of h as compared to KAOKV . Note that we do not consider linear transformations in our description , but they can be applied to all three input matrices in KAOKV and KAOQKV as shown in Figure 2 . 3.2 KRONECKER ATTENTION MODULES AND NETWORKS . Attention models have not been used in compact deep models to date , primarily due to their high computational cost . In this section , we design a family of efficient Kronecker attention modules based on KAOs as illustrated in Figure 3 in the appendix . BaseModule : MobileNetV2 ( Sandler et al. , 2018 ) is mainly composed of bottleneck blocks with inverted residuals . Each bottleneck block consists of three convolutional layers ; those are , 1 × 1 convolutional layer , 3 × 3 depth-wise convolutional layer , and another 1 × 1 convolutional layer . Suppose the expansion factor is r and stride is s. Given input X ( ` ) ∈ Rh×w×c for the ` th block , the first 1× 1 convolutional layer outputs rc feature maps X̃ ( ` ) ∈ Rh×w×rc . The depth-wise convolutional layer uses a stride of s and outputs rc feature maps X̄ ( ` ) ∈ Rhs ×ws ×rc . The last 1 × 1 convolutional layer produces d feature maps Y ( ` ) ∈ Rhs ×ws ×d . When s = 1 and c = d , a skip connection is added between X ( ` ) and Y ( ` ) . BaseSkipModule : To facilitate feature reuse and gradient back-propagation in deep models , we improve the BaseModule by adding a skip connection . Given input X ( ` ) , we use an expansion factor of r − 1 for the first 1× 1 convolutional layer , instead of r as in BaseModule . We then concatenate the output with the original input , resulting in X̃ ( ` ) ∈ Rh×w×rc . The other parts of the BaseSkipModule are the same as those of the BaseModule . Compared to the BaseModule , the BaseSkipModule reduces the number of parameters by c × c and computational cost by h ×w × c. It achieves better feature reuse and gradient back-propagation . AttnModule : We propose to add an attention operator into the BaseModule to enable the capture of global features . We reduce the expansion factor of the BaseModule by 1 and add a new parallel path with an attention operator that outputs c feature maps . Concretely , after the depth-wise convolutional layer , the original path outputs X̄ ( ` ) a ∈ R h s × w s × ( r−1 ) c. The attention operator , optionally followed by an average pooling of stride s if s > 1 , produces X̄ ( ` ) b ∈ R h s × w s ×c . Concatenating them gives X̄ ( ` ) ∈ Rhs ×ws ×rc . The final 1 × 1 convolutional layer remains the same . Within the attention operator , we only apply the linear transformation on the value matrix V to limit the number of parameters and required computational resources . In this module , the original path acts as locality-based feature extractors , while the new path with an attention operator computes global features . This enables the module to incorporate both local and global information . Note that we can use any attention operator in this module , including the regular attention operator and our KAOs . AttnSkipModule : We propose to add an additional skip connection in the AttnModule . This skip connection can always be added unless s > 1 . The AttnSkipModule has the same amount of parameters and computational cost as the AttnModule . In the following sections , we perform some theoretical analysis on the proposed methods .

Summary: Akin to previous work introducing tensor regression layers to leverage the tensor/multimodal structure of inputs [Kossaifi et al., 2017], this paper proposes a Kronecker structured attention operator for self attention models acting on tensor data. The proposed operators enforce various constraints on the query key and value matrices of the attention module, where some (or all) of these matrices are obtained as horizontal and lateral averages of the query matrix; in this context, the query matrix is actually the flattening of a 3rd order tensor, whose slices are referred to as "feature maps". A theoretical analysis investigates the properties of imposing a Kronecker structured covariance matrix on the feature maps. The main theorem provides a way to efficiently construct feature maps satisfying this covariance structure. Experiments on image classification and segmentation are given where the proposed approach shows competitive performances while maintaining a number of parameters on par with efficient architectures such as neural nets.