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.
=== 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.