Higher Order Transformers: Efficient Attention Mechanism for Tensor Structured Data

We thank the reviewer for the valuable feedback and acknowledging our contributions and the effectiveness of the proposed method. We especially thank the reviewer for pointing out the unclarity in the writing about the rank analysis and the choice of attention kernel. First we address the concern about the novelty of our work:

• Regarding the argument about the novelty and innovation in our work, although both techniques used in our method, Kronecker decomposition and kernel attention, have already been introduced earlier and widely used, they have never been studied in the context of efficient yet expressive high order attention. Moreover, we provided the general formulation, mathematical formalism, and theoretical guarantee of high order attention any high-order input data with no specific assumptions about the size and structure of the data. Furthermore, we provided comprehensive comparisons with baselines, ablation studies of various aspects of the proposed method, and efficiency analysis on two higher-order tasks, namely multivariate timeseries forecasting and 3D medical image classification.

Now, we answer the questions asked by the reviewer:

• Question 1: As shown in Section 5.3.1, we studied the effect of different numbers of attention heads (a.k.a. attention rank) for both tasks. Although the effective rank for the Kronecker decomposition varies for different tasks and datasets and is highly dependent on the data size and its characteristics, in practice, we see that low-rank approximation (where the rank or the number of heads is not a large number) is expressive enough to yield good results. To that end, we treat the rank as a hyper-parameter of the model and choose it by hyper-parameter search based on validation metrics for each dataset. This approach makes it very easy to achieve a proper performance, without any need for exhaustive analysis on the data. For the choice of attention kernel, we fixed the feature map function to SMReg from the Performer paper for all our models, as it was shown to be the most stable, fast converging and expressive by the authors. It is worth mentioning that we have added Section 5.3.3 in the revised paper to study the effect of the proposed factorization and using the kernel trick on the performance and efficiency. As shown in Tables 4 and 5, with proper hyper-parameter tuning (including on rank), HOT can yield as good results as the full attention with much less memory footprint and training time.

• Question 2: As shown in Section 4.3 the computational complexity of HOT is $\mathcal{O}(D^2\Pi N_i)$, where D is the attention dimension and $N_i$s are the input tensor dimensions, scaling linearly with respect to the input dimensions and quadratic with respect to the model’s hidden dimension.

We have updated Section 5.3.3 and Implementation Details in the appendix to address the ambiguity pointed out in Q1. We hope we have addressed all your concerns, in which case we kindly ask you to consider increasing your score. We remain of course available to answer any further questions.

We thank the reviewer for the valuable feedback and specially appreciate the reviewer for acknowledging the importance of the research question. First we address the concern about the novelty of our work:

• Regarding the argument about the novelty and innovation in our work, although both techniques used in our method, Kronecker decomposition and kernel attention, have already been introduced earlier and widely used, they have never been studied in the context of efficient yet expressive high order attention. Moreover, we provided the general formulation, mathematical formalism, and theoretical guarantee of high order attention any high-order input data with no specific assumptions about the size and structure of the data. Furthermore, we provided comprehensive comparisons with baselines, ablation studies of various aspects of the proposed method, and efficiency analysis on two higher-order tasks, namely multivariate timeseries forecasting and 3D medical image classification.

Now, we address the weaknesses as follows:

• Both of the papers propose an attention mechanism to reduce the memory and computation complexity of attention networks applied to high order data by using some notion of tensor operation in their formulation. However, there are several key differences:

  • KAN [1] assumes a 2D structure and a multivariate normal distribution for the input tensor and models the covariance using a kronecker sum of covariance matrices for rows and columns. HOT, on the other hand, does not make any prior assumptions on the distribution or the tensor structure of the input data and it can be applied to any-order multi-modal tensors.
  • KAN does not actually perform any Kronecker Decompositions or Kronecker Products in practice. Instead, it proposes to take the mean over the rows and columns of query/key/value matrices instead of their vectorized version which is at the end equivalent to the kronecker sum. On the other hand, HOT utilises a rank R Kronecker Decomposition of the high order full attention, making it easy to control the complexity-expressivity tradeoff and interpret first-order attention matrices.

  In conclusion, even though their names are very similar the two methods are drastically different, hence the novelty of our contribution is not diminished by this previous work. We thank the reviewer for bringing this reference to our attention, we added a citation and discussion of this work in the Related Work section.

• Although the use of the proposed factorized linear attention was only studied in the context of the Transformers architecture, its capabilities in other architectures such as Attention U-net [2] is yet to be explored. We have added this as one of the possible future directions of our work. As mentioned in the conlusion, we believe that high-order attention can be a powerful alternative to the current attention modules used in image and video diffusion models (specifically within the U-Net module) to enhance the spatial and temporal coherence of the outputs.

We hope we have addressed all your concerns, in which case we kindly ask you to consider increasing your score. We remain of course available to answer any further questions.

[1] Gao, Hongyang, Zhengyang Wang, and Shuiwang Ji. "Kronecker attention networks." In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 229-237. 2020.

[2] Oktay, Ozan, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazunari Misawa, Kensaku Mori et al. "Attention u-net: Learning where to look for the pancreas." arXiv preprint arXiv:1804.03999 (2018).

Thank you for your follow-up comment.

Regarding weakness 1, first, we would like to provide a deeper comparison between KAN and HOT on the theoretical level:

• KAN models the covariance matrix of the input data using a Kronecker sum operation, assuming that the rows and columns of the input data follow matrix-variate normal distribution which may not be necessarily accurate in most cases and also imposes independence between rows and columns. In practice, KAN reduces the full attention matrix $S \in \mathbb{R}^{N_1 N_2 … N_k \times N_1 N_2 … N_k}$ to an attention matrix $\tilde{S} \in \mathbb{R}^{(N_1+N_2+…+N_k) \times (N_1+N_2+…+N_k)}$ if extended to higher orders. This dramatic reduction, while makes KAN very efficient, loses a lot of information and prevents the attention from capturing any high-order interactions between axes more than first and second order even when applied to data with higher orders. Thus, KAN does not naturally extend due to its reliance on 2D assumptions. On the other hand, HOT relies on theoretical guarantee of rank R Kronecker decomposition while also incorporating RoPE on every axis to preserve structural and positional information. It can also balance expressivity and scalability by allowing flexible trade-offs through the rank.

• KAN and HOT can be implemented with the same number of parameters and yield the same computational complexities.

On the experimental level, while providing performance comparison with KAN on the Imagenet-1k dataset is not feasible for us within the short time window of the discussion period, we will report HOT performance on this benchmark in the final revision. However, we have provided partial comparison of HOT and KAN on three timeseries datasets (as 2D datasets) in the following. These results are obtained using this implementation and integrating it into the same model configuration used in all our experiments:

Regarding Weakness 2, we would like to emphasize that the scope of this work was intentionally constrained to the Transformer architecture for the following reasons:

• Transformers offer a reliable and modular baseline for systematically testing new attention mechanisms. Using Transformers ensures consistent experimental conditions, enabling a clear assessment of the proposed Kronecker factorized attention's performance, scalability, and efficiency.
• Transformers are highly versatile and have been successfully applied across multiple domains, including NLP, computer vision, and time-series analysis. This versatility makes them an ideal choice to showcase the utility of our proposed method on tasks such as multivariate time series forecasting and 3D medical image classification where best performing methods also follow Transformers architecture.
• Validating Kronecker factorized attention within the Transformer architecture establishes a groundwork for extending this mechanism to other models, such as Attention U-Net, as suggested, or further specialized architectures in areas like generative modeling or biomedical applications.

We hope we have addressed all your concerns, in which case we kindly ask you to consider increasing your score. We remain of course available to answer any further questions.

We thank the reviewer for the valuable comments on our paper and appreciating the contribution and effectiveness of our proposed method. Below we answer the questions raised by the reviewer:

• Questions 1, and 2: The architecture and the number of model parameters stay the same for each dataset regardless of the attention order and other configurations as shown in ablation studies in Section 5.3.2. Additionally, we have added Section 5.3.3 in the paper to compare the factorized attention with the naive full attention (aka. flattened attention). This study shows that HOT is as easy to use and scalable as the flattening approach while being much more efficient and less memory intensive while using the same model weights. We have also added a new column to the Tables 1 and 2 to compare the number of parameters with the baselines for both tasks and HOT is always among the models with the fewest parameters.

• Question 3: As already mentioned in the conclusion and future direction, we acknowledge that conducting more experiments on other high-order datasets including image and video could showcase even more the effectiveness of HOT. While it is not feasible for us to provide results on such datasets within the short time window of the discussion period, we will, however, explore HOT performance on ImageNet-1K dataset compared to popular baselines such as ViT in the final revision.

We hope we have addressed all your concerns, in which case we kindly ask you to consider increasing your score. We remain of course available to answer any further questions.

We thank the reviewer for the valuable feedback and acknowledge our technical contributions and the effectiveness of the proposed method. We address the concerns raised by the reviewer as follows:

• Question 1: As an addition to MAE and MSE we have added the reported values for SMAPE or symmetric MAPE in Table 11 in the appendix as an alternative to MAPE to make the results more interpretable. The are two main reasons for this choice:
  • MAPE treats positive and negative errors asymmetrically, potentially leading to biased interpretations, especially in long-range forecasting where both over- and under-predictions are important.
  • When actual values are zero or close to zero, MAPE can produce unstable or out-of-bound values..

We are currently working on adding confidence intervals visualisations and we will post a comment with updated results as soon as they are available.

• Question 2: The Kronecker decomposition we use for HOT is actually equivalent to a CP decomposition, this is the key element of the proof of our theorem. More precisely, the CP decomposition of the attention matrix reshaped into a tensor of size $N_1^2\times N_2^2 \times \cdots \times N_k^2$ is exactly equivalent to the Kronecker factorization of the attention matrix we use for HOT. We added a discussion to clarify this point in section B in the appendix after the proof. Other decomposition approaches such as Tucker, Tensor-train, etc. yield different structures that could be investigated as future work but are out of scope of the present study. The choice of Kronecker decomposition was initially motivated by three main reasons:

  • Simplicity and nice properties that enable us to link it easily to the first-order attention matrices that are already ubiquitous and easy to interpret.
  • Approximation guarantee when the rank increases.
  • Integrates well with the kernel trick to make attention linear in complexity over all axes while maintaining flexiblity.

• Question 3: We adopted the metric computation function for 3D medical image classification task from the MedMNIST python package to be consistent with the rest of the baselines and updated the results accordingly in Tables 2 and 4. This change slightly improved AUC over multiple datasets and made HOT, the best performing model on Organ, Fracture, Adrenal, and Vessel datasets and the second best on the Synapse dataset. Additionally, it is worth mentioning that the main promise of HOT is to give acceptable performance while being flexible, scalable, and highly efficient compared to the rest of the baselines.

• Question 4: We have added the robustness analysis of HOT over both tasks in the appendix in Tables 10 and 11. The results were obtained from five different random seeds.

We hope we have addressed all your concerns, in which case we kindly ask you to consider increasing your score. We remain of course available to answer any further questions.

We appreciate your acknowledgment of the improved clarity and theoretical foundation of our approach. While we recognize the importance of exploring various tensor decomposition methods, the choice of Kronecker decomposition in this work was motivated by its simplicity, scalability, and expressivity. A comprehensive numerical comparison would require additional experimentation beyond the scope of this study. However, we plan to investigate these methods in future work to explore their potential impact on performance and efficiency. For now, to further justify our choice of Kronecker decomposition, we provide a detailed analysis of how other well-known tensor factorization methods could be integrated into higher-order attention and their trade-offs compared to our proposed factorization.

CP Decomposition

As highlighted in Section B of the appendix, CP decomposition is mathematically equivalent to the proposed Kronecker decomposition. Specifically:
$S_{\text{attention}} = \sum_{r}^R S_r^{(1)} \otimes S_r^{(2)} \otimes \cdots \otimes S_r^{(k)}$
where $S_r^{(i)} \in \mathbb{R}^{N_i \times N_i}$ represents the first-order attention matrix along axis $i$.

• Number of parameters: $\mathcal{O}(D^2)$
• Computational complexity: $\mathcal{O}(D (\sum N_i) \prod N_i)$

Since the parameterization and complexity of CP and Kronecker decompositions are identical in our setting, CP decomposition does not provide additional benefits in terms of expressivity or efficiency.

Tucker Decomposition

For Tucker decomposition with a fixed rank $R$ for all axes, the attention matrix is expressed as:

$S_{\text{attention}} = \sum_{i_1, i_2, ..., i_k}^R \mathcal{G_{i_1, i_2, ..., i_k}} S_{i_1}^{(1)} \otimes S_{i_2}^{(2)} \otimes \cdots \otimes S_{i_k}^{(k)}$

where $S_{i_j}^{(j)} \in \mathbb{R}^{N_j \times N_j}$ is the first-order attention matrix for axis $i$, and $\mathcal{G} \in \mathbb{R}^{R^k}$ is a core tensor that enhances expressivity by introducing additional parameters.

While the addition of the core tensor $\mathcal{G}$ increases expressivity, it significantly impacts scalability and efficiency:

• Number of parameters: $\mathcal{O}(D^2 + R^k)$
• Computational complexity: $\mathcal{O}(D (R^k + \sum N_i) \prod N_i)$

The exponential growth in parameters and complexity with $R^k$ makes Tucker decomposition less suitable for highly scalable and efficient models, especially as the order $k$ of the tensor increases.

Tensor Train Decomposition

For tensor train decomposition with a fixed rank $R$ for all factor matrices, the attention matrix is expressed as:
$S_{\text{attention}} = \sum_{i_1, i_2, \ldots, i_{k-1}}^R S_{i_1}^{(1)} \otimes S_{i_1, i_2}^{(2)} \otimes S_{i_2, i_3}^{(3)} \otimes \cdots \otimes S_{i_{k-1}}^{(k)},$ where $S^{(i)} \in \mathbb{R}^{R \times N_i \times N_i}$ for $i = 1, k$, and $S^{(i)} \in \mathbb{R}^{R \times R \times N_i \times N_i}$ for $i \in [2, k-1]$.

While tensor train decomposition adds only a small overhead in computational complexity:

• Number of parameters: $\mathcal{O}(D^2)$ or potentially higher due to factor tensors.
• Computational complexity: $\mathcal{O}(D (R^2 + \sum N_i) \prod N_i)$.

However, constructing fourth-order factor tensors $S^{(i)} \in \mathbb{R}^{R \times R \times N_i \times N_i}$ is non-trivial and may require additional parameters or modifications to align with our use case. This lack of simplicity makes tensor train decomposition less appealing compared to Kronecker decomposition.

To conclude, the choice of Kronecker decomposition in our work is motivated by its:

1. Simplicity: It avoids additional overhead, such as core tensors or complex factor tensors, making it easy to implement and interpret.
2. Efficiency: Kronecker decomposition scales efficiently with the tensor order $k$, avoiding exponential growth in parameters or computations.
3. Expressivity: It balances expressivity and scalability by allowing flexible trade-offs through the rank $R$.

Moreover, our empirical evaluations confirm the effectiveness of this approach, achieving strong performance across tasks while maintaining computational efficiency. We hope this analysis provides further clarity and justifies our methodological choice, in which case we kindly ask you to consider increasing your score. We remain of course available to answer any further questions.