Higher Order Transformers: Efficient Attention Mechanism for Tensor Structured Data

Authors: Soroush Omranpour, Guillaume Rabusseau, Reihaneh Rabbany
Primary area: foundation or frontier models, including LLMs
Keywords: transformers, multihead attention, high order tensor, kronecker decomposition, multivariate timeseries forecasting, 3D medical image classification

Abstract:
Transformers are now ubiquitous for sequence modeling tasks, but their extension to multi-dimensional data remains a challenge due to the quadratic cost of the attention mechanism.  In this paper, we propose Higher-Order Transformers (HOT), a novel architecture designed to efficiently process data with more than two axes, i.e. higher-order tensors. 
To address the computational challenges associated with high-order tensor attention, we introduce a novel Kronecker factorized attention mechanism that reduces the attention cost to quadratic in each axis' dimension, rather than quadratic in the total size of the input tensor. To further enhance efficiency, HOT leverages kernelized attention, reducing the complexity to linear. This strategy maintains the model's expressiveness while enabling scalable attention computation.
We validate the effectiveness of HOT on two high-dimensional tasks, including multivariate time series forecasting, and 3D medical image classification. Experimental results demonstrate that HOT achieves competitive performance while significantly improving computational efficiency, showcasing its potential for tackling a wide range of complex, multi-dimensional data.

Decision: reject
Reviewer scores: 3, 6, 3, 3

Reviews and discussion:
<thread>
creation date:
1729267207040

thread_type:
official_review

<thread_object>
creator: reviewer_ccuz

<content>
summary:
The authors propose a high-order Transformer (HOT) for high-order tensor input. The authors apply Kronecker tensor decomposition and kernel attention in Performer to propose a novel Kronecker factorized attention. The experimental results in time-series forecasting and 3D medical image classification illustrate the efficient of HOT.

strengths:
1. The authors skillfully perform tensor operations.

weaknesses:
1. This paper demonstrates a limited degree of innovation. The proposed methodology in this paper is merely based on Kronecker tensor decomposition [1] and kernel attention [2].

[1] Phan, K. (2012). On Revealing Replicating Structures in Multiway Data: A Novel Tensor Decomposition Approach. In Latent Variable Analysis and Signal Separation (pp. 297–305). Springer Berlin Heidelberg.

[2] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, \& Adrian Weller (2021). Rethinking Attention with Performers. In International Conference on Learning Representations.

questions:
1. How does the use of low-rank approximation via Kronecker decomposition in Section 4.2 impact the trade-off between computational complexity and precision? The authors need to provide more analysis on this and offer guidelines for choosing the appropriate rank R or kernels to help practitioners adapt the method to different datasets and needs.

2. The authors need to provide more insights or proof on how HOT maintains low complexity when handling datasets with more than 3 dimensions. It would be helpful to understand how the model scales and whether the current complexity reduction techniques, like Kronecker decomposition remain effective as the dimensionality increases.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730173623530

thread_type:
official_review

<thread_object>
creator: reviewer_b21d

<content>
summary:
This work proposed a transformer architecture to efficiently deal with higher-order structure data.

strengths:
1. The research question is important since the attention operator is very computationally expensive on higher-order data.
2. The usage of Kronecker to factorize attention is sound.
3. Empirical studies demonstrate the effectiveness of proposed work on higher-order data.

weaknesses:
1. The major problem is the novelty. This work is essentially about attention instead of transformer. The Kronecker factorization is applied to attention, which is part of the transformer architecture. Actually, applying Kronecker factorization to attention on higher-order data has been explored in [1]. Currently, the novelty is very marginal from my current understanding. It is important for authors to demonstrate the distinct contributions from [1]. 

2. Since the proposed attention is not necessarily integrated with the transformer architecture, it would be better to show its effectiveness in other architectures, such as Attention U-Net [2].

3. [Added after discussion with authors] The evaluation is actually conducted on 1D data. The time series data is reshaped into 2D to mimic higher-order data. However, this does not make sense since there is no locality in the feature dimension of 1D data. It does not make sense to apply this on 1D data although it is mathematically correct. To improve, the evaluation on real 2D and 3D data is preferred instead of using 1D data.

[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).

questions:
Please see weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730892153671

thread_type:
official_review

<thread_object>
creator: reviewer_gjkf

<content>
summary:
This paper proposes a novel Transformer architecture to extend traditional Transformers for processing higher-order data, while reducing the attention cost to quadratic in each axis' dimension and the complexity to linear.

strengths:
* Considering the importance of higher-order data, designing specialized Transformer architectures for higher-order data is meaningful.

* The proposed Higher-Order Transformers (HOT) can efficiently capture essential cross-dimensional dependencies with lower complexity, thus overcoming the limitations of previous methods.

* The experimental results also support the superiority of the proposed HOT, both in effectiveness and efficiency.

weaknesses:
* Different form previous works which directly flatten the input tensor into a sequence, it's impossible for HOT to use the same architecture process data with different orders, thus limiting its scalability.

* Consideing the high-order feature of HOT, it's necessary to compare the model parameters with its baselines.

* Besides, it's also necessary to conduct more experiments to demonstrate the generality of the proposed method on various kinds of higher-order data, such as typical 2D image datasets and video dataset.

questions:
Please see the **weaknesses**.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730938636440

thread_type:
official_review

<thread_object>
creator: reviewer_2h8g

<content>
summary:
This paper proposes a new transformer-based model, i.e., Higher-Order Transformers (HOT) which can efficiently process data with more than two axes (i.e., higher order tensors). In addition, the authors propose a novel Kronecker factorized attention mechanism that reduces the attention cost and provide theoretical proof of Kronecker decomposition. The experiments over multivariate time series data and 3D medical images are straightforward and clear.

strengths:
+ The paper is clear and well written. It is easy to follow and presents a concise idea of Kronecker factorized attention mechanism based on Kronecker decomposition. 
+ The paper also do a good job in describing existing research/roadmap in transformer-based models.
+ For multivariate forecasting tasks, the proposed HOT model always outperforms baselines.

weaknesses:
- For multivariate forecasting tasks, the model's effectiveness is assessed only through MAE and MSE. Can the authors also report other metrics, e.g., MAPE? In addition, I am curious about the confidence intervals for prediction of the HOT model compared to other baselines.
- Why Kronecker decomposition: there are also other higher-order structure decomposition approach, e.g., CP decomposition, Tucker decomposition, etc. I wonder if the authors have compared Kronecker decomposition with other tensor decomposition methods? 
- Improvements are minor: although the HOT model achieves the best prediction accuracy across all datasets for multivariate forecasting tasks, improvements are not significant. Also, for 3D medical images, the HOT model's performances are less competitive. 
- It would be helpful if the authors could conduct robust experimental comparisons, i.e., if the proposed Kronecker factorized attention mechanism can help improve model's robustness.

questions:
See my questions from the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732086267618

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers

comment:
We sincerely thank all the reviewers for their time and insightful comments. We are excited that all of the reviewers acknowledged the effectiveness and simplicity of the proposed method, and we appreciate that they found our paper is well motivated (Reviewers **Gjkf, B21D**), well presented (Reviewers **B21D, Gjkf, 2H8G**), and mathematically sound (Reviewer **B21D, ccUZ**).


To the best of our efforts, we have provided detailed comments to address all concerns raised by each reviewer. Meanwhile, we have carefully revised the paper and appendices where the main modifications are highlighted in blue. Specifically, the main revisions we made are as follows:

 - Updated Section 5.3.1 by adding explanation about the choice of rank and the effectiveness of the low-rank approximation.
 - Updated the AUC scores for 3D medical image classification results in Tables 2 and 4 using the same function as the original baselines to make the results more consistent for fair comparison.
 - Added Section 5.3.3 to study the effect of using the proposed Kronecker factorization and the attention kernel trick on the performance and efficiency. 
 - Updated the efficiency analysis in Figure 4 based on the experiments ran on an Nvidia A100 GPU.
 - Added Robustness Analysis of the HOT performance on both tasks in the Tables 10 and 11 in the appendix.
 - Reported SMAPE as additional metric on timeseries forecasting in Table 11.
 - Added a comparison on the number of parameters with the baselines in Tables 1 and 2.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this submission, the authors introduce a high-order Transformer associated with a factorization-based acceleration strategy. However, the reviewers have concerns about the experimental settings and the performance of the proposed model. In particular, introducing high-order decomposition may increase runtime and memory costs significantly, but the authors did not provide comparisons for various methods on their runtime and memory costs. Compared with the increase in computations, whether the improvements are significant is questionable. Although the authors made efforts to solve the concerns of the reviewers, the reviewers still think the authors need to enhance the experimental part, and accordingly, the submission requires a next-round review.