Self-attention has recently been adopted for a wide range of sequence modeling problems. Despite its effectiveness, self-attention suffers quadratic compute and memory requirements with respect to sequence length. Successful approaches to reduce this complexity focused on attention to local sliding windows or a small set of locations independent of content. Our work proposes to learn dynamic sparse attention patterns that avoid allocating computation and memory to attend to content unrelated to the query of interest. This work builds upon two lines of research: it combines the modeling flexibility of prior work on content-based sparse attention with the efficiency gains from approaches based on local, temporal sparse attention. Our model, the Routing Transformer, endows self-attention with a sparse routing module based on online k-means while reducing the overall complexity of attention to O(n 1.5 d) from O(n 2 d) for sequence length n and hidden dimension d. We show that our model outperforms comparable sparse attention models on language modeling on Wikitext-103 (15.8 vs 18.3 perplexity) as well as on image generation on ImageNet-64 (3.43 vs 3.44 bits/dim) while using fewer self-attention layers.
Summary: This paper proposes a new mechanism for computing the attention scores efficiently in Transformers to avoid the quadratic computational cost in conventional Transformers (i.e., when computing the QK^T for self-attention). With a content-based "routing" technique ("group/cluster attention" may be a more accurate term to differentiate with the dynamics of Capsule Networks), the time steps of each layer form different clusters, and the self-attention mechanism is only performed within each such cluster. The authors verify their approach on very large-scale and challenging sequence benchmarks such as WikiText-103 (WT103) and enwik-8.
Pros:
+ The math notations are generally clear (e.g., dimensionality) and the paper is well-organized.
+ The paper addresses a very important (efficiency) question in Transformers, and gave a proper review to prior related works (such as Child et al.).
+ The experiments used to evaluate the model are all large-scale, high-dimensional and challenging tasks. The routing Transformer seems to be able to do very well on the WikiText-103 word-level language modeling task.
+ I especially like the analysis of the local vs. routing attention via the JSD score. It's a very interesting and useful way to compare the attention maps without visualizations (which many other works use). 
--------------------------------
Cons/Questions:
1. Please number your formulas and equations. 
2. Based on your definition of X' (at the end of page 3, $X_i' = \sum_{j < i}^n A_{ij} V_j$) and the subsequent Eq. (2), it seems that you did not add layer normalization and residual connection to the self-attention block. However, almost all prior works that use the Transformer architecture, such as Vaswani et al. [1], Child et al. [2] and Dai et al. [3] has two residual blocks: 1) X' = layernorm(SelfAttention(X) + X); and 2) X'' = layernorm(mlp(X') + X'). Is there any reason that you abandoned this design and used Eq. (2)? Could that have impacted the model performance? In addition, it seems you didn't scale QK^T by 1/\sqrt{d}, as did in prior works?
3. How are the k-means centroid vectors (\mu_1, ..., \mu_k) initialized? Picking the initial centers for k-means clustering has been long known as an interesting/challenging problem, especially in a high-dimensional space (which tasks like language modeling deal with). This challenge spawned works such as K-means++ [4]. I suggest the authors discussing their initialization scheme as well as the initialization's impact on the effectiveness of their approach.
4. In Section 4.2, you said "$FG^\top$ where $F$ and $G$ are binary matrices denoting cluster memberships of queries and keys respectively". However, your previous description seems to suggest that the clusters are formed based on the R matrix (with shape n x d), which contains both Q and K. Does this mean F = G^T? Or could position i in query and position i in key be in different clusters? (I think $ltr(FG^\top * A)$ would only make sense if F=G^T.) This part is confusing in the paper; please clarify.
5. The inequality at the bottom of page 5 (**feel free to clarify if I missed some assumptions from the paper**): given that $\|a-b\|^2 = (a-b)^\top(a-b) = a^\top a + b^\top b - 2a^\top b$, it's unclear in this case ($a = Q_i-K_j$ and $b = Q_j-K_i$) whether $2a^\top b$ is positive or negative (even if $Q_i, K_i$ are unit vectors). Therefore, I'm not sure why it is a >= relation here.
If this is incorrect indeed, then the subsequent analysis the authors made on \epsilon would be a bit problematic, too. (Additionally, the authors seem to suggest that $Q_i^\top K_j = Q_j^\top K_i$ in the subsequent analysis, which should not be the case, unless W_Q = W_K.) Also, just curious, how did you regularize $W_R$ to be orthonormal while training?
6. The authors acknowledge at the end of Section 4 that their method of assigning clusters "does not guarantee that each point belongs to a single cluster". This means that even with balanced cluster sizes, each cluster will have more than n/k routing vectors (i.e., time steps). This makes the $O(nkd + n^2d/k)$ bound the authors mentioned potentially inaccurate. On the other hand, without forcing such cluster balance, in the worst case, the complexity is still O(n^2d), correct?
7. The paper proposes to use $\sqrt{n}$ clustering centers, which makes sense from the perspective of minimizing total computations. However, since the cluster centroids are learned parameters, does this mean the trained routing Transformer cannot easily generalize to different sequence lengths easily? (e.g., train on sequence length 256 but test on sequence length 2000, or simply 15?) Moreover, on longer sequences, won't the sorting/clustering be increasingly inefficient?
8. What is the motivation of using half of the heads for local attention and the other half for routing attention (cf. the intro paragraph of Section 5)? Why not just use all of the heads with routing attention? Without ablations on this, it's hard to tell which part is contributing how much exactly.
9. Could you also show the # of parameters used by your model in Table 1-3 (as did in prior works)?
10. Do you observe better performance with deeper layers (which seems to be the case for the sparse Transformer [2])?
11. Except the computational concern, I still don't quite exactly see the motivation behind clustering attention. For language modeling, for instance, the time steps should be mutually dependent, and it's more like a connected, chain-like structure (where each word depends strongly on its neighbors). However, the cluster attention seems to group the words together, and there is no cross-cluster attention (i.e., the graph is broken into smaller components). In addition, the authors said the centroid parameters (\mu_1, ..., \mu_k) are shared across sequences. So in a certain sense, are they like the "anchors" in the (transformed) word embedding space? Are we then, in language modeling tasks, performing attentions among words only around these "anchors"?
12. While the theoretical complexity may be lower, how does the wall-clock time of an L-layer routing Transformer compare to an L-layer conventional Transformer (on the same sequence length, batch size)? In particular, although the routing mechanism avoids computing A, having to deal with each of the $\sqrt{n}$ clusters means that one will need to process the cluster-based attentions sequentially (and, as mentioned, you have 8 heads for local attention and 8 for routing attention, which I think are also processed in two steps?). Would this make the model actually slower on GPUs?
=====================
Minor issues that didn't impact the score:
13. In the equation at the bottom of page 3 (the definition of X_i'), shouldn't it be (j <= i) instead of (j < i), in the case of autoregressive sequence modeling? E.g., in language modeling, you do use the current word when predicting the next word.
14. At the end of Section 5.1, you said: "while having fewer [...] attention heads". I believe the released Transformer-XL model also used 16 heads? (https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_large.sh) 
15. How does the number of cluster centers influence the performance of the model? I think this would be an interesting ablation study to make.
=====================
While I'm impressed with the result the routing Transformer achieved on WikiText-103, I am not fully convinced the effectiveness of the approach (e.g., on the other tasks, the routing Transformer seems to be slightly worse than the SOTA transformer; this is less surprising, as the authors didn't change the underlying design of the Transformers). I think the paper can be improved by including more elements, such as ablative experiments and runtime benchmarking. Also, there seem to be some problems with the derivations that the authors made in the current version of the paper. 
[1] https://arxiv.org/abs/1706.03762
[2] https://arxiv.org/abs/1904.10509
[3] https://arxiv.org/abs/1901.02860
[4] http://ilpubs.stanford.edu:8090/778/1/2006-13.pdf