Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: Focusing on the multi-head self-attention (MHSA) structure, this paper proposes an extension of the tensor diagram to denote self-attention (SA) structures more intuitively. Then inspired by the Tucker format, this paper also proposes a new form of SA named Tucker-Head Self-Attention (THSA), which can also be illustrated with the tensor diagram extension. Although THSA derives some success, THSA may lack theoretical guarantee and experiments seems insufficient.
main_review: Pros:
- Clear writing.
- The proposed idea is somewhat novel and interesting.
	
Cons:
- Motivation of THSA is not clear. Although the paper claims "inspired by a Tucker tensor form", what is the advantage of the Tucker tensor form? And why MHSA can be benefited from the Tucker tensor form?
- The ability of extended tensor diagram is limited. A traditional tensor diagram need not consider the combination order for diagram nodes. However, the proposed extension requires contracting the nodes in the softmax box firstly. I agree that the diagram extension is sufficient to present the proposed method. Nevertheless, it may need to elaborate the limitation clearly, for example,
	- Difficult to denote multi-layers due to the order of softmax box;
	- More attention when applying the tensor diagram extension to other cases besides the self-attention condition. 
- More situations are supposed to be considered as the diagram extension is firstly proposed:
	- Is there a probability to take a node out from the box? If possible, what transformation needs to be applied?
	- Will there be a case multi-box overlap? When this case happened, how to do it?
- In Definition 2, the expressive power seems not to make sense. Under the definition, a model with higher expressive power may not perform well. For instance, an over-fitting model may have large expressive power, however, its performance is lower, compared to a suitable model. Therefore, it may not provide a guarantee to THSA performance. By the way, I guess there is a typo in Definition 2, "$f \in \mathcal{X}$" - > "$f \in F$".
- In Table 1, the paper seems not to mention that the results are all derived through re-implementing. And the results in the original paper are higher (e.g., ALBERT: MNLI-91.3%, QNLI-99.2%), which makes a concern that the experiment is not prepared well.
- Insufficient experiments. There are plenty of new works proposed recently. And if THSA is really consistently better than MHSA, THSA is supposed to achieve SOTA results on a task.
- More training run-time for THSA. Will THSA (same num of param) be slower when compared with MHSA?
summary_of_the_review: Although this paper has a good presentation, for lacking theoretical guarantee and insufficient experiments, I regret to give a weak rejection.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper investigates the multi-head self-attention mechanism (MHSA) of transformer networks through the lens of tensor decompositions via tensor diagram notation. The authors propose an extension to MHSA inspired by the Tucker decomposition (termed THSA), analyze its expressive power, and demonstrate that it belongs to a class of more expressive functions than MHSA. Further, the authors show the positive effect of this drop-in replacement on the downstream tasks.
main_review: Questions
===

How can a "super node" in Fig.4 be seen as a single tensor if a contraction of it crosses the boundary of a non-linear operation? This overloads the tensor diagram notation in an undesirable way, violating the multilinearity assumption. Taking this notation to the extreme, an MLP can be represented as a single "super-node" matrix, which spans multiple nonlinearities. Such a convention for tensor diagrams would be not much more useful than a typical block scheme of most neural architectures.

Tucker decomposition of a 4d-tensor leads to a 4d core tensor. If the core is a matrix, the authors should either carefully specify which "variant" of the decomposition is used (with specific references to external sources) or use a different term for the proposed parameterization.

The notation of T-junctions used to connect "stacked" modes to other nodes (such as to the matrix C from the top in Fig.4b) is crucial for understanding the proposed module; however, it is missing. 

"The concept of the number of heads in THSA generalizes to the stable rank of the core matrix" -- why can't we draw a similar connection with the standard MHSA? 

Related to the previous question, formulas (2) operate an integer number of heads (R) - how does this go along with the claim of a learned number of heads?

Formula (2b) rewrites a product of QK^T matrices for a single head through indexing along the R mode using s subscript. This corresponds to the rank factorization of the said product. This also seems to be equivalent to a single-head self-attention with R features instead of D, due to the fact that the values C_{rs} can be absorbed into either W^Q or W^K (which are in turn head-specific and may need an _{[r]} subscript). Or, if matrices W^Q and W^K are not head-specific anymore, this looks like a weight sharing scheme, in which one of the matrices is head-specific, and the other one is reused across heads. The main questions are: (1) where is Tucker decomposition in this formulation, (2) how does this equation correspond to Fig.4b, and (3) how exactly is this different from a standard MHSA?

Could the empirical advantage be due to a changed parameterization of W^K and W^Q matrices? I would like to see an experiment with the same configuration as MHSA in terms of dimensions of activations (probably corresponds to +25% configuration from the paper) but done in a way that the fifth matrix C is introduced into the QK product (using small surgical changes to the original MHSA module code). Such an added matrix experiment would level the changes with the proposed THSA and allow us to see whether the improvement is due to such parameterization of the QK product. Since such decomposition does not increase the expressive power of the QK product, any change in performance is due to either a different initialization or learning rate of the product elements resulting from such parameterization. If weight sharing is the case, additional studies of sharing W^Q or W^K in an MHSA setting are required.

The claim about the possibility of initializing THSA from MHSA should be followed with a remark that this is only possible in the case of matching dimensions. 

The variance across runs in Fig 5 looks small - was RNG seed changed explicitly between runs to affect weight initializations, as opposed to relying on the intrinsic variance of SGD and randomness of multi-processing data loading in pytorch?

Around formulas (2), the first dimensions of Q and X^Q are F and N, respectively, but how is this possible if Q=X^Q W^Q?


Delivery
===
Abstract: "the number of heads is trainable" is misleading, as this number is discrete and is not directly optimized.

Abstract: "allows initialization from existing pre-trained models" is too broad, given that this is only possible in cases with a specific R.

The first occurrence of "expressive power" (except abstract) requires a clarification of the subject (a class of functions) - either as a citation or as a forward reference to the definition from the supplementary materials. It should also be made clear how expressive power is compared to explain when used together with "higher".

"(2) The number of heads is not trainable": not clear how this is a limitation or a problem.

"re-parameterization of the weight matrices via a variant of Tucker Decomposition or Tucker Representation" - the provided reference is too broad given the word "variant". A default Tucker decomposition assumes a core tensor of the same dimensionality as the decomposed tensor, which is not the case in this paper.

The proposed visualization of the softmax operation in the tensor diagram is good but requires clarification that normalization is applied along the mode containing the "alpha" point.

The "Evaluatuon order" section would benefit from a remark that the order affects the computational complexity of the overall contraction.

While the notation of stacking matrices to obtain a 3d-tensor is clear, connecting these stacked modes together requires additional explanation. Furthermore, if this notation is used purely to denote the stacked nature of nodes, it should be made clear, too (and perhaps such edges can be color-coded differently). 

The related work section misses an overview of tensor decompositions used in machine learning.


Writing
===

Usage of the word "guaranteed" throughout the text requires explaining: by what means? If through a proof - the word "proven" is better suited.

"(2) The number of heads is not trainable ... capture different context information through each head": if this is a known fact, the authors should consider adding a reference where this is stated.

The first occurrence of "tensor representation theory" - a reference is required (e.g., Kolda)

"we also propose a novel data-driven structure that" - too vague.

"re-parameterization of the weight matrices" - which ones / of what?

"for instant, in Section 2.1" -> "for instance, in Fig. 2.3"

notation for subscript "_{[h]}" not introduced

Texts of "Single-head self-attention in tensor diagram" and "Multi-head self-attention in tensor diagram" sections are nearly identical; Fig. 3a is never referenced nor explained. 


Code
===
The provided source code is not readily usable for reproducibility. A few problems and recommendations regarding the file with the implementation referred to in the readme:

1. The code contains multiple python syntax errors
2. Local variable tucker_heads is not defined in a few places; perhaps "self." is missing
3. Unused code clutter (all branches where self.tucker_heads is False) 
4. Unused variables (e.g., c = self.c_proj(...))
summary_of_the_review: The approach to analyzing MHSA taken by the authors is interesting and worth digging deeper, especially due to the clear empirical advantage of the proposed modification. However, the paper contains several issues in claims, theory, and connections between them. On the side of delivery, the paper has good potential, but in its current state, it is not well-polished and requires a major improvement to facilitate better understanding.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper attempts for a better understanding of the mechanism of multi-head self-attention (MHSA). It proposes to interpret MHSA weight matrices as factors of the Tucker Decomposition and adds an additional core weight tensor to MHSA to obtain Tucker-head self-attention (THSA).
main_review: The main idea of the paper, to interpret the weight tensors of MHSA as factors of a decomposition, is new and THSA appears to be supported by experimental results (although the results are so small that I wonder whether they are significant relative to hyper-parameter tuning and stochastic fluctuations of the training). However, no evidence is provided to justify the choice - why should it be Tucker Decomposition, and not Tensor Train or Tensor Ring decomposition? And how does that choice influence the properties of the attention mechanism?

The authors also claim the following contributions:
1. tensor diagram notation for attention mechanism
2. design space for attention block together with expressiveness analysis framework
3. initialisation scheme from pre-trained Transformer

While (3) has important applications in reusing pre-trained transformer models, (1) doesn't help much to understand the attention mechanism, as it will be readable only for people whose multi-linear algebra background is advanced enough to also reach the same understanding without diagrams. Whereas (2) effectively proves that, given some network, adding an additional parametric matrix somewhere in the middle of that network does not decrease its expressiveness, which is trivial

In my view, the general idea is valid and potentially interesting, but it would make the paper much stronger to revise it in terms of methodology and experiments, so as to actually study how different tensor decompositions influence the attention mechanism, and which ones might be suited best for specific layouts and tasks. However, this would be a completely different paper that would have to be reviewed again, so I am on the fence about the present version.

The presentation is mostly clear, but in a few places it is sloppy and appears unfinished and rushed. This would need to be reworked:
page 4 - duplicate paragraphs at the end of the page
page 5 - duplicate paragraphs in 3. and 3.1.
Definition 2: f \in X should be probably f \in F ?
(page 3: For instant -> For instance)
summary_of_the_review: The paper address an interesting and relevant problem, namely generalisation of the MHSA mechanism, and it presents interesting, if somewhat terse and preliminary, experimental results. However, while the view as tensor decomposition is promising, the seemingly ad-hoc commitment to Tucker without a clear rationale or systematic study is unsatisfactory and leaves an unfinished impression. The paper does put forward a novel idea and has high potential, but it would require a major overhaul to fulfil that potential and turn it into a possibly really strong and useful contribution, rather than a fleeting quick-shot.

I find it difficult to give a recommendation for this paper: If one follows the conventional criteria, there is novelty and the potential to lead to interesting follow-up, so one might argue for acceptance, even though there isn't a strong case for future impact. On the other hand, it feels a bit premature and a bit more work could make it a lot better and more significant.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: 
The paper proposes a THSA that generalizes the standard self-attention. THSA empirically outperforms the standard self-attention.
main_review: Tensor network diagram seems unnecessary to the paper. Using Eq. (2a,3b,2c) and Eq. (1a,1b,1c) is enough for readers to understand the difference between MHSA and THSA. The diagrams in figure 4 convey the same idea but involve more technical content. The first contribution claimed by the author could be safely erased without affecting the completeness of this work.  Let me know if I am wrong.

In Sec. 3.2, as a reader, I am more curious about the proof for theorem 3 than definition 2. But the authors put definition 2 in the main paper and proof for theorem 3 in the appendix -- why not do the opposite?


The author claimed "THSA *significantly* improves MHSA", I am wondering whether the author actually did the significant test (as I encourage the authors to do it), or it would better use another word like "notably" instead.


Please consider *quantitatively* showing the inference time.

Subscripts in Figure (3) seem wrong, please have a check. A_mr B_nr not  A_mr B_nr. The same happens in the text related to Eq 2, please double-check the superscripts there (the shape of X^Q, W^Q, Q, Q_r).


The technical content is quite close to  Cordonnier et.al 2020 https://openreview.net/forum?id=bK-rJMKrOsm , please consider explaining the difference. See figure 2  (left) in Cordonnier et.al 2020, the standard SAN can be considered as a special case of theirs. This is quite similar to  Proposition 1 in this paper -- of course, they are not the same, but with some differences. 




summary_of_the_review: The paper introduces a general Self-attention, which is interesting and insightful. However, a general model that costs more (time or space) but slightly performs better is not that attractive.
review_title: null

Meta-review: This paper presents a tensor diagram view of the multi-headed self-attention (MHSA) mechanism used in Transformer architectures, and by modifying the tensor diagram, introduces a strict generalization of MHSA called the Tucker-head self attention (THSA) mechanism. While there is some concern regarding the incremental nature of the proposition, the identification of where to usefully add the additional parameter that converts from MHSA to THSA was nontrivial, and the experimental results on the performance benefits across multiple tasks is convincing.
Decision: Accept (Poster)