Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: In this paper, the authors propose to adaptively learn the sparsity of attention module through modeling the mask matrix as a bipartite graph generated by SBM.

Specifically, they discuss how they can apply the new attention mechanism to both forward and backward propagation.
- They use Eqn. (5), (6), (7), and (8) to model the mask matrix in forward propagation, which introduce a trainable cluster embedding matrix $C$ and two MLPs to decide the node membership of tokens in $Q, K$.
- They "incorporate a Straight-Through Estimator (STE) to pass the gradient beyond the discrete sampling", which works as if $M \odot \mathbb E M$ is used in forward propagation.
- They also show the new attention mechanism is a universal approximator, and evaluate their proposed methods on Long-range-arena (LRA) benchmarks.
strengths_and_weaknesses: ### Originality

The idea to model the sparsity of attention by SBM is new, though formally similar to the classical idea in pruning that directly model the whole mask matrix.

### Quality

Pros: 

- The author propose a new data-adaptive sparse attention mechanism.


Cons:
- The experiments on LRA is not that convincing nowadays given the recent paper "Efficiently Modeling Long Sequences with Structured State Spaces" (ICLR 2022), which shows current Transformers are way worse than their state-space-based method. The effectiveness of the experimental results would be in doubt whether in another setting/benchmarks which Transformers are more good at, the conclusion would still hold.
- The actual runtime of the new method is not reported and compared to exisitng methods.


### Clarity

The clarity of this paper is awesome. There is a corner concern: to my knowledge, "SBM" should be short for "stochastic block model" rather than stochastic blockmodel. 

Also, the statement of Theorem 1 is not very informative. It would be better if the content in Appendix A that Eqn. (10) can hold with $O(n)$ connections can be moved to the main paper.
 
### Significance

The problem studied is important: how to sparsify and therefore accelerate attention. However, the GPU-unfriendly sparse matrix multiplication makes the significance limited. 
questions: Could the authors show the extra time and space cost of adding the SBM mechanism? How is that compared to another efficient attention mechanism?
limitations: N/A
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper proposes a Transformer variant where the self-attention matrices are masked using random adjacency matrices generated using a stochastic block model (SBM). The goal here is to reduce the computational cost of producing each N x N self-attention matrix. In the forward pass, the proposed method computes SBM parameters from the query and key embeddings using an MLP and samples a self-attention mask using a fast random graph sampling method. In the backward pass, gradients are passed through this discrete sampling step using a straight-through estimator. The authors show that the proposed SBM-Transformer is a universal approximator in the setting considered by Yun et al. (2020) and Zaheer et al. (2020). Empirically, the authors demonstrate that the SBM-Transformer (1) is able to emulate a standard Transformer with dense attention on a synthetic task and (2) outperforms several efficient Transformer variants on average on the Long Range Arena benchmark.
strengths_and_weaknesses: ### Originality

- As noted in the paper, the use of sparse attention masks to reduce the computational cost of the self-attention operation has previously been explored in several existing papers (e.g., axial attention (Ho et al. 2019), BigBird (Zaheer et al., 2020)). The main methodological novelty here is the use of a stochastic block model to generate a random attention mask in each forward pass. To my knowledge, this approach has not been studied in prior work.

### Quality

- **Theoretical results.** The paper establishes a universal approximation property for the proposed transformer variant (Theorem 1). This result is helpful as a sanity check that the proposed method is a reasonable approach. The overall technical contribution here is relatively minor since the proof of the result is primarily an application of existing machinery from Yun et al. (2020) and Zaheer et al. (2020). A minor weakness of the analysis is that it makes use of a modified version of the SBM-Transformer architecture relative to that used in the empirical evaluation -- in particular, the analysis requires that M_ii = 1, unlike in the version used to obtain the results in Table 1.
- **Empirical evaluation.** A strength of the paper is the result that the proposed SBM-Transformer architecture outperforms the baseline efficient Transformer architectures on average over the LRA benchmark tasks (Table 1). However, a significant weakness of this comparison is that it does not appropriately account for the accuracy-cost tradeoff corresponding to each of the evaluated methods. For instance, it is not obvious from the paper whether the SBM-Transformer outperforms the given baselines when controlling for a measure of computational cost such as the per-example FLOP count during inference. Such controls are important since the SBM-Transformer introduces additional operators (such as a 2-layer MLP used for computing the SBM parameters), the costs of which will have to be offset by the increased sparsity of the self-attention matrix.

### Clarity

- The presentation was generally clear and easy to follow.
- The specification of the masked attention computation in Section 3.2 is problematic since simply masking the attention logits to 0 is not equivalent to zeroing-out the attention scores after the softmax operator. 

### Significance

- As noted in Section 6, the unstructured sparsity of the resulting attention maps in the SBM-Transformer formulation is not amenable to fast matrix multiplication on current GPU hardware. Thus, as a practical matter, the proposed method requires longer wall clock times for inference, which consequently limits the significance of the work for practitioners. The authors argue that the method (1) "enables multi-head attention with long sequences", and that (2) practitioners can tune the computational cost of inference by imposing constraints on the number of edges to sample. I don't find these arguments to be convincing as is -- in particular, claim (1) requires additional justification in light of (a) existing efficient transformer architectures that incur subquadratic computation, and (b) techniques such as memory-efficient attention [1] that require subquadratic memory. As for claim (2), the merit of this additional flexibility (e.g., in terms of accuracy vs. cost) should be evaluated empirically.

[1] Rabe & Staats, 2021. Self-attention Does Not Need O(n^2) Memory. 
questions: - As discussed in the previous section, does the SBM-Transformer push out the accuracy-cost Pareto frontier vs. existing efficient Transformer architectures?
- Nondeterministic prediction at inference time is typically considered to be undesirable by practitioners. However, the setup as described requires random sampling under the SBM even during inference. Can this sampling step by replaced with a deterministic operation?

limitations: The authors highlight the limitation of their work with regard to hardware support for unstructured sparsity in Section 6.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The work under revision proposes SBM-Transformer. A novel variant of a transformer architecture adjusts a sparsity of its attention blocks dynamically, based on input sequence data. It is done so forward and backward evaluation cost is linear in the number of edges. Those are adaptively chosen by deploying a generative model for community detection, stochastic block model (SBM), combined with a fast sampling method, fastRG (Alg. 1, developed in [30]). Such an architecture is then flexible and capable of using a mixture of sparsities of attention blocks across the layers yet still capable of learning a full attention model if required. On average, however, it reduces computational complexity significantly, as documented in the well-chosen experimental section on the benchmark Long Range Arena (LRA) dataset. 

The paper builds on previous theoretical expressivity results of sparse attention transformers and applies this result to proposed architecture in two Lemmas, resulting in Theorem 1 with proofs in the Appendix.

The paper concludes on the limitation of the method. Unfortunately, GPU architectures used to train large transformers are not efficient for sparse tensor operations. Proposed SMB-Transformer thus " ...often results in longer runtimes compared to dense counterparts even though the memory usage is lower ... ". 
strengths_and_weaknesses: Overall, this is an interesting application of sampling methods (an algorithm fastRG[30]) combined with gradient descent optimization with theoretical efficiency guarantees supported by experiments while not sacrificing the performance. To the reviewer's knowledge, this is a novel approach that I believe is of high interest to the ML community.

Short term, the practical impact/significance of the work is limited due to NOT being well suited for large GPU clusters - the nowadays platform of choice for large transformer models. Also mentioned in a nicely written Related Work, there are several methods in place, e.g., Reformer [18], using hashing similarity instead of SBM, similar in the idea, be it more constrained.

The paper could also be improved in terms of accessibility and readability, especially for a reader not familiar with SBM and generative modeling. For instance, it could help to elaborate in detail how a stochastic block model version from paper [2] (I believed that one is used), translates to SMB-transformer proposed. Especially in section 4.1.

Theorem 1: The idea of the proof presented in the main body of the paper is still technical and does not spark the intuition of why a “star graph” structure ensures expressibility after rephrasing the idea of the theorem from [38] on line 213. Moreover, the proof in Appendix relies on previous works, mainly by Yun et al.,(2020) [14], and Theorem 1 is proven by showing that SMB-Transformer meets the necessary assumptions of Yun et al.,(2020) [14]. Proof of Lemma 2 shows expressibility for p=3 sparsity patterns (expressibility is not guaranteed for p=2 as noted in line: 215) and thus proves it for higher p’s. However, it rather seems like a direct application of previously mentioned works on this special case of SMB-Transformer. Could you elaborate on what is a novelty here? Otherwise, I suggest being specific and making it an application of Theorem XY from [] and [] ...

line 285 and Table 2. Given the experimental std/errors (in brackets?) presented in the table, I am not sure we can conclude that “… applying a small sparsity regularizer helps in boosting overall sparsity as well as performance …”.  I would claim it is rather insignificant looking at average results in the last column! I agree, however, the computation saving effect of regularization vs. no reg. in Figure 5 still holds. So I suggest considering rephrasing or omitting the claim above perhaps.

Minor comments: 
typo: typo Table 1 “the the”
questions: To help accessibility of the article (see above) I suggest rewriting/extending the lines between 145 - 164. For instance why bipartite structure? Or what is 2-layer MLPd_h to d_h and why two layers (most likely to retain dimensionality but could you put some motivating lines of thoughts down?)

Are there any numerical or other limitations incurred by the use of SBM? In particular, this generative model assumes graph edges follow (mean parameterized) Poisson distribution. Besides one mentioned (line 182) regarding a degeneration of the mask matrix M, I'd encourage authors to elaborate more on this topic. For instance, is there any limitation in terms of learning "outliers" that is small clusters but highly connected clusters/tokens? (SMB is known to work well for the balanced size of clusters, but may struggle for imbalanced ones.)

Please find more suggestions in the previous block of comments.
limitations: A large part of the conclusions is dedicated to a major limitation - despite using fewer memory run-times on GPU clusters are often longer for SBM-Transformer than for dense models. In my opinion, this is unfortunate concerning the practical impact of this work because a majority of large transformer models GPT-3, PaLM, OPT175B, etc. have been trained in more and more accessible High-Performance Computing (HPC) centers, i.e. supercomputers, with massive GPU partitions.

Secondly all results hold "in expectation" (in a weak sense). Could authors elaborate more on the numerical stability/deterioration of the algorithm?
review_title: null

Meta-review: metareview: This paper focuses on sparse attention modules for improving the computational cost of the transformer. The authors propose to adaptively learn the sparsity of the attention module through modeling the mask matrix as a bipartite graph generated by the stochastic block model (SBM).

While there are some concerns, such as the method not being well suited for large GPU clusters, overall, all the reviewers find the proposed method interesting and novel, so I recommend accept. But the authors are advised to revise the presentation of the SBM to improve the accessibility and readability and incorporate other reviewers' suggestions. 
confidence: Less certain
Decision: Accept