Title: No title found

Abstract: Attention operators have been applied on both 1-D data like texts and higher-order data such as images and videos. Use of attention operators on high-order data requires flattening of the spatial or spatial-temporal dimensions into a vector, which is assumed to follow a multivariate normal distribution. This not only incurs excessive requirements on computational resources, but also fails to preserve structures in data. In this work, we propose to avoid flattening by developing Kronecker attention operators (KAOs) that operate on high-order tensor data directly. KAOs lead to dramatic reductions in computational resources. Moreover, we analyze KAOs theoretically from a probabilistic perspective and point out that KAOs assume the data follow matrix-variate normal distributions. Experimental results show that KAOs reduce the amount of required computational resources by a factor of hundreds, with larger factors for higher-dimensional and higher-order data. Results also show that networks with KAOs outperform models without attention, while achieving competitive performance as those with original attention operators.

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 (Veličković 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.1. ATTENTION OPERATOR:

The inputs to an attention operator include a query matrix Q = [q 1 , q 2 , ⋯, q m ] ∈ R d×m , a key matrix K = [k 1 , k 2 , ⋯, k n ] ∈ R d×n , and a value matrix V = [v 1 , v 2 , ⋯, v n ] ∈ R p×n . The attention operation computes the responses of a query vector q i 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 as O = V softmax(K T Q). Matrix multiplication between K T and Q results in a coefficient matrix E = K T Q, in which each element e ij is calculated by the inner product between k T i and q j . This coefficient matrix E computes similarity scores between every query vector q i , and every key vector k j and is normalized by a column-wise softmax operator to make every column sum to 1. The output O ∈ R p×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(m 2 × d) and O(m 2 ), respectively. There are several other ways to compute E from Q and K, 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((W K K) T W Q Q), where W V ∈ R p ′ ×p , W K ∈ R d ′ ×d , and W Q ∈ R d ′ ×d .

For notational simplicity, we omit linear transformations in the following discussion. 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 ∈ R h×w×c , where h, w, and c denote the height, width, and number of channels, respectively. The tensor is first converted into a matrix X (3) ∈ R c×hw by unfolding along mode-3 (Kolda & Bader, 2009) , as illustrated in Figure 1 . Then we perform the attention operation by setting

2.2. NON-LOCAL OPERATOR:

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(h 4 × c). The memory used to store the coefficient matrix incurs O(h 4 ) 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 ( ) ∈ R h×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

EQUATION

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 juxtaposing H T and L T as C = [H T , L T ] ∈ R c× (h+w) . In the regular attention operator (a), the input tensor is unfolded into a mode-3 matrix and fed into the attention operator. The output of the attention operator is folded back to a tensor as the final output. In KAO KV (b), we juxtapose the horizontal and lateral average matrices derived from the input tensor as the key and value matrices. We keep the mode-3 unfolding of input tensor as the query matrix. In KAO QKV (c), all three input matrices use the juxtaposition of two average matrices. In contrast to KAO KV , we use an outer-sum operation to generate the output third-order tensor.

Q K V Q K V Q K V c c c c c c c c c h h h w w w h h h w w w c c c c c c c c c c c w w h h w h  w h  w h  w h  w h  w h  hw hw hw hw hw hw

Based on the horizontal and lateral average matrices contained in C, we propose two Kronecker attention operators (KAOs), i.e., KAO KV and KAO QKV . In KAO KV as shown in Figure 2 (b), we use X ( )

(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 ( ) ∈ R h×w×c by considering the column vectors in O as mode-3 fibers of Y ( ) . KAO KV uses Y ( ) as the output of layer . If h = w, the time and space complexities of KAO KV are O(hw

× c × (h + w)) = O(h 3 × c) and O(hw × (h + w)) = O(h 3 ), respectively.

Compared to the original local operator on 2-D data, KAO KV 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 KAO QKV . In KAO QKV as shown in Figure 2 (c), we use C as the query, key, and value matrices as h+w) .

[ H h , L w ] = O = attn(C, C, C) ∈ R c×(

(3)

The final output tensor

Y ( ) ∈ R h×w×c is obtained as Y ( ) ∶∶i = HT i∶ LT i∶ ,

where Hi∶ and Li∶ 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

KAO QKV are O((h + w) × c × (h + w)) = O(h 2 × c) and O((h + w) × (h + w)) = O(h 2 )

, respectively. Thus, the time and space complexities have been reduced by a factor of h 2 as compared to the original local operator, and by a factor of h as compared to KAO KV . Note that we do not consider linear transformations in our description, but they can be applied to all three input matrices in KAO KV and KAO QKV 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 ( ) ∈ R h×w×c for the th block, the first 1 × 1 convolutional layer outputs rc feature maps X( ) ∈ R h×w×rc . The depth-wise convolutional layer uses a stride of s and outputs rc feature maps

X( ) ∈ R h s × w s ×rc . The last 1 × 1 convolutional layer produces d feature maps Y ( ) ∈ R h s × w s ×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( ) ∈ R h×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( ) ∈ R h s × w s ×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.

3.3. FROM MULTIVARIATE TO MATRIX-VARIATE DISTRIBUTIONS:

We analyze our solutions for attention operators on high-order data from a probabilistic perspective. We take the non-local operator on 2-D data as an example. Formally, consider a self-attention operator with

Q = K = V = X (3)

, where X (3) ∈ R c×hw is the mode-3 unfolding of a third-order input tensor X ∈ R h×w×c , as illustrated in Figure 1 . The ith row of X (3) corresponds to vec(X ∶∶i ) T ∈ R hw , where X ∶∶i ∈ R h×w denotes the ith frontal slice of X (Kolda & Bader, 2009) , and vec(⋅) denotes the vectorization of a matrix by concatenating its columns (Gupta & Nagar, 2018) .

The frontal slices X ∶∶1 , X ∶∶2 , . . . , X ∶∶c ∈ R h×w of X are usually known as c feature maps. In this view, the mode-3 unfolding is equivalent to the vectorization of each feature map independently. It is worth noting that, in addition to vec(⋅), any other operation that transforms each feature map into a vector leads to the same output from the non-local operator, as long as a corresponding reverse operation is performed to fold the output into a tensor. This fact indicates that unfolding of X in local operators ignores the structural information within each feature map, i.e., the relationships among rows and columns. In addition, such unfolding results in excessive requirements on computational resources, as explained in Section 2.2.

In the following discussions, we focus on one feature map X ∈ {X ∶∶1 , X ∶∶2 , . . . , X ∶∶c } by assuming feature maps are conditionally independent of each other, given feature maps of previous layers. This assumption is shared by many deep learning techniques that process each feature map independently, including the unfolding mentioned above, batch normalization, instance normalization (Ulyanov et al., 2016) , and pooling operations (LeCun et al., 1998) . To view the problem above from a probabilistic perspective (Ioffe & Szegedy, 2015; Ulyanov et al., 2016) , the unfolding yields the assumption that vec(X) follows a multivariate normal distribution as vec(X) ∼ N hw (µ, Ω), where µ ∈ R hw and Ω ∈ R hw×hw . Apparently, the multivariate normal distribution does not explicitly model relationships among rows and columns in X. To address this limitation, we propose to model X using a matrix-variate normal distribution (Gupta & Nagar, 2018) , defined as below.

Definition 1. A random matrix A ∈ R m×n is said to follow a matrix-variate normal distribution MN m×n (M , Ω ⊗ Ψ) with mean matrix M ∈ R m×n and covariance matrix Ω ⊗ Ψ, where

Ω ∈ R m×m ≻ 0 and Ψ ∈ R n×n ≻ 0, if vec(A T ) ∼ N mn (vec(M T ), Ω ⊗ Ψ).

Here, ⊗ denotes the Kronecker product (Van Loan, 2000; Graham, 2018) .

The matrix-variate normal distribution has separate covariance matrices for rows and columns. They interact through the Kronecker product to produce the covariance matrix for the original distribution. Specifically, for two elements X ij and X i ′ j ′ from different rows and columns in X, the relationship between X ij and X i ′ j ′ is modeled by the interactions between the ith and i ′ th rows and the jth and j ′ th columns. Therefore, the matrix-variate normal distribution can incorporate relationships among rows and columns.

3.4. THE PROPOSED MEAN AND COVARIANCE STRUCTURES:

In machine learning, (Kalaitzis et al., 2013) proposed to use the Kronecker sum to form covariance matrices, instead of the Kronecker product. Based on the above observations and studies, we propose to model X as X ∼ MN h×w (M , Ω ⊕ Ψ), where M ∈ R h×w , Ω ∈ R h×h ≻ 0, Ψ ∈ R w×w ≻ 0, ⊕ denotes the Kronecker sum (Kalaitzis et al., 2013) , defined as and I [n] denotes an n × n identity matrix. Covariance matrices following the Kronecker sum structure can still capture the relationships among rows and columns (Kalaitzis et al., 2013) . It also follows from (Allen & Tibshirani, 2010; Wang et al., 2017) that constraining the mean matrix M allows a more direct modeling of the structural information within a feature map. Following these studies, we assume X follows a variant of the matrix-variate normal distribution as

Ω ⊕ Ψ = Ω ⊗ I [w] + I [h] ⊗ Ψ,

EQUATION

where the mean matrix M ∈ R h×w is restricted to be the outer sum of two vectors, defined as  and 1 [n] denotes an all-one vector of size n. Under this model, the marginal distributions of rows and columns are both multivariate normal (Allen & Tibshirani, 2010) . Specifically, the ith row vector X i∶ ∈ R 1×w follows X T i∶ ∼ N w (µ i +υ T , Ω ii +Ψ), and the jth column vector X ∶j ∈ R h×1 follows X ∶j ∼ N h (υ i +µ, Ψ ii +Ω). In the following discussion, we assume that Ω and Ψ are diagonal, implying that any pair of variables in X are uncorrelated. Note that, although the variables in X are independent, their covariance matrix still follows the Kronecker covariance structure, thus capturing the relationships among rows and columns (Allen & Tibshirani, 2010; Wang et al., 2017) .

M = µ υ = µ1 T [w] + 1 [h] υ T , where µ ∈ R h , υ ∈ R w ,

3.5. MAIN TECHNICAL RESULTS:

Let X row = (∑ h i=1 X T i∶ ) h ∈ R w and X col = (∑ w j=1 X ∶j ) w ∈ R h be the average of row and column vectors, respectively. Under the assumption above, X row and X col follow multivariate normal distributions as

EQUATION

where µ = (∑ h i=1 µ i ) h, Ω = (∑ h i=1 Ω ii ) h, υ = (∑ w j=1 υ j ) w, and Ψ = (∑ w j=1 Ψ jj ) w. Our main technical results can be summarized in the following theorem.

Table 1 : Comparisons between the regular attention operator, the regular attention operator with a pooling operation (Wang et al., 2018) , and our proposed KAO KV and KAO QKV in terms of the number of parameters, MAdd, memory usage, and CPU inference time on data of different sizes. The input sizes are given in the format of "batch size × spatial sizes × number of channels". "Attn" denotes the regular attention operator. "Attn+Pool" denotes the regular attention operator which employs a 2 × 2 pooling operation on K and V input matrices to reduce required computational resources. 5) and ( 6) with diagonal Ω and Ψ, if (a) r 1 , r 2 , . . . , r h are independent and identically distributed (i.i.d.) random vectors that follow the distribution in Eq. ( 5), (b) c 1 , c 2 , . . . , c w are i.i.d. random vectors that follow the distribution in Eq. ( 6), (c) r 1 , r 2 , . . . , r h and c 1 , c 2 , . . . , c w are independent, we have

Input:

X ∼ MN h×w M , Ψ+Ω w ⊕ Ω+Ψ h , where X = [r 1 , r 2 , . . . , r h ] T + [c 1 , c 2 , . . . , c w ], M = (µ υ) + (µ + υ).

In particular, if h = w, the covariance matrix satisfies tr Ψ+Ω w ⊕ Ω+Ψ h = 2 h tr (Ω ⊕ Ψ), where tr(⋅) denotes matrix trace. The proof of Theorem 1 can be found in the appendix. With certain normalization on X, we can have µ + υ = 0, resulting in M = µ υ. As the trace of a covariance matrix measures the total variation, Theorem 1 implies that X follows a matrix-variate normal distribution with the same mean and scaled covariance as the distribution of X in Eq. ( 4). We build KAOs based on this conclusion and the process to obtain X from X.

4. EXPERIMENTAL STUDIES:

In this section, we evaluate our methods and networks on image classification and segmentation tasks.

4.1. COMPARISON OF COMPUTATIONAL EFFICIENCY:

According to the theoretical analysis in Section 3.1, our KAOs have efficiency advantages over regular attention operators on high-order data, especially for inputs with large spatial sizes. We conduct simulated experiments to evaluate the theoretical results. To reduce the influence of external factors, we build networks composed of a single attention operator, and apply the TensorFlow profile tool (Abadi et al., 2016) to report the multiply-adds (MAdd), required memory, and time consumed on 2-D simulated data. For the simulated input data, we set the batch size and number of channels both to 8, and test three spatial sizes; those are, 56 × 56, 28 × 28, and 14 × 14. The number of output channels is also set to 8. Table 1 summarizes the comparison results. On simulated data of spatial sizes 56 × 56, our KAO KV and KAO QKV achieve 31.1 and 305.8 times speedup, and 96.18% and 99.73% memory saving compared to the regular attention operator, respectively. Our proposed KAOs show significant improvements over regular attention operators in terms of computational resources, which is consistent with the theoretical analysis. In particular, the amount of improvement increases as the spatial sizes increase. These results show that the proposed KAOs are efficient attention operators on high-dimensional and high-order data.

4.2. RESULTS ON IMAGE CLASSIFICATION:

With the high efficiency of our KAOs, we have proposed several efficient Kronecker attention modules for compact CNNs in Section 3.2. To further show the effectiveness of KAOs and the modules, we build novel compact CNNs known as Kronecker attention networks (KANets). Following the practices in (Wang et al., 2018) , we apply these modules on inputs of spatial sizes 28×28, 14×14, and 7 × 7. The detailed network architecture is described in Table 6 in the appendix due to space constraint. We compare KANets with other CNNs on the ImageNet ILSVRC 2012 image classification dataset, which serves as the benchmark for compact CNNs (Howard et al., 2017; Zhang et al., 2017; Gao et al., 2018; Sandler et al., 2018) . Details of the experimental setups are provided in the appendix. The comparison results between our KANets and other CNNs in terms of the top-1 accuracy, number of parameters, and MAdd are reported in Table 2 . SqueezeNet (Iandola et al., 2016) has the least number of parameters, but uses the most MAdd and does not obtain competitive performance as compared to other compact CNNs. Among compact CNNs, MobileNetV2 (Sandler et al., 2018) is the previous state-of-theart model, which achieves the best trade-off between effectiveness and efficiency. According to the results, our KANets significantly outperform MobileNetV2 with 0.03 million fewer parameters. Specifically, our KANet KV and KANet QKV outperform MobileNetV2 by margins of 0.9% and 0.8%, respectively. More importantly, our KANets has the least computational cost. These results demonstrate the effectiveness and efficiency of our proposed KAOs.

The performance of KANets indicates that our proposed methods are promising, since we only make small modifications to the architecture of MobileNetV2 to include KAOs. Compared to modules with the regular convolutional layers only, our proposed modules with KAOs achieve better performance without using excessive computational resources. Thus, our methods can be used widely for designing compact deep models. Next, we show that our proposed KAOs are as effective as regular attention operators. We perform experiments to compare our proposed KAOs with regular attention operators. We consider the regular attention operator and the one with a pooling operation in (Wang et al., 2018) . For the attention operator with pooling operation, the spatial sizes of the key matrix K and value matrix V are reduced by 2 × 2 pooling operations to save computation cost. To compare these operators in fair settings, we replace all KAOs in KANets with regular attention operators and regular attention operators with a pooling operation, denoted as AttnNet and At-tnNet+Pool, respectively.

4.3. COMPARISON WITH REGULAR ATTENTION OPERATORS:

The comparison results are summarized in Table 3 . Note that all these models have the same number of parameters. We can see that KANet KV and KANet QKV achieve similar performance as AttnNet and AttnNet+Pool with dramatic reductions of computational cost. The results indicate that our proposed KAOs are as effective as regular attention operators while being much more efficient. In addition, our KAOs are better than regular attention operators that uses a pooling operation to increase efficiency in (Wang et al., 2018) .

Under review as a conference paper at ICLR 2020 4.4 ABLATION STUDIES In order to show the efficiency and effectiveness of our KAOs in broader application scenarios, we perform additional experiments on image segmentation tasks using the PASCAL 2012 dataset (Everingham et al., 2010) . With the extra annotations provided by (Hariharan et al., 2011) , the augmented dataset contains 10,582 training, 1,449 validation, and 1,456 testing images. Each pixel of the images is labeled by one of 21 classes with 20 foreground classes and 1 background class. We re-implement the DeepLabV2 model (Chen et al., 2018) as our baseline. Following (Wang & Ji, 2018) , using attention operators as the output layer, instead of atrous spatial pyramid pooling (ASPP), results in a significant performance improvement. In our experiments, we replace ASPP with the regular attention operator and our proposed KAOs, respectively, and compare the results. For all attention operators, linear transformations are applied on Q, K, and V . Details of the experimental setups are provided in the appendix. 

5. CONCLUSIONS:

In this work, we propose Kronecker attention operators to address the practical challenge of applying attention operators on high-order data. We investigate the problem from a probabilistic perspective and use matrix-variate normal distributions with Kronecker covariance structure. Experimental results show that our KAOs reduce the amount of required computational resources by a factor of hundreds, with larger factors for higher-dimensional and higher-order data. We employ KAOs to design a family of efficient modules, leading to our KANets. KANets significantly outperform the previous state-of-the-art compact models on image classification tasks, with fewer parameters and less computational cost. transformation is only performed on the value matrix V to limit the number of parameters and computational resources.

3. ILLUSTRATION OF KRONECKER ATTENTION MODULES:

Conv 1 x 1 Stride = 1 BN + ReLU6 Conv 1 x 1 Stride = 1 Output DWConv 3 x 3 Stride = BN + ReLU6 + (a) Conv 1 x 1 Stride = 1 BN + ReLU6 c DWConv 3 x 3 Stride = BN + ReLU6 Conv 1 x 1 Stride = 1 Output + (b) Conv 1 x 1 Stride = 1 BN + ReLU6 DWConv 3 x 3 Stride = BN + ReLU6 c KAO Conv 1 x 1 Stride = 1 Output + (c) Conv 1 x 1 Stride = 1 BN + ReLU6 DWConv 3 x 3 Stride = BN + ReLU6 c Conv 1 x 1 Stride = 1 Output KAO + + (d) d s w s h   d s w s h   d s w s h   d s w s h   d w h   d w h   d w h   d w h  

4. THE KANETS ARCHITECTURE:



5. PROOF OF THEOREM 1:

The fact that Ω and Ψ are diagonal implies independence in the case of multivariate normal distributions. Therefore, it follows from assumptions (a) and (b) that

[r 1 , r 2 , . . . , r h ] T ∼ MN h×w M r ,

EQUATION

where M r = µ + [υ, υ, . . . , υ] T = µ + 1 [h] υ T , and

EQUATION

where M c = υ + [µ, µ, . . . , µ] = υ + µ1 T [w] . Given assumption (c) and X = [r 1 , r 2 , . . . , r h ] T + [c 1 , c 2 , . . . , c w ], we have

EQUATION

where M = M r + M c = (µ υ) + (µ + υ).

If h = w, we have

EQUATION

and

EQUATION

This completes the proof of the theorem.

REFERENCES:

All citations refer to the references in the main paper.

:

BACKGROUND AND RELATED WORKIn this section, we describe the attention and related non-local operators.

Figures and Tables:

Figure 1: Conversion of a third-order tensor into a matrix by unfolding along mode-3. In this example, a h × w × c tensor is unfolded into a c × hw matrix.

Figure 2: Illustrations of regular attention operator (a), KAO KV (b) and KAO QKV (c) on 2-D data.In the regular attention operator (a), the input tensor is unfolded into a mode-3 matrix and fed into the attention operator. The output of the attention operator is folded back to a tensor as the final output. In KAO KV (b), we juxtapose the horizontal and lateral average matrices derived from the input tensor as the key and value matrices. We keep the mode-3 unfolding of input tensor as the query matrix. In KAO QKV (c), all three input matrices use the juxtaposition of two average matrices. In contrast to KAO KV , we use an outer-sum operation to generate the output third-order tensor.

Figure 3: Architectures of the BaseModule (a), BaseSkipModule (b), AttnModule (c), and AttnSkip-Module (d) as described in Section 3.2. The skip connections indicated by single dashed paths are not used when s > 1 or c ≠ d. Those indicated by double dashed paths are not used when s > 1.

Comparisons between KANets and otherCNNs in terms of the top-1 accuracy on the Im-ageNet validation set, the number of total parameters, and MAdd. We use KANet KV and KANet QKV to denote KANets using KAO KV and KAO QKV , respectively.

Table content: <table><tr><td>Model</td><td colspan="3">Top-1 Params MAdd</td></tr><tr><td>VGG16</td><td>0.715</td><td colspan="2">128m 15300m</td></tr><tr><td>AlexNet</td><td>0.572</td><td>60m</td><td>720m</td></tr><tr><td>SqueezeNet</td><td>0.575</td><td>1.3m</td><td>833m</td></tr><tr><td>MobileNetV1</td><td>0.706</td><td>4.2m</td><td>569m</td></tr><tr><td>ShuffleNet 1.5x</td><td>0.715</td><td>3.4m</td><td>292m</td></tr><tr><td>ChannelNet-v1</td><td>0.705</td><td>3.7m</td><td>407m</td></tr><tr><td>MobileNetV2</td><td colspan="3">0.720 3.47m 300m</td></tr><tr><td>KANet KV (ours)</td><td colspan="3">0.729 3.44m 288m</td></tr><tr><td colspan="4">KANet QKV (ours) 0.728 3.44m 281m</td></tr></table>

Comparisons between KANets with regular attention operators (denoted as AttnNet), KANets with regular attention operators with a pooling operation (denoted as AttnNet+Pool) and KANets with KAOs in terms of the top-1 accuracy on the ImageNet validation set, the number of total parameters, and MAdd.

Table content: <table><tr><td>Model</td><td>Top-1</td><td>Params</td><td>MAdd</td></tr><tr><td>AttnNet</td><td>0.730</td><td>3.44m</td><td>365m</td></tr><tr><td colspan="2">AttnNet+Pool 0.729</td><td>3.44m</td><td>300m</td></tr><tr><td>KANet KV</td><td>0.729</td><td>3.44m</td><td>288m</td></tr><tr><td>KANet QKV</td><td>0.728</td><td>3.44m</td><td>281m</td></tr></table>

Comparisons among KANet KV , MobileNetV2, MobileNetV2 with KAOs KV (MobileNetV2+KAO KV ), and KANet without KAO (KANet w/o KAO) in terms of the top-1 accuracy on the ImageNet validation set, the number of total parameters, and MAdd. On the other hand, KANet KV outperforms KANet w/o KAO by a margin of 0.8%, while KANet w/o KAO has more parameters than KANet KV . KANet KV achieves the best performance while costing the least computational resources. The results indicate that our proposed KAOs are effective and efficient, which is independent of specific network architectures.

Table content: <table><tr><td>Model</td><td>Top-1 Params MAdd</td></tr><tr><td>MobileNetV2</td><td>0.720 3.47m 300m</td></tr><tr><td colspan="2">MobileNetV2+KAO 0.726 3.46m 298m</td></tr><tr><td>KANet KV</td><td>0.729 3.44m 288m</td></tr><tr><td>KANet w/o KAO</td><td>0.721 3.46m 298m</td></tr></table>

shows the evaluation results in terms of pixel accuracy and mean intersection over union (IoU) on the PASCAL VOC 2012 validation set. Clearly, models with attention operators outperform the baseline model with ASPP. Compared with the regular attention operator, KAOs result in similar pixel-wise accuracy but slightly lower mean IoU. From the pixel-wise accuracy, results indicate that KAOs are as effective as the regular attention operator. The decrease in mean IoU may be caused by the strong structural assumption behind KAOs. Overall, the experimental results demonstrate the efficiency and effectiveness of our KAOs in broader application scenarios.

Table content: <table/>

describes the detailed KANets architecture. We use KAOs in every AttnSkipModule. In KAOs, we use multi-head attention with 4 heads and concatenate results for output. The linear

Table content: <table/>

Details of the KANets architecture. Each line describes a sequence of operators in the format of "input size / operator name / expansion rate r / number of output channels c / number of operators in the sequence n / stride s". "Conv2D" denotes the regular 2D convolutional layer. "AvgPool" and "FC" denote the global average pooling layer and the fully-connected layer, respectively. All depth-wise convolutions use the kernel size of 3 × 3. For multiple operators in a sequence denoted in the same line, all operators produce c output channels. And the first operator applies the stride of s while the following operators applies the stride of 1. k denotes the class number in the task.

Table content: <table><tr><td>Input</td><td>Operator</td><td>r</td><td>c</td><td>n</td><td>s</td></tr><tr><td>224 2 ×3</td><td>Conv2D 3 × 3</td><td>-</td><td>32</td><td>1</td><td>2</td></tr><tr><td colspan="2">112 2 ×32 BaseSkipModule</td><td>1</td><td>16</td><td>1</td><td>1</td></tr><tr><td colspan="2">112 2 ×16 BaseSkipModule</td><td>6</td><td>24</td><td>2</td><td>2</td></tr><tr><td>56 2 ×24</td><td>BaseSkipModule</td><td>6</td><td>32</td><td>2</td><td>2</td></tr><tr><td>28 2 ×32</td><td>AttnSkipModule</td><td>6</td><td>32</td><td>1</td><td>1</td></tr><tr><td>28 2 ×32</td><td>BaseSkipModule</td><td>6</td><td>64</td><td>1</td><td>2</td></tr><tr><td>14 2 ×64</td><td/><td>6</td><td>64</td><td>3</td><td>1</td></tr><tr><td>14 2 ×64</td><td>AttnSkipModule</td><td>6</td><td>96</td><td>3</td><td>1</td></tr><tr><td>14 2 ×96</td><td>BaseSkipModule</td><td>6</td><td>160</td><td>1</td><td>2</td></tr><tr><td>7 2 ×160</td><td>AttnSkipModule</td><td>6</td><td>160</td><td>2</td><td>1</td></tr><tr><td>7 2 ×160</td><td>AttnSkipModule</td><td>6</td><td>320</td><td>1</td><td>1</td></tr><tr><td>7 2 ×320</td><td>Conv2D 1 × 1</td><td>-</td><td>1280</td><td>1</td><td>1</td></tr><tr><td colspan="2">7 2 ×1280 AvgPool + FC</td><td>-</td><td>k</td><td>1</td><td>-</td></tr></table>

Peer review (iclr2020):
=== Summary === The authors propose the Kronecker Attention Operator (KAO), a novel efficient attention mechanism for high order data such as images (height H and width W dimensions). Regular self-attention on images produce HW queries attend to HW keys and their values, resulting in quadratic time/space complexity O((HW)^2). In contrast, KAO applies self-attention after averaging keys, values (and potentially queries) along width or height dimension, resulting in (H+W) keys and values and thus achieving linear time/space complexity in O(HW). The authors analyze their approach from a probabilistic perspective by assuming that the rows and columns follow matrix-variate normal distributions with Kronecker covariance. Experimental results on ImageNet classification and PASCAL segmentation show that adding Kronecker Attention Operators to baseline architectures yields improvements with a negligible increase in parameters/computations. === Recommendation === Attention has been shown to provide improvements over convolutions but can indeed incur significant memory costs. Hence, discovering efficient attention mechanisms in an important research direction with practical impact. The proposed method is sound and simple and may achieve significant speedups over regular self-attention. However, there are many experimental presentation/evaluation issues in the draft: - An important assumption of the method is that Kronecker factorization preserves structure well while reducing computations. Previous works reduce memory costs by pooling the input to the attention operator (and resizing accordingly the output). A natural baseline would be to pool in such a way that the number of queries is the same as when applying Kronecker factorization. Can the authors report results on such a baseline? - Section 4.1 and Table 1 present computational efficiency (memory and latency) in simulated cases. Can the authors additionally report results for the gains obtained on the entire architectures? The rest of the architecture isn*t an external factor to ignore. - Can the authors also report memory consumption and latency time of KANet vs MobileNetv2? - The authors ignore a vast literature on (self-)attention mechanisms used in vision, some of which are relatively cheap. These methods should also be cited as related work and potentially use as baselines (Squeeze-and-Excitation for example). The probabilistic analysis does not convincingly motivate the Kronecker Attention Operator. Notably, the assumptions for Theorem 1 are also quite unrealistic for images (nearby pixels are correlated via locality, so rows and columns aren*t independent). Can the authors clarify the assumptions and explain why they*re reasonable? In conclusion, the authors only compare their approach against regular self-attention and do not refer to relevant work on attention mechanisms for vision. They do not provide memory/consumption gains on actual architectures, which makes it hard to evaluate the significance of the work. The theoretical motivation for the method is not convincing. I argue for rejection. === Additional questions/comments === - Section 2.1: Self-attention is better described as Q = W^Q X, K = W^K X, V = W^V X, O = V softmax(K^TQ). As opposed to Q=K=V, O = W^V V softmax((W^KK)^TW^QQ) - Since the base architecture KANet is MobileNetv2, there is no need to include VGG, AlexNet, SqueezeNet, in Table 2.