Higher Order Transformers With Kronecker-Structured Attention
Authors: Soroush Omranpour
Guillaume Rabusseau
Reihaneh Rabbany
Venue: Accepted by TMLR ()

Abstract: Modern datasets are increasingly high-dimensional and multiway, often represented as tensor-valued data with multi-indexed variables. While Transformers excel in sequence modeling and high-dimensional tasks, their direct application to multiway data is computationally prohibitive due to the quadratic cost of dot-product attention and the need to flatten inputs, which disrupts tensor structure and cross-dimensional dependencies.
We propose the Higher-Order Transformer (HOT), a novel factorized attention framework that represents multiway attention as sums of Kronecker products or sums of mode-wise attention matrices. HOT efficiently captures dense and sparse relationships across dimensions while preserving tensor structure. Theoretically, HOT retains the expressiveness of full high-order attention and allows complexity control via factorization rank.
Experiments on 2D and 3D datasets show that HOT achieves competitive performance in multivariate time series forecasting and image classification, with significantly reduced computational and memory costs. Visualizations of mode-wise attention matrices further reveal interpretable high-order dependencies learned by HOT, demonstrating its versatility for complex multiway data across diverse domains.

Published in Transactions on Machine Learning Research (11/2025)
Higher Order Transformers With Kronecker-Structure At- tention
Soroush Omranpour
soroush, wwwpssh.mmmmgourrmmii......u
Mila
McGill Universily
Reihanch Rabbany
haltamnch.rabbbnnymrrrti.... quebbec
Mila, CIFAR AI Chair
McGill University
Guillaume Rabusseau
rabussgu@rnila quebee
Mila DIRO, CIFAR Al Chair
University of Montreal
Reviewed on OpenReview: https: /ttps:// optpr/vee...eeee.eom net/ forum? QNOaXcKFkT
Abstract
Modern datasets are increasingly high-dimensional and multiway, often represented as lensor-valued data with multi-indexed variables. While Transformers excel in sequence modeling and high-dimensional tasks, their direct application multiway data is compu- tationally prohibitive due to the quadratic cost of dot-product attention and the need to flatten inputs, which disrupts tensor structure and cross-dimensional dependencies. We propose the Higher-Order Transformer (HOT), a novel factorized attention framework that represents multiway attention as sums of Kronecker products or sums of mode-wise atten- tion matrices. HOT efficiently captures dense and sparse relationships across dimensions while preserving tensor structure. Theoretically, HOT retains the expressiveness of full high-order attention and allows complexity control via factorization rank. Experiments 2D and 3D datasets show that HOT achieves competitive performance in multivariate time series forecasting and image classification, with significantly reduced computational and memory costs. Visualizations of mode-wise attention matrices further reveal interpretable high-order dependencies learned by HOT, demonstrating its versatility for complex multi- way data across diverse domains. The implermnnattinn of our proposed method is publicly available at /https:// com/ omranpourHHOT
1 Introduction
Multiway data, represented as multidimensional arrays or tensors, are pervasive across scientific, engineer- ing, and real-world applications. In climate modeling, multidimensionall time series capture spatiotemporal variations for weather and climate prediction (Nguyen al. 22023.. Modical imaging uses 3D modalities like MRI and CT scans to reveal anatomical details beyond traditional 2D images (Yang al. 2023). Recom- mendation systems model user-item interactions alongside contextual factors (e.g., time, location) naturally as tensors (Frolov & Oscledets 2016 Such examples underscore the ubiquity of multiway data and the need for models that capture complex interdependencies efficiently.
Tensors generalize matrices to higher dimensions, enabling structured representation and analysis of multi- dimensional relationships. They are intogral in domains ranging from quantum physics (Montangero) 2018) and algebraic geometry (Landsberg] 2012) data science and machine learning, where they encode im- ages, videos, medical scans, and more. Tensors facilitate higher-order correlations, multiway clustering, and

---

Published in Transactions on Machine Learning Research (11/2025)
dimensionality reduction via decompositions Lu et al. sis (Lu et al., 2003), dictionary learning (Bahri et al. 2019), and tensor regression (Guo et al. 2012 extend meatrx-bassed methods, mitigating the curse of dimensionality and enabling efficient high dige-dimensional mod- cling. Among the most powerful tools for multiway data are tensor factorizations like Canonical Polyadic (CP) and Tucker decompositions, which generalize PCA and SVD to higher-order settings. These met extract latent patterns and have broad applications in signal processing (e.g., noise reduction, EEG and fMRT analysis), computer vision (e.g,, compression, recognition), and machine learning (e.g., feature learning in recommendation systems, NLP, and social network analysis) (Lu et al., 2011). In fields such chemometrics and neuroscience, tensor factorization uncovers hidden structures and enables scalable analysis of complex data.
2011
Techniques such tensor component analy-
Despite these advances, applying mod- ern deep learning architectures to multiway data remains challenging.
Transformers (Vaswani et. al., 2017) have revolutionized sequence modeling in vision (Dosovitskiy et al. 2020).
Hends
Proling
speech (Dong et al. (018)), and reinforce- ment learning (Parisotto et al. 2020
N
owing to their self-attention mechanism's 's capacity for modeling long-range depen- dencies. However, their quadratic cost in time and memory hinders direct applica- tion to high-dimensional lensors, limiting their use in domains like video analysis, high-dimensional forecasting, and 3D imaging. Several strategies attempt to adapt Transformers for multiway data. Flattening tensors into sequences allows standard architectures (Dosovitskiy et al. 2020), but sacrifices structural information and local dependencies, making positional encodings insufficient. Axial attention (Ho et al.] 2019)) and spatiotemporal attention (Song et al. 2016) mitigate computational costs by processing one mode at time but may fail to capture full cross-dimensional interactions.
SwiGLE
N
Vider
RMS Nara
Product Softran
Pridect Sodmma
Q1
Krnneder Fastertaal
Mutribead Afrretion
02
N
Py
Py
Poel Projer.
Pnject
Projection Layer
N
Figure 1: Overall structure of High Order Transformer (HOT) de- picting the proposed method for 2D data with size N₁ x N₂ x D. The model shares the same arrangement as the Transformer en- coder while employing Kronecker Factorized Multihead Attention to reduce the computational complexity. Each mode of the ten- sor (c.g., N2) has its own attention matrix, combined using Kronecker product operations.
We introduce Higher-Order Transformers (HOT), a novel architecture for efficienlly modeling tensor- structured data. Inspired by Kronecker- based covariance models in statistics (Song et al. 2023 Wang et al., 2022), HOT employs a factorized high-order attention mechanism that models both dense and sparse interac- tions via Kronecker-structure formulations. Each attention head computes independent attention matrices for each modality, which are then combined through Kronecker products or sums. This design preserves the expressivity of full attention while significantly reducing computational costs. Our key contributions are follows:
We introduce Higher-Order Transformers (HOT), a factorized attention architecture that extends self-attention to tensor inputs, efficiently capturing dense and sparse cross-dimensional dependencies while preserving the tensor structure.
We provide theoretical analysis showing that Kronecker factorized attention retains the expressive- ness of full high-order attention with controllable computational complexity via factorization rank.
2

---

Published in Transactions on Machine Learning Research (11/2025)
We validate HOT across diverse Cenchmarks-iincheee multivariate time series forecasting, 3D medical image classification, and multispectral sigmentationn-demnnntaannnn competitive or superior performance with significantly reduced computation and memory usage.
2 Related Work
Traditional multiway data analysis relies tensor decompositions such as CP and Tucker extract latent structures (Kolda & Bader 2009b). While effective for linear interactions, these methods struggle with com- plex dependencies. Kronecker structures have been applied to covariance modeling, enabling large matrices to be represented as products of smaller ones (Tsiligkaridis & Hero 2013 Greenewald & Hero 2015). Recent
advances include iterative algorithms for approximating separable structures (Song et al. 2023 and sparse corrections for improved robustness (Greenewald & Hcro 2015).
Deep learning has driven significant progress in multiway data analysis through CNNs and hierarchical models (Leuun et al. 2015). However, such networks often require millions of parameters. Tensor decom- positions have been used to compress network weights while maintaining performance (Novikov et al., 2015 Lebedev et al. 2015). Transformers (Vaswani et al. 2017) excel in sequence modeling but. face challenges with high-dimensional data due to computational costs.
A common strategy flattens multi-dimennsioall data into sequences for Transformer inputs. Vision Trans- formers (ViT) divide images into patches tokens (Dosovitskiy et al. 2020), while VIVIT extends this to spatiotemporal video patches (Arnab et al. 2021). However, flattening disrupts intrinsic multiway rela- tionships. To preserve structure, tensor decomposition techniques compress Transformer models (Ma et al. 2019), and Kronecker Attention Networks (KAN) use Kronecker operators to capture second-order covari- ances without flattening (Gao et al. 2020), though with strong distributional assumptions. TimeSFormer [Bertasius et al. 2021) introduces a factorized attention mechanism for video understanding that separately applies atlention spatial and temporal dimensions in order. Multiscale Vision Transformers (MViT) (Fan) et al.) 2021) further improve efficiency by progressively reducing spatial resolution while expanding channel capacity, enabling scalable representation learning across different gramularities.
New models explicitly capture oross-dinensional dependencies. iTTranformer enhances multivariate time series forecasting by applying attention across variables (Liu et al. 2024), while Crossformer targets spatial- temporal relationships (Zhang & Yan. 2023). In medical imaging, SegFormer3D employs hicrarchical cross-
scale altention for 3D seginentation (Perera et al., 2024), and CdTransformer distills attention across or- thogonal planes to capture anatomical dependencies (Zhu al. 2024
The computational demands of self-attention have inspired effciiency-focused innovations. Tensor Product
Attention (TPA) compresses queries, keys, and values via tensor decompositions (Zhang et al. 2025
Sparse
attention Child et al. 2019 Zaheer et al. 2020) and linearized methods (Katharopoulos et al. 2020 Choro-
manski et al. 2021) reduce complexity for longer sequences. Performers approximate attention via kernel
methods (Choromanski et al. 2021), while Reformer uses locality-sensitive hashing for efficient memory
use Kitaev et al. 2020).
3 Preliminaries
In this soction, we introduce key tensor notations and operations that are fundamental to the high-order attention mechanism proposed in this work. Vectors are denoted by lowercase letters (e.g., v), matrices by uppercase letters (e.g., M), and tensors by calligraphic letters (e.g., T). We la to denote the d x d identity matrix and x; to denote the tensor product along mode i (Kolda & Bader 2009a). The notation k} for any integer k.
K] refers to the set 11.2.....k
Definition 3.1 (Tensor). k-th order tensor T E R Ni N generalizes the concept of a matrix to higher dimensions. A tensor can be viewed as a multidimensional array, where each element indexed by k distinct indices, representing multiway data that spans across l multiple modes. The number of modes is equal to the order of the tensor.
3

---

Published in Transactions on Machine Learning Research (11/2025)
Definition 3.2 (Tensor Mode and Fibers). A mode-i fiber of tensor T is the vector obtained by fixing all indices of T except the i-th one, e.g,, "n -1 E RRN.
Definition 3.3 (Tensor Slice). A tensor slice is two-dimensional section of a tensor, obtained by fixing all but two indices, e.g., Tai par.aaaann-iiiiin htw.::----..eeee.oaaaa..aa....
N
Slices and fibers extend the familiar concept of matrix rows and columns higher-dimennionall tensors, providing powerful ways to analyze and manipulate multi-way data.
Definition 3.4 (Tensor Matricization). The i-th mode matricization of a tensor rearranges the mode-i fibers of the tensor into a matrix. It is denoted as T(a E MMN:
Definition 3.5 (Mode tensor product). The mode product between tensor T
and matrix A E RdxNa is denoted by x,, A E IR- x.Nxx..xx...... X...xx...--.111 XdxNx+1 X.xXN and defined by (7 xn A) Σ,T,, An
Ain for all i1 E [Ni], ... i [Na).
Definition 3.6 (Rank-1 Tensor). An order-N tensor is of rank-1 if it can be strictly decomposed into the outer product of N vectors. A rank-1 matrix can therefore be written as X a C = ab and a rank-1 order-3 tensor as X' = a @ Q which can be extended to any-order tensors.
We conclude this part by defining the Kornecker product and sum and stating useful identity relating matricization, mode n product, and the Kronecker product.
Definition 3.7 (Kronecker Product and Sum). The Kroneeker product of two matrices A E R and B E RPx is the mp x ng block matrix
ann B
01,n B
A B
ami
Om
The Kronecker sum of two square matrices A E R and B E Raxx is the mn x mn square matrix defined by A B A @ I + Im & B.
Proposition 3.8. For any lensor T E BMIxNox--XMxx of order k + and any matrices A e R Mr XN .....,,, A E R M M we have (T X1 A1 X2 A2 X3. X As) = (1 @ A2 @ @ Ax)
4
Method
In this section, first review the self-attention mechanism in Transformer layers (Vaswani et al., 2017). which we extend to higher orders by tensorizing queries, keys, and values. Then, we explain the theoretical backbone of the proposed low-rank factorized attention and architecture details.
4.1 Generalizing Attention to Higher Order Tensors
Given input tensor X -M where N1, N N& the sizes of the positional modes and D is the hidden dimension, we start by generalizing the attention mechanism to operate over all positional modes collectively.
The simplest and most naive approach to flatten the multidimensional input into a sequence, apply standard attention on the resulting matrix, and then reshape it back to the original dimensions. We first compute the query (Q), key (K), and value (V) tensors for each head h by linear projections along the hidden dimension:
2h X X (W E -
T E RIX---NOXDD 1hh X XK+1 (
DA X X T M.x
where Xk+1 denotes multiplication along the (k + 1)-th mode (the hidden dimension).
4

---

Now we delve into the computation of the factor matrices sh.) As mentioned before, each matrix s represents first-order attention weights over the mode i. Thus, they can be computed independently using the standard scaled dot-product attention mechanism. Since the input to the attention module is a high- order tensor, computing first-order attention matrices requires reshaping of the input query, key, and value xDx + N that takes a high-order tensor as input and only preserves the i-th mode and the hidden dimension. In this Published in Transactions on Machine Learning Research (11/2025)
The scaled dot-product attention scores Sh No.. No)..Ne)
are then given by
(K(+))) (((ttt)) K (((+t)))((+))) (K+1)
thh = Softmax
(1)
vDH
where Q/--1) E M.) and K(R+1) E No are the matricization of the query and key lensors, and the Softmax function is applied row-wise. Each positional index is considered as a single entily in the attention calculation. The output of the high-order attention function hA : Xxxx I No Xxxxx"XxNDDD is computed by applying the attention weights to the value tensor:
hann (aum(d)a- = ) VV S*
(2)
Lastly, the output is reshaped back to the original tensor shape N1 No X ... X N D.
Although scaled dot-product attention widely used and has shown great promise across various domains, it comes with limitations that highly impact its scalability. While the above formulation leads compu- tational and memory complexity of .. N2))), impractical for large tensors, it is also originally designed for ID sequences and can not directly handle multiway data (e.g., images, videos, etc.) without destroying the tensor structure through flattening. These limitations motivate the development of high-order atlention mechanisms that can naturally handle high-order lensor data in an elficient manner.
4.2 Kronecker Decomposition
We parameterize the potentially large high-order attention matrix Sh (N N.. No...N using a Kronecker product or sum structure (Figure 2 of the form:
Sh
prod
(2) ... 8 S
(k)
(3)
1 (1) (s
Sh
1 + $(()))
(4)
sum
where each ((0) E NNN, is a factor matrix corresponding to the attention weights over the i-th mode for head h, and k. is a scalar normalization to ensure the row-stochastic property. Based on these formulations, we propose two variants of our method, namely HOT (product) and HOT (sum).
Theorem 4.1 (Row stochastic property). Let S E NNN and T E MM be row-stochastic matrices (i.e. Σj =1 for i, and likewise E, The for all k). Then S T NNNNM and 2 (S NMNNM are also row-stochastic.
Proof. Proof is presented in the appendix A.1
Corollary 4.2. Let (S E R.N.N. be a collection of row-stochastic matrices. Then (2))
and & (8()) @ ((2) ... + S(K) are also row-stochastic.
Proof. Proof is presented in the appendix A.1
(2) tensors. We propose to use a permutationninnariial pooling functions Spool x
work, we consider summation over all modes except the i-th and last one as the pooling function, i.c.,
Σ ( (0paai(()) (T) ier 1-
Th Te de-iddeerr d.
5

---

Published in Transactions on Machine Learning Research (11/2025)
Nj x ... x NE
R
Σ
S
=
S
N
S (
T 2
N2
S
r k
1
N1
R
R
R
Ni
N₂
Nx
Figure Visualization of rank R Kronecker Decomposition of a high-order full attention matrix E R& Na..N No....N with factor matrices S; RNxNN Note that the actual full attention matrix the LHS can be potentially much larger than what is depicted in the figure.
We then construct the query and key matrices Q¹ KP RNNDD for mode i and head h using weights With query: key DuxD as:
Q (1)
Opnol (X) W
(5)
query
K h poont W
(6)
key
Then compute the attention matrix (()) for the ith modality
(A() VDH
Sh Softmax
(7)
computational cost of O(N)DH DH 1111,))) N2).
Explicitly constructing the full attention matrix Shh using the Kronecker product (Eq. equation (Eq. equation formulation from the factor matrices sy² would negate the computational savings of the factorization. Instead, exploit properties of the Kronecker product and the associative law of matrix and tensor multiplication to apply the attention without forming Shh directly. Formally, it is easy to check that
3 or sum
Sh Sprod ((ro))( ((++a))
(s(1) 8... @ ((at (V(++1)) ((D* X1 5t-5) X2 ... x s(())
(8)
(k+1)
and
Sh sum (V (1+1))
((()) #.. ((.
1 (1"*x + +.. (yt Xx )(
(9)
(k+1)
We can, thus, multiply attention matrices one by one with the value tensor. The operation on each mode i yields computational complexity of O(,,Dn(II, N,)), resulting in an overall complexity of O(D (T, NNj)). For factorized attention layer with R heads of width Dy D/R, this leads to total complexity is O(D(E, NO),, N2)).
In addition to reducing computational and memory complexity, the proposed factorization of the high-order attention matrix provides a convenient framework for studying and customizing attention on each mode independently. This not only simplifies post-training interpretation of attention maps along each axis but also enables the use of custom attention masks tailored to specific modalities and applications. For instance, causal mask can be casily applied to temporal attention in autoregressive generative setting. However, in this paper, focus solely the use case of HOT Transformer encoder.
4.3 Architecture
Our final architecture is built on three main components (Fig.
A projection layer implemented as a convolution layer followed by a ReLU activation function that projects the input tensor to the hidden space. Sctting the stride equal the kernel size, this module divides the input into non-overlapping temporal spatial patches, reducing the input data resolution.

---

Published in Transactions on Machine Learning Research (11/2025)
2. The Transformer encoder layer consists of alternating layers of our proposed Kronecker (product or sum) self-attention and MLP blocks. Layernorm is applied before every block, and residual connections after every block. The MLP contains two layers with a GeLU non-linearity.
3. A global average pooling layer, followed by single linear layer as the projection head.
4.4 Theoretical Analysis
4.4.1 Attention Matrix Rank
Understanding the rank of attention matrices is key to analyzing Transformer expressivity and efficiency. Although we model high-order allention via low-rank factorization, this does not imply that the computed matrices Sprod and Stum from Eqs. 3 and 4 are themselves low-rank. Instead, we analyze their geometry and spectral spread as proxies for effective rank.
While attention matrices are theoretically full-rank, since softmax of random matrices typically yields linearly independent rows, their effective rank is often much lower due to row-stochasticity, dominance of few singular values from peaked attention distributions, and training dynamics encouraging low-rank structures. Moreover, computing the exact rank of an n x n attention matrix is computationally expensive at O(n²) and numerically unstable. Therefore, we adopt the stable rank as a robust measure, defined for any nonzero A E gmxa and its singular values (A) as:
|A E(A((A (A) = [AI3 si(A)
(10)
We analyze stable rank for the factorized attention matrices Sprod and Shh Shum as well the mode-wise matrices " For the Kronecker product, the stable rank factorizes: sr(S sr() x sr(S)
(11)
For the Kronocker sum formulation, the stable rank is
(I m,)l P + E., (1bz n3) tx((())
sr (r Shh
(12)
EE (s())) 2)
Before deriving bounds, we recall for any square A E Iuxx;
1 < | S n,
(13)
1 1 tr( S n,
(14)
0 < (A) S 1,
0< s1()) $
(15)
implying
1 S sr) < n.
(16)
Thus, the stable ranks satisfy
S sI (podd)
- II
No.
(17)
i=1
and
Σh (Isss Ny) + E/ (Tlsfi Mi)
S Sr sum
II N₁. 1=1
(18)
122
While sr(Sh shares the same bounds as the original non-factorized attention matrix, sr(S% has a higher lower bound, independent of the stable ranks of the modo-wise matrices s((i) This means that even if all s collapse to rank-one, Sh typically avoids collapsing, especially when the N, are large relative to k. In the next chapter, empirically compare stable ranks for the Kronecker product and sum attention matrices.
7

---

Published in Transactions on Machine Learning Research (11/2025)
4.4.2 Low Rank Factorization Approximation Guarantee
The summation across all heads appearing in Eq. equation 2 when using Kronecker product formulation from Eq. functions analogously to rank R Kronecker decomposition (when W is the same for all heads) and the factorization rank corresponds to the number of heads. Note that the factorization rank is different and independent from the attention matric rank (explained in the previous section) as the former a hyper- parameter and the latter is to be computed in practice. In Theorem 4.3 below we show that any attention matrix can be docomposed as sum of such Kronecker products. The following thcorem shows that rank R Kronecker decomposition is capable of approxinating any high-order atlention matrix arbitrarily well, as R increases, ensuring that significant interactions are missed. This theoretical aspect is crucial for ensuring that the attention mechanism can potentially adapt to any dataset or task requirements.
Theorem 4.3 (Universality of Kronecker product decompasition). Given any high-order allention malric S No. Nz)..Ne), there erists an R E N such that can be expressed as a rank R Kronecker de- composition, i.e., S E r & St2 @...& S( As R approaches minj=1..... Iug N2, the approzimation is guaranteed to become exact, meaning the Kronecker decomposition is capable of universally represenling any high-order attention matria S.
Proof. Proof is presented in the appendix A.1
While a sum of Kronecker products can express any tensor (given enough summands), is not the case for sums of Kronecker sums (which are still Kronecker sums). Still, our experiments showcase that the Kronecker decomposition
almost always match the performance of the Kronecker product decomposition, suggesting that the attention matrices needed to obtain good accuracy multi-modal data are highly structured. In the appendix, we have discussed the choice of Kronecker factorization in section A.2 and the inductive biases our method section A.3
5 Experiments
We thoroughly evaluate HOT on three high-order data tasks, validating the generality of the proposed framework. At each subsection, we introduce the lask, benchmark datasets, and baselines used, and discuss the performance results. Implementation details are presented in the appendix. We close the section by reviewing ablation studics that further confirm our thcory and design choices.
Table 1: Multivariate forecasting results with prediction lengths S E (96, 192, 336, 720} and fixed lookback length = 96 with the best in Bold and second-best best in underline. Results are averaged from all prediction lengths. FLOPS are calculated for a fixed input size (100, 96). Results reported from the original papers; others are reproduced by us.
ECL
Weather
Traffic
Solar
Models
Params
GFLOPS
MSE
MAE
MSE
MAE
MSE
MAE
MSE
MAE
AutoFormer!
15M 21M
0.227
0.338
0.338
0.382
0.628
0.379
0.885
0.711
FedFormer!
0.214
0.327
0.309
0.360
0.610
0.376
0.291
0.38
Croesformert
0.244
0.334
0.259
0.315
0.550
0.304
0.611
0.639
TimesNet!
301M
0.192
0.295
0.259
0.287
0.620
0.336
0.301
0.319
PatchTSTt
1.5M
0.205
0.290
0.259
0.281
0.481
0.304
0.270
0.307
Transformer (spatial)t
13M
1.03
0.178
0.270
0.258
0.278
0.428
0.282
0.233
0,262
Transformer (temporal)t
1.00
0.202
0.300
0.258
0.283
0.863
0.499
0.285
0.317
Transformer (non-factorized
620K
4.22
0.162
0.261
0.245
0.275
OOM
0OM
0.219
0.257
HOT (product)
680K
1.51
0.169
0.268
0.245
0.275
0.420
0.278
0.221
0.257
HOT (sum)
680K
1.51
0.167
0.266
0.245
0.275
0.428
0.280
0.220
0,260
8

---

Published in Transactions on Machine Learning Research (11/2025)
5.1 Long-range Time-series Forecasting
Given historical observations X = (x1 xx) E RTxN with T time steps and N variates, we predict the future S time steps Y (x.+1....,, XT+S} E RSSNN
Datasets We include four real-world datasets in our experiments, including ECL, Traffic, Weather used by Autoformer (Wu et al. 2021), and Solar-Energy proposed in LSTNet (Lai et al. 2017). Further dataset details are in the Appendix.
Baselincs We choose seven Transformer-basee modols as our baselines, inchuding Crossformer (Zhang & Yan 2023), Autoformer (Wu et al. 2021). FEDformer (Zhou et al. 2022), PatchTST (Nie et al. 2023). and three variations of the vanilla Transformer based on how processes timeseries data: Transformer (spatial) with attention applied on spatial axis (Liu al.) (2024) 2. Transformer (temporal) with attention applied on the temporal axis (Liu et al. 2024 3. Transformer (non-factorized) with full spatiotemporal atlention on flattened input corresponding to non-factorized high-order attention.
Results Table presents the lorecasting results, with the best and second-best performances highlighted. The non-factorized Transformer achieves the best results across all datasets except Traffic, where it runs out of memory. While this demonstrates its strong expressive power, it comes at a significant computational cost due to the quadratic complexity of attention with respect to the product of data dimensions. Our proposed HOT models offer a more efficient alternative. HOT (product) matches the performance of the non-factorized Transformer on most datasets while requiring only third of the dooting-poiit operations. HOT (sum) consistently ranks second, closely following HOT (product). Both models also use significantly fewer parameters, making them computationally efficient. Among the baselines, the spatial Transformer (i.e., iTransformer (Liu et al. 2024) outperforms the temporal Transformer and several methods like Autoformer, FedFormer, and Patch IST, showcasing the importance of capturing spatial correlations in high-dimensional data.
Table 3D image classification results on MedMNIST3D with the best in Bold and second- best in underline, FLOPS calculated for single input image. RResults reported from the original papers; others reproduced by us.
Organ
Nodule
Fracture
Adrenal
Vessel
Models
Params
GFLOPS
AUC
ACC
AUC
ACC
AUC
ACC
AUC
ACC
AUC
ACC
ResNet-18t
12M
2.32
99.4
90.0
86.3
84.4
71.2
50.8
82.7
72.1
82.0
74.5
ResNet-50t
26M
3.19
99.4
88.9
88.6
84.1
75.0
51.7
82.8
75.8
91.2
85.8
MDANet
7M
98.9
89.7
86.8
86.00
83.9
81.5
90.1
92.9
CdTransfrrmer
88.6
71.6
51.7
87.7
81.5
89.3
ViT-3Dt
130M
91.4
86.7
64.8
82.0
82.6
ViViT-st
50M
86.6
65.5
53.8
81.0
79.9
83.9
88.7
TimeSformer
0.71
90.4
90.6
89.8
68.0
52.3
81.7
82.5
85.3
91.3
MVIT
1.6
0.72
99.0
88.1
90.1
86.8
66.7
53.7
82.5
80.1
80.6
89.2
Transforner (actoric))))
1.5M
0.72
99.7
92.6
91.6
87.8
74.4
58.6
88.8
83.9
92.0
91.1
HOT (product)
2M
0.67
99.8
94.4
92.0
89.1
73.6
58.3
88.6
83.8
85.7
91.1
HOT (sum)
2M
0.67
99.6
92.2
90.7
87.7
74.9
58.7
87.3
82.9
84.9
90.7
5.2 3D Medical Image Classification
Given a 3D image X E H with width W, height H, and depth D, we predict the image class probability y E RC over a set C' classes.
Dataset MedMNIST v2 Yang et al. 2023) is a large-scale benchmark for medical image classification on standardized MNIST-like 2D and 3D images with diverse modalitics, datasct scales, and tasks. We primarily experiment on the 3D portion of MedMNIST v2, namely the Organ, Nodule, Fracture, Adrenal, and Vessel datasets. The size of each image 28 x 28 x 28 (3D).

---

Published in Transactions on Machine Learning Research (11/2025)
Baselines We choose seven medical image classifier models, including ResNet-18 ResNet-500 (He et al,
2015 Yang et al. 2023), MDANet (Huang et al.) 2022), CaTraanformer (Zhu et al. 2024). VrT-3D and
ViViT-S (Lai et al. 2024), TimeSFormer (Bertasius et al. 2021), MVTT (Fan et al. 2021 and vanilla Transformer with full attention over all three axes applied on flattened input as our baselines.
The classification results on MedMNIST3D summarized in Table 2 with the best and second- best performances highlighted. HOT (product) achieves state-of-the-art results on most datasets, with the highest ACC and accuracy for Organ, Nodule, and Adrenal datasets, while requiring minimal computational resources. HOT (sum) closely follows, demonstrating competitive performance, particularly on the Fracture dataset, where achieves the best accuracy. Both HOT variants utilize only 2M paramelers and maintain low FLOPS count, showcasing their efficiency. The full Transformer achieves strong results, often ranking second matching HOT (product), especially on Adrenal and Vessel datascts. Among other transformer baselines, TimeSFormer and MVIT also deliver competilive results with very low FLOPS, with TimeeForrer reaching second-best. AUC Vessel and MVTT showing balanced performance across Nodule and Fracture, though both remain behind HOT overall. Additionally, ResNet-50 performs well on certain datasets, par- ticularly Fracture, but lags in computational efficiency. CdTransformer and ViT-3D demonstrate strong performance on specific datasets (e.g., Nodule and Vessel for CdTrransfrrmer), but their parameter-heavy designs make them less efficient compared HOT. Overall, HOT models strike the best balance between accuracy and efficiency across diverse 3D image classification tasks.
Results
Table 3: Multispectral image segmentation results on SSL4EO-L Benchmark dataset (using OLI/TIRS-H -S-- scnsor product) with the best in Bold and sccond-best in underline. *Results reported from the original papers; others are reproduced by us.
Models
Pretraining
Params
GFLOPS
CDL
NLCD
ResNet18f ResNet50t
11.2M
3.42
68.0
67.0
MoCo
23.5M
7.12
65.9
67.4
VIT-S16
22.5M
5.07
64.1
66.8
HOT (product)
4.3M
3.80
66.2
HOT (sum)
4.3M
3.89
64.5
67.2
5.3 Multi- Speti-Spectral Image Segmentation
Given a multisepetral satellite image X' E WxxS with width W, height. H, and spectral channels S, we predict the pixel-wise class probability y E XXXX over a set of C classes.
Dataset The SSL4EO-I. Benchmark dataset (tewarrt et al. 2023 is collection of Landsat images paired with land cover classification masks. The images are captured using different sensors; for this experiment, we focus only on images from the OLI/TIRS-EE sensor product. The dataset is available in two versions: NLCD with 17 classes and CDL with 134 classes. The size of each image is 264 x 264 7.
Baselines We compare HOT variants with VIT and ResNet results reported in the benchmark paper (Stewart et al. 2023). We only report models pretrained with MoCo (He et al. 2019) and subsequently fine- tuned on the SSL4EO-L Benchmark, as these consistently yield the best performance across all backbones.
Results Table reports multispectral segmentation results on SSL4EO-L. Among MoCooprrerriinct base- lines, ResNet18 (11.2M params, 3.42 GFLOPS) achieves the best CDL score (68.0), while ResNet50 (23.5M params, 7.12 GFLOPS) attains the sccond-best NLCD score (67.4), suggesting that convolutional backbones remain strong for high-resolution image segmentation. ViT-S16, despite a similar parameter count (22.5M, 5.07 GFLOPS), underperforms both ResNet variants. Our proposed HOT models, trained from scratch with only 1M parameters and 3.89 GFLOPS, outperform ViT on both benchmarks, achieve the best NLCD re- sult (68.6), and remain competitive with ResNet18 on CDL (66.2), highlighting their efficiency in balancing accuracy and computational cost.
10

---

Published in Transactions on Machine Learning Research (11/2025)
HOT (sum) On Weather Dataset
HOT (preduct) On Weather Datbott
HOT touna On Organ Dataset
HOT (producti On Organ Dataset
2.0
3.00
4 M
1.23
2.75
Asctpp Stdde falk
2.30
1.20
1.20
Made
Nede
1/-
1.6
1.15
1.5
2.00
4m
133
1.05
1.00
12
25
25
1.00
125
0
1000
125
50
200
50
1.00
200
Figure Evolution of the average stable rank for mode-wise attention matrices across training steps for HOT (product) and HOT (sum) models.
5.4
Ablation Study
Mholfer Sulauet
Organ Dolaaet
5.4.1 Attention Rank
We examine the average stable rank of ion ma- trices across heads and layers, including mode-wise matrices xh, factorized matrices 8h xrodd, Stmm. and the full altention from the Transformer base- line. On the Weather dataset (). HOT (prod- uct) shows higher mode-wise stable ranks than HOT (sum), with consistently lower ranks in the time mode compared to the variate mode mode-eespite simi- lar dimensions-sigeeeee suggesting richer spatial dependen- cies. Stable ranks plateau early and remain steady during training. For the Organ dataset, stable ranks are similar across all three spatial modes in both models, reflecting uniform spatial structure in the 28 x 28 x 28 volumes. Ranks stay stable throughout training, with slight increase validation. Training ranks are generally higher due to dropout-induced variability. Looking at full attention matrices (Fig 4. HOT (sum) has consistently higher stable ranks than HOT (product), consistent with thcory. Both maintain steady ranks during training. The non-factorized Transformer, however, exhibits an early spike in stable rank before converging near HOT (product) levels a trend also seen on the Organ dataset and deserving further analysis.
120
Figure 4: Evolution of the average stable rank for full attention matrices across training steps. Results are shown for the Transformer fuct-iact,,ized), HOT (sum), and HOT (product) models.
5.4.2 Factorization Rank
We analyze how the number of attention heads (i.e., factorization rank) influences HOT's perfor- mance on medical imaging and time series datasets.
The factorization rank critical for approximat- ing high-order attention and capturing diverse cross- dimensional patterns. Following common practice, treat the rank (i.e., number of attention heads) as hyperparameter and determine the best value through hyperparameter search. Figure 5 shows that increasing the rank generally improves per- formance, with optimal results between 8 and 16 heads. Beyond this range, higher ranks degrade per- formance and increase errors, likely due to overfit-
Figure 5: Effect of increasing the number of heads (i.e.. factorization rank) on model performance with Kro- necker product, Kronecker sum, and non-factorized attention. Left: Multivariate time series datasets. Right: 3D medical imaging datasets.
11

---

Published in Transactions on Machine Learning Research (11/2025)
ting and reduced head dimension d2 D/h, which eventually becomes too small to model interactions effectively. Optimal ranks vary by dataset: HOT (product) and HOT (sum) perform best at rank 8 on Traffic and Organ, while rank 16 yields superior results on Nodule and Weather.
5.4.3 Efficiency Analysis
We evaluate the computational efficiency of HOT against. ViT, ResNet50, and ResNet18 across training time, inference time, and GPU memory foolprint (Figure 7 HOT requires 27.4 per sample for training higher than ResNet18 (11.0 ms) and ResNet50 (22.0 ms), but lower than VIT (34.2 and and achieves 9.8 ms per sample for inference, placing it between R8 (3.4 ms) and ViT (11.5 ms), while ResNet50 is slightly faster at 7.2 ms. Notably, HOT is the most memory- efficient, using only 0.21 GB of GPU memory, significantly less than VIT (0.74 GB) and even ResNet18 (0.39 GB). We fur- ther analyze HOT's Kronecker attention by comparing FLOP requirements against Full (non-factorized) and Divided (from [Bertasius et al., 2021]) attention mechanisms (Figure 6 showing that for an input of size 2243, Kronecker attention requires only 371.6 GFLOPS, versus 775.3 GFLOPS for Full attention and 1210 GFLOPS for Divided attention, with the efficiency gap consistent across smaller inputs. Overall, HOT maintains competitive training and inference times, minimizes memory usage, and provides substantial computational savings at scale, making it strong candidate for resource- eesource-constrainee medical Al applications. All experiments were implemented in PyTorch and conducted on a single NVIDIA A100 GPU (80 GB) with an x86 64 CPU (6 cores) and 64 GB of RAM.
12-
Attetton Type
320 1255
300
175
Irpult Siz Oemaraie C20
Figure 6: Comparison of computational cost. measured in GFLOPS versus 3D in- put dimension for different attention mech- anisms. The Kronecker, Full, and Divided attention types exhibit distinct scaling be- haviors as the input dimension increases. with the Kronecker attention having signif- icantly lower computational cost.
Training Imsisample)
Interence (msisampple)
GPU Merrory
35
34.22
11.43
0:74
0.7
3.60
27/55
0.6
0.11
7.23
20
0.59
15
10
1191
2.4
0.2
921
Rodletta
Resiedd
10T
Restatto
Rededdd
NOT
Rusterta
Reffedd
Figure 7: Comparison of model efficiency across four architectures: HOT, VIT, ResNet50, and ResNet18. The figure shows three metrics: (a) training time per sample (ms), (b) inference time per sample (ms), and (c) GPU memory footprint (GB). HOT exhibits the highest training time but low GPU memory usage, while VIT and ResNet variants offer a better trade-off between speed and memory consumption.
5.5 Visualization of Attention Maps
We visualize the spatial and temporal attention maps learned by the HOT model on the ECL dataset for time series forecasting. To reveal these patterns, we compute average attention weights across batch of
12

---

Published in Transactions on Machine Learning Research (11/2025)
test samples for each head and layer. Overall, observed no significant differences between the patterns learned by the HOT (Product) and HOT (Sum) models. However, clear distinctions emerge between the temporal attention maps and those along the spatial (variable) axis.
Spatial maps (Fig exhibit two main patterns: clustered regions and sparse vertical stripes. The clustered patterns indicate strong dependencies among groups of variables, suggesting the model identifies meaningful partitions in the data. In contrast, the sparse vertical stripes highlight specific variables that exert significant inffuence others, reflecting their importance in prediction.
Temporal maps (Fig ! reveal two characteristic behaviors. First, many attention heads display diagonal line structures, capturing periodic temporal dependencies; the distance of these lines from the main diagonal corresponds to the length of the learned period. Additionally, some allention maps, especially in later layers, become sparse, focusing attention on key time steps such as the beginning or end of the sequence. Complete attention visualizations and additional ablations are provided in the appendix D
Layer 1, Head 4
Layer 1, Head
Layer 2, Hoad
Layer 1. Head 1
20
56
10-
112
176
112
140
15
160
188
230
198
210
152
256
280
02 RR38
Figure Visualization of the spatial attention maps learned by HOT
Layer I, Hoad
Layer Head 2
Layer Head
Layer 2. Head 3
2
00
91
24
54
2-
R
R
34
10
10
15
Figure 9: Visualization of the temporal attention maps learned by HOT
6 Conclusion
We addressed the challenge of applying Transformers to high-dimensional, multiway tensor data, where standard attention mechanisms are computationally expensive and misaligned with tensor structures. To overcome these limitations, we introduced the Higher-Order Transformer (HOT), which uses Kronecker fac- torization to express multiway attention sums of mode-wise attention matrices. This approach efficiently models dependencies across dimensions while preserving tensor structure. Our analysis shows HOT retains the expressiveness of full high-order attention with adjustable complexity through a rank paramcter. Experi- ments on diverse lasks, including multivariate time series forecasting, 3D medical imaging, and multispectral image segmentation demonstrate HOT's strong performance and scalability, while attention visualizations highlight its ability to learn mcaningful, interpretable patterns. Future work includes extending HOT generative modeling, applying it to large-scale image and video datasets, exploring its transfer learning po- tential across modalities, and integrating it with other tensor decompositions to further enhance efficiency
13

---

Published in Transactions on Machine Learning Research (11/2025)
and modeling capabilities. HOT offers promising foundation for advancing multiway data modeling across diverse domains.
Broader impact statement This work introduces Higher- Higher-Order Transformers (HOT) as an efficient and interpretable framework for modeling multiway data, with experiments on 3D medical imaging benchmarks. Potential benefits include faster and more accurate image analysis, lower computational costs that broaden access, and improved transparency through mode-wise attention maps. However, risks include bias from non-diverse datasets, limited generalization, privacy concerns, and automation bias in clinical use. Since ovaluations limited to rescarch benchmarks, real-world performance and interpretability remain uncertain. Future work should focus on fairness audits, robust validation on diverse clinical data, uncertainty quantification, and human-ing see-in-ther to------loppp deployment to ensure safe and ethical use.
Acknowledgments G. Rabusseau and R. Rabbany acknowledges the support of the CIFAR Al chair program. This work made use of compute resources by the Digital Research Alliance of Canada.
References
Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lucic, and Cordelia Schmid. Vivit: A video vision transformer. CoRR, abs/2103. 2021. URL
15691
Djalel Bahri, Bertrand Pasdeloup, Vincent Gripon, and Dominique Pastor. Robust kronecker component analysis for tensor dictionary lcarning. IEEE Transactions on Signal Processing, 67(24):6400-6420, 2019. 2019.. doi: N0.1109/TSP:2019 2993399..
Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-time attention all you need for video under- standing? CoRR, abs/2102.05095, 2021. URL https://arxiv..rrg http://ogg///bb//ee/eeee 2502/
Rewon Child, Scott Gray, Alec Radford, and Πya Sutskever. Generating long sequences with sparse trans- formers. Xie preprint arxi::194. 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. International Conference on Leurning Representations, 2021.
Linhao Dong, Shuang Xu, and Bo Xu. Speech-trannsfommer: A no-recurrence sequence- to-enncc-t---uuuncee model for speech recognition. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5888, 2018. doi: N0.1109/TCASSP:200 2018.64650666
A Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. CoRR, abs/2010. 1192. 2020. URL attts::/arrivv.or//ab//////
Haoqi Pan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. CoRR, abs/ Abs/2104.11227, 2021. URL https://arxiv.org org/ abs/2104.1n 11227
Evgeny Frolov and Ivan V. Oseledets. Tensor methods and recommender systems. CoRR, abs/1603.06033, 2016. URL httpp://axxvvvorr//ee
06038
Hongyang Gao, Zhengyang Wang, and Shuiwang Ji. Kronecker attention networks. CoRR, abs/2007 200.00007,8 2020. URL htttps://axxiv.rrgg wwtpp////oroaaa//////////// 08442
Kristjan Greenewald and Alfred O Hero. Robust kronecker product pea for spatio-temporal covariance estimation. IEEE Transactions on Signal Processing, 63((3):6368 6378, 2015.
Ying Guo, Hongtu Zhu, Yi Li, Min Chen, Martin Styner, and Weili Lin. Tensor regression models for longitudinal imaging data. Statistica Sinica, 23(3):1401-1433 2012.
14

---

Published in Transactions on Machine Learning Research (11/2025)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition.
CoRR, fbs/1511.03365, 2015. URL http::://aaxii.og/aaabaa..
03385
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised
visual representation learning. CoRR, abbs1911.0550 2019. URL httt:://aaxxiv..ogg/aa//////////
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention multidimensional
transformers. CoRR, abs/ 19ss1992.,2180, 2019. URL http:///arxi..oog///////1111111.
Yixiang Huang, Lifu Zhang, and Ruoxi Song. Semantic attention guided multi-dimennionn information
complementary network for medical image classification, 2022. URL http://///oo..0111..211.1. 3.r8-2205040v11
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers rmns: Fast
auloregressive transformers with linear attention. arXiv preprinl arXiv:2006. 16236, 2020.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization, 2017. URL https
//arxiv. org/ www.o///////.//////////// 6980
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Rev., 51:455-500
2009a. URL
Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM revieu, 5133):555-500 2009b.
Guokun Lai, Wei-Cheng Chang, Yiming Yang, and Hanxiao Liu. Modcling long- and short-term temporal patterns with deep neural networks. CoRR, abs/ 1793100000705, 2017. URL http://arxiv http:///axxxv.oorraa///1100000 07015
Zhixin Lai, Jing Wu, Suiyao Chen, Yucheng Zhou, and Naira Hovakimyan. Residual-based language models are free boosters for biomedical imaging, 2024. URL /ttpps://axxvv.rrtt httbb///443..1// 17343
J. M. Landsberg. Tensors: Geometry and Applications, volume 128 of Graduate Studies in Mathematics. American Mathematical Society, 2012. ISBN 978-0-8218-6907-9.
Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition. In International Conference on Learning Representations (ICLR), 2015.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, (71533)):4336 444, 2015. doi: 10.1038/natto 10.1038/nature1455 2tt8//aauur14439.
Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, and Mingsheng Long. itransformer: Inverted transformers are effective for time series forecasting. In The Twelfth International Conference on Learning Representations, 2024. URL
Haiping Lu, Konstantinos N. Plataniotis, and Anastasios N. Venetsanopoulos. Multilinear Subspace Learning: Dimensionalily Reduction of Multidimenssinal Dala. CRC Press, 2003.
Zifeng Lu, Konstantinos N Plataniotis, and Anastasios N Venetsanopoulos. A survey of tensor methods for multiway data analysis. arXiv preprint orXiv:1107.2849, 2011.
Xindian Ma, Peng Zhang, Shuai Zhang, Nan Duan, Yuexian Hou, Ming Zhou, and Dawei Song. A ten- sorized transformer for language modeling. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Accé-Buuc, E. Fox, and R. Garnett (ods.), Advances in Neural Information Processing Systems, volume 32, Cur- ran Associates, Inc., 2019. URL https://prroeeeding: www.odiinnnn.euiis.... hoturi/p..c/ppppr..eee////pp/p/p C.50445scsaats5555c7ddedddddd--2
15

---

Published in Transactions on Machine Learning Research (11/2025)
Simone Montangero. Introduction to Tensor Network Methods. Springer International Publishing, Cham, 2018. ISBN 978-3-033-011408-- doi: 10.1007/976--3330-0044443 URL http:///iinnk.ssinnerr.. com/ 0000//10..00//978---0000-0010090
Tung Nguyen, Jason Jewik, Hritik Bansal, Prakhar Sharma, and Aditya Grover. Climatelearn: Benchmarking machine learning for weather and climate modeling, 2023. URL https://arxii:
wtt. htttpe//aarxii.orr//bb//0001 01909
Yuqi Nie, Nam H. Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series worth 64 words: Long-term forecasting with transformers. In International Conference on Learning Representations, 2023.
Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov. Tensorizing neural networks. In Advances in Newral Information Processing Systems (NewrIPS), pp. 442 450, 2015.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphaäl Lopez Kaufman, Aidan Clark, Seb Noury, Matthew Botvinick, Nicolas Hoess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. In Hal Daumé III and Aarti Singh (ods.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Reseurch, pp. 7487-777777 7498. PMLR, 13 18 Jul 2020. URL http:::///pproociingg... mlr.
prees//11//prrisoto2aa. html
Shehan Perera, Pouyan Navard, and Alper Yilmaz. Segformer3d: an efficient transformer for 3d medical image segmentation, 2024. CRL httth:://aartt https/www..rg/b////4...../
Dogyoon Song, Zeyu Sun, and Alfred O Hero. On separability of covariance in multiway data analysis. arXiv preprint 2rXte:2301.00001, 2023.
Sijie Song, Cuiling Lan, Junliang Xing, Wenjun Zeng, and Jiaying Liu. An end-to-end spatio-temporal attention model for human action recognition from skeleton data. CoRR, abs/1611 abs/1611.06067, 2016. URL
http://arxiv http://aaxivoooo/aaaaae wwwaaa/ae......aaeeeee..e...e..e 06067
Adam J. Stewart, Nils Lehmann, Isaac A. Corley, Yi Wang, Yi-Chia Chang, Nassim Ait Ali Braham, Shradha Sehgal, Caleb Robinson, and Arindam Banerjee. Ss14co-1: Datasets and foundation models for landsat imagery, 2023. URL htttps///arxxt wwt. htrg//bs/233000 09424
Jianlin Su, Yu Lu,, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, 09942404.,9864, 2021. URL httph:ttt//aitt wttp/ab////....00000000 09864
T Tsiligkaridis and A O Hero. Covariance estimation in high dimensions via kronecker product expansions. IEEE Transactions on Signal Processing, (21):53337 2013.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukası Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st International Conference Newral Information Processing Systems, NIPS'17, pp. 6000 6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Yu Wang, Zeyu Sun, Dogyoon Song, and Alfred Hero. Kronecker-structurre covariance models for multiway data, 2022. URL
01721
Haixu Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. Autoformer: Decomposition transformers with auto-correlation for long-term scries forecasting. CoRR, abs/2106.13008, 2021. URL https://arxiv.
org/abs/2106. 13008
Haixu Wu, Tengge Hu, Yong Liu, Hang Zhou, Jianmin Wang, and Mingsheng Long. Timesnet: Temporal 2d-variation modcling for general time series analysis, 2023. URL httts://arxiv..or org/ harg/abs/200 aas//bbb/222.... 0bss/220000288
Jianeheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspoter Pfister, and Bingbing Ni. Medmnist v2-a largo-scale lightweight benchmark for 2d and 3d biomedical image classification. Scientificc Dala, 10(1):41, 2023.
16

---

Published in Transactions on Machine Learning Research (11/2025)
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Outanon, Philip Pham, Anirudh Ravula, Qifan Wang, and Li Yang. Longformer: The long-document transformer. arXiv preprint Axte:2004.05150 2020.
Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, and Andrew Chi-Chih Yao. Tensor product attention is all you need. arXiv preprint arXi::2001.04425, 2025.
Yunhao Zhang and Junchi Yan. Crossformer: Transformer utilizing cross-dimension dependency for mul- tivariate time series forocasting. In The Eleventh International Conference on Learning Representations, 2023. URL htttps://openrevieer http:p///ppaaaeeeee...ae/////s// wtt.:/ooruuuiiiecesss//aaae...
Tian Zhou, Ziqing Ma, Qingsong Wen, Xue Wang, Liang Sun, and Rong Jin. Fedformer: Frequency enhanced
decomposed transformer for long-teray series forecasting. CoRR, abs/2201. 2022. URL https:
1/azzv.org/abs//201.112
Qikui Zhu, Chuan Fu, and Shuo Li. Class-consistent contrastive learning driven croes-dimensional trans- former for 3d medical image classification. In Kate Larson (ed.), Proceedings of the Thirty-Third In- ternational Joint Conference on Artificial Intelligence, IJCAI-24, pp. 1807-1815. International Joint Conferences on Artificial Intelligence Organization, 8 2024. doi: /144966//cai..224//000 URL https: r/dii.org/10.24963//jaaaaa 2024/200 Main Track.
17

---

Published in Transactions on Machine Learning Research (11/2025)
A Theoretical Analysis
A.1 Proofs
Theorem A.1 (Row stochastic property). Let E RNNN and T E RRMM be row-stochastic matrices (L.e. Σj Sig =1 for all i, and likewise Ee The = 1 for all k). Then S S E RNMXNM and (S E R MMMMM are also row-stochastic.
Proof. For the Kronecker product, recall by definition
(S & (..... Siz The
......,,.... for i,j = = and k; l ,,..,,,,,,,........ M.
Fix an output-row index (2, k). Then
N M
M
Es (TT)) =
(Sy The) =
Siy
Σ 1=1
The) = =1x1=1
(.
Since this holds for every row ki, S T row-stochastic.
For the Kronecker sum, by definition S STT= = S IM + IN ST. Again, fix a row index (i, k).. We have:
Σ [(S & Ja))u....G.) + (IN @ TuNro. = Σ (S IND + Σ @ (1..).) (19
(7.
(3))
while (S JM) = Siz oke, and (IN @ T)(A,)0 dig The Hence,
Es,
Sig dee Σ
Sis
1 x 1=1, ōkt
j.
Em
and similarly
N
M
Σ 3.2
8ig The =
8Gs)
Σ
Tke = x 1=1.
j=1
Thus,
Σ 1 5 (S T) T) (1 (1+1)=1
(.
Corollary A.2. Let (S() E NRN.N. j be collection of row-stochastic matrices. Then S(1) @ 5(2) .... 8(k) and T (5(1 e ( e S(A) are also row-stochaslic.
Proof. The Kronecker product associative, and from the theorem, we know that the Kronecker product of two row-stochastic matrices row-stochastic. Now, we proceed by induction. The base case k is given by the theorem. Assume that the Kronecker product of k row-stochastic matrices is row-stochastic. Then (8(1) @ ... @ 5(+-1)) S(k) is the Kronecker product of row-slochastic matrix with ((k), which by the base theorem is again row-stochastic.
By definition, the Kronecker sum of matrices is given by a sum of terms, cach of the form:
1 = IN, & ... & IN_, & s(i) & IN 8 ... & IN
Each term 1 is a Kroneeker product of identity matrices and one row-stochastic matrix s(0, so by the product result above, each A is row-stochastic. Summing & such matrices yields matrix whose rows sum to k, so dividing by l ensures that the normalized sum D A00 is row-stochastic.
18

---

Published in Transactions on Machine Learning Research (11/2025)
Theorem A.3 (Universality of Kronecker product decomposition). Given any high-order attention matrix S' E NL) Ne), there erisls an R E N such that S can be expressed as rank R Kronecker de- composition, i.e., S - S, &.S &.. As R approaches ming=1,".. . I1 Ni, the upprozimation is guaranteed to become ezact, meaning the Kronecker decomposition is capable of universally representing any high-order attention matriz S.
Proof. Let 8 be the tensor obtained by reshaping the attention matrix S into tensor of size N1 N₂ x
x N& x N1 X N₂ x ................................ N& and let T E x X be the tensor obtained by merging cach pair of
modes corresponding to one modality Let R be the CP rank of T and let T T= =E== R (1) (2)
8r 8 o... (R)
be a CP decomposition, where o denotes the outer product and each 5 E R for i 1, ... A (see, e.g. (Kolda& & Bader 2009a) for an introduction the CP decomposition). By reshaping each of E R. into matrix St RN: RNNN,, one can check that
R Di @
Sa... Aedi Jx = (S
2 ... ....0 (5(
r=1
from which it follows that S E S( @ S(2) 8 ... @ SK as desired.
The second part of the theorem comes from the fact that min;=1,-. Dist dj is a well known upper bound on the CP rank of tensor of shape dj d2 x... d2 (see again (Kolda & Bader 2209aa))
A.2 Discussion on the Choice of Factorization
We consider various tensor factorization methods to decompose the high-order attention matrix:
CP Decomposition: Based on the given proof, the CP decomposition of the attention matrix reshaped into a tensor of size N? x N3 x matrix used HOT.
x Nº is equivalent to the Kronecker factorization of the attention
Tucker Decomposition: The attention matrix could be expressed as rank R Tucker decomposition as:
n Σ fL
Sattention
G
((1) @
(2) ... 2 (K)
(19)
where s(g is the first-order attention matrix for axis and G E R is core tensor that enhances exprossivity. However, introducing the core tensor, would come at a cost of O(R*) more paramoters and a computational complexity of O(DR* NE IIN)) growing exponentially with R, strongly limiting its scalability and efficiency.
Tensor Train: For tensor train decomposition with a fixed rank for all factor matrices, the attention matrix is expressed as:
R Σ f1.8g....iu-1
Satteution
((1)
St2) 11 @
5(3) x
@ s(() 1
(20)
where s(1) E R for i (1. and SOO E RRXRXN,N, otherwise. Constructing fourth- order factor tensors s(i) E RRXRXN,N, is non-trivial and may require additional parameters or modifications to align with our use case. This lack of simplicity makes tensor train dccomposition less appealing design choice for our use case.
The choice of Kronecker decomposition our work motivated by its
liee in pytorch would be obtained obtained by permuting modes S and reshaping: tarch.transpose(S, [0, A++1, 1, 2, 1).reshape([N2 .N2)
19

---

Published in Transactions on Machine Learning Research (11/2025)
1. Simplicity: Avoiding additional parameters, such as core tensors, making it easy to implement and interpret.
Efficiency: Scales efficiently with the tensor order, avoiding exponential growth parameters or computations.
3. Expressivity and Performance: Comes with theoretical guarantee, while empirically performing well.
A.3 Inductive Biases
The Higher-Order Transformer (HOT) introduces several key inductive biases that enhance its ability to process multiway data efficiently:
Preservation of Tensor Structure: Unlike couventional Transformer models that require flat- tening or reshaping of input tensors into sequences, HOT inherently preserves the original multi- dimensional structure of the data. This is crucial for applications such images, videos, and volumetric medical scans, where 2D or 3D neighborhood structures encode critical spatial and tem- poral relationships. By maintaining the tensor form, HOT enables direct modeling of multiway interactions without losing structural dependencies.
Locality and Spatial Awareness: In many real-world datasets, relationships between neighboring elements in a tensor are often stronger than those between distant elements. This property, known as locality, is particularly evident in vision tasks, where adjacent pixels in image or neighboring frames in video exhibit. high correlation. While self-attention mechanisms are inherently global, HOT retains weak locality by leveraging positional encodings, which helps the model differentiate between nearby and distant elements in the lemsor space.
Global Interaction via Self-Attention: Although locality is an important inductive bias, captur- ing long-range dependencies is equally crucial in high-dimensional data. Self-attention allows HOT to model global interactions across all tensor modes, enabling it to capture complex dependencies beyond local neighborhoods. This capability is particularly beneficial in tasks such as multivariate lime-series forecasting, where dependencies between distant time steps play a key role in prediction accuracy.
Separable Structure through Kronecker Factorization: A core inductive bias in HOT is the assumption that interactions across different tensor modes can be decomposed into separable components. This is achieved through Kronecker factorization, which models multiway dependencies in a structured and efficient manner. By assuming a factored representation of attention matrices, HOT significantly reduces computational complexity while still capturing meaningful relationships across dimensions.
B Datasets Details
B.1 Longe Time-series Forecasting
We evaluate the performance of the proposed HOT model on five real-world datasets: Weather (Wu et al. 2021), consisting of 21 metcorological variables rocorded every 10 minutes in 2020 at the Max Planck Bio- geochemistry Institute, ECL (Wu al., 2021), which tracks hourly electricity consumption for 321 clients, Traffic (Wu et al. 20210, collecting hourly road occupancy data from 862 sensors San Francisco Bay area freeways between Jamuary 2015 and December 2016, and Solar-Energy (Lai et al.] 2017), recording solar power production from 137 photovoltaic (PV) plants, sampled every 10 minutes in 2006.
split protocol used in TimesNet (Wu et al. 2023 ensuring datasets are chronologically split to prevent any data leakage. For forecasting tasks, we use fixed lookback window of 96 time steps for the Weather, ECL, Solar-Energe and Traffic datasets, with prediction lengths of 96, 192, 336, 720. Further dataset details are presented in Table 4
We follow the data processing and train-valiee wrain-validttion-eee
20

---

Published in Transactions on Machine Learning Research (11/2025)
Table Timeseries forecasting dataset details.
Dataset
Variables
Prediction Length
Train/Val/Tess Size
Sample Frequency
Weather
21
(96, 192, 336, 720]
(36792, 5271, 10540)
10min
ECL
321
(96, 192, 336, 720]
(18317, 2633, 5261)
Hourly
Traffic
862
(96, 192, 336, 720]
(12185, 1757, 3509)
Hourly
Solar-Energe
137 (96, 192, 336, 720)
(36601, 5161, 10417)
10min
B.2 3D Medical Image Classification
We conduct experiments on the 3D subset of the Medical MNIST datasel Yang et al. 2023 All dalasets have an image size of 28x28 28 voxels, allowing for consistent 3D image classification across different medical domains. The images come from various sources; ranging from human CT scans to animal microscopy, and have been adapted create challenging classification tasks. Details are presented Table
OrganMNIST33D is based the same CT scan data used for the Organ A, C, datasets, but instead of 2D projections, it directly uses the 3D bounding boxes of 11 different body organs. The dataset is adapted for multiclass classification on organ identification from volumetric medical data.
NoduleMNIST3D originates from the LIDC-IDRI dataset, public repository of thoracic CT scans designed for lung nodule segmentation and malignancy classification. For this study, the dataset has been adapted for binary classification of lung nodules based on malignancy levels, excluding cases with indeterminate malignancy. The images are center-cropped and spatially normalized to retain a consistent voxel spacing.
AdrenaIMNIST3D features 3D shape masks of adrenal glands collected from patients Zhongshan Hospital, Fudan University. Each shape is manually annotated by expert endocrinologist using CT scans, though the original scans are not included in the dataset to protect patient privacy. Instead, the dataset focuses on binary classification of normal versus abnormal adrenal glands based on the processed 3D shapes derived from the scans.
FractureMNIST33D is derived from the RibFrac dataset, which contains CT scans of rib fractures. The dataset classifies rib fractures into three categories (buckle, nondisplaced, and displaced), omit- ting segmental fractures due to the resolution of the images.
VesscIMNIST3D uses data from the IntrA dataset, which includes 3D models of intracranial aneurysms and healthy brain vessels reconstructed from magnetic resonance angiography (MRA) images. The dataset focuses classifying healthy vessel segments versus aneurysms, with the models voxelized into 3D volumes.
Table 5: MedMNIST 3D datasets details.
Dataset
Modality
Number of Classes
Train/Val/Test Size
Organ
Abdominal CT
11
(972, 161, 610)
Nodule
Chest CT
2
(1158, 165, 310)
Adrenal
Shape from Abdominal CT
2
(1188, 98, 298)
Fracture
Chest CT
3
(1027, 103, 240)
Vessel
Shape from Brain MRI
2
(1335, 192, 382)
21

---

Published in Transactions on Machine Learning Research (11/2025)
B.3 Multispectral Image Segmentation
We evaluate our models on the SSL4EO-I, Benchmark dataset. (Stewarl et al. 2023], large-scale collection of multispecttal Landsat images with corresponding land cover classification masks. For this work, we focus exclusively on images capturod by the OLI/TIRS-RS sensor to ensure consistent spectral charactcristics.
Each image has a resolution of 264 x 264 pixels (approximately 7.92 km 7.92 km at 30 m/px) with spectral bands, encompassing visible and infrared channels. The dataset includes 25,000 labeled images divided into training (20,000), validation (2,500), and test (2,500) splits, ensuring balanced data distribution.
The benchmark provides two versions of land cover annotations: the NLCD version with 17 broad land cover classes, and the CDL version with 134 finer-grained classes. Images are preprocossod to retain all multispectral channels, facilitating pixel-wise prediction of land cover classes in segmentation tasks. This setup allows us to evaluate HOT and baseline models a controlled, consistent, and realistic multispectral segmentation setting.
B.4 Implementation Details
Table 6: Hyperparameter Search Space.
Hyperparameter
Value List
Number of Blocks (timeseries forecasting)
[]
Number of Blocks (3D image classification)
[6]
Number of Blocks (multispectral image segmentation)
[6]
Number of Hidden Dimensions
[128]
Dropout
[0, 0.1, 0.2, 0.2, 0.3]
Weight, Decays
[0.01, 0.1, 0.3, 0.5)
Number of Attention Heads
[1, 2, 4, 8, 16, 32)
Patch Size
[2, 4]
Pooling Function
[Mean]
Timeseries Forecasting The convolution encoder a single 1D convolution layer with kernel size and stride both set to patch size and applied on the temporal axis. This equivalent to splitting the input timescries into patches and applying a linear projection into the latent space of the model. Rotary positional encoding (Su al.] 2021) is applied only on the time axis. The output of the transformer is pooled before being fed to the final MLP layer by either taking the average over time. We conduct forecasting experiments by training models on each datasct. Following the same split of training/ saaininn/vallaaaiinnno tost sets as in (Liu et al. 2024 the model weights from the epoch with the lowest MAE on the validation set are selected for comparison on the test set.
3D Medical Image Classification The convolution encoder is implemented as a single-layer 3D convo- lution with a kernel size and stride equal patch size. The output of the Transformer is pooled before being fed to the final MLP classifier by taking the average. We conduct classification experiments by training models on each dataset. We follow the official split of training aoaining/valiaationn// validationnteess sets. The model weights from the epoch with the highest AUC score on the validation set are selected for comparison on the test set.
Multispectral Image Segmentation The convolution encoder is implemented as single-layer 2D con- volution with a kernel size and stride equal to patch size. The output of the Transformer is directly fed to the final MIP classifier without pooling. We conduct classification experiments by training models on each dataset. We follow the official split of training/valiin validation/test test sets. The model weights from the epoch with the highest ACC score the validation set selected for comparison the test. set..
22

---

Published in Transactions on Machine Learning Research (11/2025)
All the experiments implemented in PyTTrrch and conducted on single NVIDIA A100 GPU with 80 GB of memory, x86 64 CPU with 6 cores and 64GB of RAM. We utilize ADAM (Kingma & Ba 2017) with an initial learning rate of 2 x 10-4 and L2 loss for the timeseries forecasting task and cross-entropy loss for the medical image classification task. Our experiments show that using weight decay is crucial for avoiding overfitting most cases. The batch size uniformly set to 32, and the number of training epochs is fixed 100. We conduct hyperparameter tuning basod on the search space shown Table
C Baselines Details
C.1 Timeseries Forecasting
We select seven Transformer-based models as benchmarks for our study, each offering unique approaches time series forecasting:
Crossformer (Zhang & Yan 2023]: This model explicitly captures dependencies across different dimensions in multivariate time series by embedding imput data into two-dimensional vector array and employing a Two-Stage Attention mechanism to model both temporal and cross-dimensional relationships efficiently.
Autoformer (Wu et al. 2021 Autoformer introduces an Auto-Correlate Correlation mechanism to replace traditional self-attention, aiming to enhance the capture of long-term temporal dependencies while reducing computational complexity in time series forecasting.
FEDformer (Zhou et al. 2022; FEDformer leverages frequency-enhance decomposition method to model time series data, focusing on extracting both trend and scasonal components to improve Torecasting accuracy.
PStchTST (Nie al. 2023 PStchTST proposes a patching mechanism that partitions time series data into patches spanning multiple time steps, enabling the model to capture local temporal patterns and improve forccasting performance.
Transformer (spatial) (Liu et al. 2024): This variation applies the attention mechanism along the spatial dimension, focusing on modeling dependencies between different variables each time step.
Transformer (temporal) (iu el al.
2024): contrast, this version applies attention along the temporal axis, emphasizing the modeling of temporal dependencies within each variable over time.
Transformer (non-factorized): This baseline employs the vanilla transformer architecture with a full self-attention mechanism over flattened input data, corresponding to non-factorized spatiotem- poral attention. We consider this model as the main baseline to beat, as it models the full high-order attention.
C.2 3D Medical Image Classification
We select seven medical image classifier models as baselines, each employing distinct methodologies for 3D medical image analysis:
ResNeet-88RReseel (He et al. 2015 Yang et al. 2023): These convolutional neural networks (CNNs) utilize residual learning to facilitate the training of deeper architectures, effectively address- ing the vanishing gradient problem. They have been widely adopted for various image classification tasks, including medical image analysis.
MDANet (Hang et al., 2022: : The Multi-Scale Discriminative Attention Network (MDANet)
integrates multi-scale feature extraction with discriminative attention mechanisms enhance the representation of critical regions in medical images, thereby improving classification performance.
23

---

Published in Transactions on Machine Learning Research (11/2025)
The Cross- Cooss-Dimensional Transformer (CdIrransformer) employs cross-dimensional attention mechanism to capture both spatial and channel-wise dependencies in 3D medical images, facilitating more comprehensive feature learning for classification tasks.
CdTransformer (Zuu et al. 2024):
ViT-3D (Lai et al., 2024p: This model extends the Vision Transformer (ViT) architecture to 3D medical image classification by dividing volumetrie data into non-overlapping 3D patches and pro- cessing them using transformer encoders, effectively capturing global context. without relying convolutional operations.
ViVIT-S [Ta et al. 2024): The Video Vision Transformer (ViViT) adapts transformer architectures for video classification by modeling spatial and temporal dimensions jointly. The ViVIT-S variant smaller-scale version designed to capture spatiotemporal features efficiently, making it suitable for 3D medical image classification tasks.
TimeSFormer [Bertasius et al. 2021]: TimeSFormer introduces factorized attention mechanism that separately processes spatial and temporal dimensions, significantly reducing the computational complexity of standard full attention for video understanding. achieves competitive performance on video benchmarks while being more memory-efficient than naive attention implementations.
Multiscale Vision Transformer (Fan al 2021 MVIT improves efficiency and scalability by progressively reducing spatial resolution while increasing channel capacity across layers. This multiscale design allows the model to capture hicrarchical representations, making it suitable for lasks requiring fine-grained and coarse-grained feature integration.
Transformer (non-factorized): This baseline employs the vanilla transformer architecture with full self-attention mechanism dependencics across all spatial dimensions without any factorization dimensional reduction. We consider this model as the main baseline to beat as it utilizes the full high-order atlention.
flattened 3D medical image data, aiming to capture complex
C.3 Multispectral Image Segmentation
We select three models as baselines, each employing distinct methodologies:
ResNeet-18/RssNeee Hee et al. 2015)
ViT Dosovitskiy et al. 2020 VIT applies standard transformer architecture directly to image patches by flattening and embedding them as a sequence. While it achieves strong performance on image classification tasks, its full attention mechanism scales quadratically with the number of patches, leading to higher memory and computational requirements compared factorized hierarchical transformer variants.
Additional Experiments
D.1 Positional Encoding
In this section, we explore the influence of various positional encoding (PF) strategies on the performance of our proposed HOT model, focusing on the MedMNIST3D dataset. We evaluate 3D adaptations of widely- used PE methods, including sinusoidal encodings, learnable absolute positional encodings, and Rotary Po- sitional Embeddings (RoPE). RoPE is applied mode-wise during the computation of the attention matrix for each corresponding mode, ensuring that positional information is effectively integrated into the model's attention mechanism.
The results reveal that the choice of positional encoding significantly impacts performance, with RoPE emerging as the optimal strategy for the Kronecker product variant of HOT. RoPE consistently achieves the highest accuracy and AUC scores across the Organ, Nodule, and Vessel datasets, with particularly strong results the Organ dataset (99.8% AUC and 94.4% accuracy). In contrast, for the Kronecker variant.
24

---

Published in Transactions on Machine Learning Research (11/2025)
Table 7: Ablations on different positional encoding strategies on 3D medical image classification task.
Organ
Nodule
Vessel
Model
PE
AUC
ACC
AUC
ACC
AUC
ACC
NoPE
99.2
88.5
87.2
85.5
80.7
89.3
3D Absolute PE
99.2
89.7
91.3
87.1
82.5
91.1
HOT (product)
3D Sin-Cos PE
99.3
90.1
89.6
86.9
85.7
90.5
RoPE
99.8
94.4
92.0
89.1
85.7
90.6
NoPE
99.2
89.2
87.9
86.1
79.8
90.1
3D Absolute PE
99.4
90.8
92.8
88.3
85.8
91.6
HOT (sum)
3D Sin-Cos PE
99.3
90.3
90.5
87.1
86.2
90.5
RoPE
99.6
92.2
90.7
87.7
84.9
90.7
3D absolute learnable positional encoding proves to be the most effective in most cases, such as on the Nodule and Vessel datasets, where achieves the highest AUC scores of 92.8% and 85.8%, respectively,
A key observation is that incorporating any form of positional encoding consistently improves performance compared to models without positional information NoPE. For instance, on the Vessel dataset, the use of 3D Sin-Cos PE boosts AUC by approximately 5% compared to NoPE. Notably, even without positional encodings, both HOT variants remain competitive with baseline models, demonstrating the robustness of the HOT framework. These findings underscore the importance selecting appropriate positional encoding strategics tailored the specific architccture and task, with RoPE and 3D absolute PE emerging as strong candidates depending on the model variant. and dataset characteristics.
D.2 Component Analysis
This scction examines the contribution of key components Attentiion and Feodforward Networks (FFN) to the performance of the HOT model in time-series forecasting tasks. By systematically removing these com- ponents, we assess their individual and combined impact on forecasting accuracy across the Solar, Weather, and ECL datasets.
Table 8: Ablations on HOT components on timescrios forecasting task. We remove the components and assess the impact on the performance of HOT (product) model. The average results of all predicted lengths are listed here.
Components
Solar
Weather
ECL
Model
Attention
FFN
MSE
MAE
MSE
MAE
MSE
MAE
x
0.240
0.269
0.265
0.283
0.193
0.288
/
0.225
0.260
0.251
0.279
0.177
0.270
HOT (product)
0.221
0.257
0.245
0.275
0.169
0.268
X
0.230
0.272
0.250
0.282
0.190
0.288
HOT (sum)
V
0.220
0.260
0.245
0.275
0.167
0.266
The results Table g reveal that the attention module plays a pivotal role in the model's performance. Removing the attention module leads to significant degradation in forecasting accuracy, with higher MSE and MAE values across all datasets. For instance, the Solar dataset, the absence of attention increases
25

---

Published in Transactions on Machine Learning Research (11/2025)
the MSE from 0.221 to 0.240, underscoring its importance in capturing temporal dependencies. In contrast, removing the FFN results in comparatively smaller performance drop, suggesting that while the FFN contributes to the model's effectiveness, it is less critical than the attention mechanism. The best performance is achieved when both components are included, with the HOT (product) variant achieving the lowest MSE and MAE values across all datasets. For example, on the ECL dataset, the full model achieves an MSE of 0.169, compared to 0.193 when attention removed.
These findings highlight the complementary roles of attention and FFN in the HOT framework. While atten- tion is indispensable for modeling complex temporal relationships, the FFN provides additional refinement, contributing to the model's overall robustness and accuracy. This analysis underscores the importance of retaining both components to achieve optimal forecasting performance.
D.3 Robustness Analysis
In this section, we evaluate the robustness of HOT by analyzing the stability of its performance across multiple runs with different random seeds. This analysis is conducted for both 3D medical image classification and time-series forecasting tasks, ensuring that the model's performance is consistent and reliable under varying initial conditions.
Table 9: Robustness of HOT performance on 3D medical image classification. The results are obtained from five random seeds.
Organ
Nodule
Fracture
Adrenal
Vessel
Variant
AUC
ACC
AUC
ACC
AUC
ACC
AUC
ACC
AUC
ACC
Product
99.8 + 0.4
94.4 + 1.5
92.0 + 0.5
89.1 + 1.7
73.6 + 0.3
58.3 + 1.5
88.6 = 0.2
83.8 1.3
85.7 + 0.7
91.14 1.8
Sum
99.6 + 0.3
92.2 = 1.
90.7 + 0.6
87.7 + 1.5
74.9 + 0.3
58.7 + 1.3
87.3 = 0.5
82.9 = 1.6
84.9 + 0.6
90.04 1.2
Table 10: Robustness of HOT performance on timeseries forecasting. random seeds.
The results are obtained from five
ECL
Weather
Variant
MSE
MAE
SMAPE
MSE
MAE
SMAPE
Product
0.169 + 0.002
0.268 + 0.003
0.520 + 0.003
0.245 + 0.004
0.275 + 0.003
0.642 + 0.005
Sum
0.167 + 0.003
0.266 + 0.002
0.516 + 0.003
0.245 + 0.005
0.275 + 0.004
0.637 + 0.003
Traffic
Solar
Variant
MSE
MAE
SMAPE
MSE
MAE
SMAPE
Product
0.420 + 0.004
0.278 + 0.006
0.515 + 0.005
0.221 + 0.002
0.257 + 0.003
0.420 + 0.004
Sum
0.428 0.003
0.280 + 0.005
0.530 + 0.004
0.220 + 0.002 0.260 + 0.002
0.416:t 0.003
For 3D medical image classification (Table both HOT (product) and HOT (sum) exhibit stable performance across five runs, with low standard deviations in AUC and accuracy metrics. For example, on the Organ dataset, HOT (product) achieves an AUC of 99.8 + 0.4, while HOT (sum) maintains an AUC of 99.6 + 0.3. In time-series forecasting (Table 10 both variants consistently achieve low error rates with minimal variability, as evidenced by SMAPE, MSE, and MAE metrics. On the ECL. dataset, HOT (product) achieves an MSE of 0.169 + 0.002, and HOT (sum) achieves an MSE of 0.167 + 0.003.
The robustness analysis confirms that HOT delivers stable and reliable performance across both medical imaging and timo-scries forccasting tasks. Its low variability in key metrics, combined with strong predictive accuracy, makes dependable choice for wolldworld applications where consistency is critical.
26

---

Published in Transactions on Machine Learning Research (11/2025)
D.4 Visualization of Attention Maps
Layer
Leyer
Layer
Layer
105
130
12
120
109
160
210
210
215
038
291
(00
Leyer Head
Layer Head
Layer 1. Need
Layer
63
112
112
354
134
150
195
535
230
200
700
03N2 82 2838888
Layer 2, Head
Layer 2, Head
Layer 2, Head 3
Layer 2, Head
10
112
132
112
154
154
150
100
190
300
302
Layer 2, Head
Layer Head
Layer 2, Head 7
Layer Headd
112
154
182
182
224
224
200
300
300
305
300
4,0 8F
2,9
180 8P88
Figure 10: Visualization of the spatial attention maps learned by HOT (product) on ECL dataset.
27

---

Published in Transactions on Machine Learning Research (11/2025)
Layer Heed
Layer I. Head
Layer 1. Head
Layer 1. Head
81
24
R
20
12
Layer
Layer
Layer
71
84
14
Leyer Heed
Heed
2. Heed
2. Heed
21
24
8
Layer Head
Layer 2, Head
Layer 2, Head
Layer 2, Head
10
14
16
10
1
16
30
36 20
15 M
10 30.00
Figure 11: Visualization of the temporal attention maps learned by HOT (product) on F.C. dataset.
28

---

Published in Transactions on Machine Learning Research (11/2025)
Layer 1. Head 1
Layer Head
Layer Head
Layer 1, Head 4
180
126
214
301
300
500
352
PS87SE880 120000 2220
0550 2E2 1040040110
75226
Hewed 5
Head
Head
Heedd
14
145
140
140
140
182
182
102
182
234
234
200
300
0.50
Head
Head
14
50
63
145
140
140
145
320
12
50
00
140
152
210
234
200
204
AR
Figure 12: Visualization of the spatial attention maps learned by HOT (sum) on the ECL dataset.
29

---

Published in Transactions on Machine Learning Research (11/2025)
Layer 1, Head
Layer Heed
Layer Head
Layer
32
12
14 14
22
12
15
18
20
10
16
20
20
14
15
Layer Head
Layer
Head
R
35
Layer 2, Head
Layer 2, Head
Layer 2, Head
Layer 2. Head
10 55
22
Layer
Layer
Layer
22
Figure 13: Visualization of the temporal attention maps learned by HOT (sum) on the ECL datasct.
30


Reviews:
[{"review_id": "wIooHmlzxw", "reviewer": "Reviewer_WHxp", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes Higher-Order Transformers (HOT), a generalized Transformer architecture designed to efficiently handle high-dimensional tensor data by factorizing the attention mechanism using Kronecker products and sums. The key contribution is decomposing the computationally expensive full attention matrix (which scales quadratically with tensor size) into mode-wise attention matrices that can be combined through Kronecker operations, reducing complexity from O(D(N₁N₂...Nₖ)²) to O(D(∑ᵢNᵢ)(∏ⱼNⱼ)) while preserving tensor structure and cross-dimensional dependencies, where D is the embedding dimension, and N₁N₂...Nₖ are the sizes of tensors. The authors provide theoretical guarantees showing that their factorized approach can approximate any high-order attention matrix with sufficient rank, and demonstrate competitive performance on multivariate time series forecasting and 3D medical image classification tasks with significantly reduced computational and memory costs compared to full attention mechanisms.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Strengths:\n1. The paper provides rigorous theoretical analysis including universality guarantees (Theorem 4.3) showing that Kronecker decomposition can represent any high-order attention matrix, along with stability analysis of the row-stochastic properties and rank bounds.\n2. HOT achieves substantial reductions in memory usage (40-60%) and training time while maintaining competitive performance, making it practical for resource-constrained applications with high-dimensional data.\n\nWeaknesses:\n1. The evaluation is restricted to relatively small-scale datasets (28×28×28 medical images, time series with modest dimensions), leaving questions about scalability to truly large tensor data that would most benefit from this approach.\n2. While the paper proves universal approximation capability, it lacks detailed analysis of how well low-rank Kronecker decompositions approximate full attention in practice, and doesn't provide clear guidance on selecting the factorization rank for different data characteristics.", "audience_yn": "Yes", "audience_explanation": "This paper addresses a fundamental computational bottleneck that affects multiple active research areas in machine learning. The quadratic scaling of attention mechanisms has been a critical challenge limiting the application of Transformers to high-dimensional data, and this work provides a principled solution through tensor factorization that maintains theoretical guarantees while achieving substantial practical speedups. The approach is particularly relevant to researchers working on video analysis, 3D medical imaging, climate modeling, and multivariate time series - domains where data naturally exists in tensor form but current methods either flatten the structure (losing important relationships) or face prohibitive computational costs. I believe that the broader TMLR community would likely find value in both the theoretical contributions and practical implications.", "requested_changes": "1. The experimental evaluation is limited to relatively small tensors (28×28×28 medical images, modest time series dimensions). The paper would benefit from experiments on larger-scale data where the computational advantages would be most pronounced - such as high-resolution video analysis, large spatiotemporal climate datasets, or higher-dimensional medical volumes.\n2. While Theorem 4.3 proves universal approximation capability, the paper lacks analysis of how approximation quality degrades with different factorization ranks in practice. Adding empirical studies showing the trade-off between rank and approximation error would provide clearer guidance for practitioners.\n3. Some complexity expressions omit the hidden dimension D. The paper should consistently include D in all complexity analyses for accuracy.", "broader_impact_concerns": "The paper demonstrates HOT on 3D medical imaging classification but doesn't address the critical ethical implications of deploying AI models in healthcare.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OE4X8KvpJC", "reviewer": "Reviewer_rjzJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces HOT, a higher order Transformer architecture using Kronecker products (sums).\nThe approach is theoretically introduced and tested on a couple of datasets showing good results as well as efficiency improvements.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper is - for the methodology and notation section - nicely written. I am not particularly an expert in this area, but the proofs and concepts check out and are consistent", "audience_yn": "Yes", "audience_explanation": "Higher order Attention for Higher order inputs is of high practical and theoretical relevance", "requested_changes": "# Strength\n- Overall, I think its a nice approach with really concise and nice presentation of the methodology\n- It is, to the best of my knowledge a novel idea\n- Very good preliminaries section with concise explanations for unfamiliar readers\n\n# Weaknesses\n- Throughout the paper there are several typos and grammatical issues. I'd ask the authors to do a proper and careful reading to smooth out all of the rough edges\n- Do I understand it correclty, that you perform hpo for the presented model, but not for the baselines and that experimental results are obtained by a single run?\n    - Why are the number of blocks and hidden dimensions fixed over HPO runs? Is it a randomized search then over different dropout, weight decay, attention heads and patch size parameters?\n- For runtime analysis please provide all details, i.e. system, GPU, epochs, model size etc.\n\n\n# Minor/Questions\n- Not sure, whether (KAN) is the best abbreviation. I know KAN as Kolmogorov Arnold Networks, but\n- Do I understand it correctly in section 4.1 that this is just flattening + reshaping? Also, maybe it could be made clearer that this is not your suggested approach, but something that is commonly used? It's rather unprominently written in the last sentence of the paragraph\n- The bullet points contributions in the beginning are rather a poor representation of the actual contributions. \n- Reporting MSE and MAE seems a bit redundant for the forecasting table\n-", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HjfcR5lKzT", "reviewer": "Reviewer_nCnc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper aims to propose an alternative attention mechnisim which replaces the standard self-attention on high-dimensional data, e.g. 3D data, with Kronecker-Structured attention. With that, the proposed method can reduce the training time and memory usage significantly while retaining the competitive performance. The proposed method is evaluated via two types of datasets, time-series prediction and classification tasks.\n\nStrengths: \n\n1. The proposed attention method reduces the training time and memory usage significantly, which opens the potential to scale the model size with reasonable training time and hardware resources. \n\n2. The derivation of HOT is clear and easy to follow, and the support of theoretical analysis on the design of HOT is good.  \n\nWeaknesses:\n1. For the experiments on 3D data classification, there is no comparison to other efficient attention methods, that are applied on the similar data, e.g. video. \n\n2. Even though the paper included empirical savings in training time and memory usage, there is no analytic analysis how the savings on FLOPs with respect to data size;  moreover, there is no comparison on inference time.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper provided theoretical analysis on the proposed HOT and evaluated on two types of datasets.", "audience_yn": "Yes", "audience_explanation": "The readers will be interested to learn more about efficient attention and learning how to apply it to 3D data.", "requested_changes": "1. See weakeness.\n\n2. The paper lacks the comparison to those 3D attention methods that are used in video modeling as the video is also a part of 3D data, e.g. TimeSformer [1] and MViT [2].\nReference:\n[1] Is Space-Time Attention All You Need for Video Understanding?\n[2] Multiscale Vision Transformers\n\n3. Any insight that the proposed method could work on video recongition? \n\n4. For section 5.3.2, doesn't the number of attention heads also affect the vanilla attention? Is that phonemonon uniqle to the proposed HOT? It will be great to include the results of original attnetion.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)