Understanding axial attention in TSFMs through single location regression

Andrei-Tiberiu Pantea; Aymeric Dieuleveut

stat.MLOct 29, 2026

Understanding axial attention in TSFMs through single location regression

Authors: Andrei-Tiberiu Pantea, Aymeric Dieuleveut

Organizations: CMAP, École polytechnique, IP Paris; Electricité de France R&D; Palaiseau, France · CMAP, École polytechnique, IP Paris; Palaiseau, France

Abstract

Recent advances in Time Series Foundation Models (TSFMs) increasingly rely on axial attention mechanisms, but despite their empirical success, the theoretical properties of axial attention remain largely unexplored. To bridge this gap, we extend the Single Location Regression framework introduced in Marion et al. [2025] to 2D structured data, in order to study sparse information settings where a target depends on a small number of tokens. We formulate an analytically tractable axial attention-type predictor and rigorously analyze its statistical properties and training dynamics. We prove that under a specific asymptotic regime and with knowledge of the underlying oracle parameters, our predictor achieves Bayes optimality, whereas a class of semi-linear simplifications strictly fails to do so, underlining the role of the inner non-linearity. We prove that for a specific temperature scaling, Projected Gradient Descent converges globally to the optimal oracle parameters. Extensive numerical experiments validate our theoretical findings and highlight the critical role of inverse temperature scheduling in empirical convergence.

Figures & tables

Figure 1: Schematic view of an axial attention layer. Raw input tokens are first processed by a time-wise self-attention layer that isolates features along the rows. A subsequent channel-wise attention layer vertically blends these extracted features, effectively capturing global cross-channel and temporal context.

Figure 2: Under normal operating conditions, the background embeddings behave as uninformative noise, or as a standard token, modeled by the standard normal distribution $\mathcal{N}(0,I_{d})$ . When an anomaly occurs in one sensor, its specific signature is encoded into the token $X_{I_{0}J_{0}}$ , characterized by a shifted mean $\ell k^{*}$ and a localization variance $\gamma^{2}I_{d}$ . The regression target $Y$ represents a downstream consequence or system-level anomaly score that depends strictly on identifying this specific embedded state via the projection vector $v^{*}$ . Our goal is to investigate whether we can recover this target through an axial attention layer.

Panel of Figure 3 — Figure 3: Dynamics of stochastic PGD to the oracle parameters over 20 independent runs. The solid curve represents the empirical median across the runs and the bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals (although not always visible). Left: The empirical risk obtained by our predictor on the training batches of data using both the training parameters $k$ , $v$ and the oracle parameters $k^{*}$ , $v^{*}$ . For computing the oracle and model losses we use the initial inverse temperatures. Middle-left: The alignment between the training and oracle parameters $\kappa=k^{T}k^{*}$ , $\nu=v^{T}v^{*}$ . Middle-right: The distance from the manifold $\mathcal{M}$ of the training parameters $(k,v)$ . Right: The schedule used for $\lambda=\lambda_{r}=\lambda_{c}$ . The hyperparameters of the experiment are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ and $\lambda_{t}=5\cdot 10^{-2}/(1+\max(t-2\cdot 10^{4},0)\cdot 10^{-4})$ .

Panel of Figure 4 — Figure 4: Leftmost: Alignment dynamics during the training of our predictor with an inverse temperature warm-up period of 20k steps using $\lambda_{c}=\lambda_{r}=0.1$ , followed by constant $\lambda_{c}=\lambda_{r}=\lambda$ for $\lambda\in\{0.01,0.03,0.05,0.07\}$ . We plot the median across 20 runs for each value of $\lambda$ . The dashed vertical line represents the warm up period. Center-left to right: The numerically computed gradients of $\mathcal{R}^{<}_{\lambda_{r},\lambda_{r}}(\kappa,\nu)$ for $\lambda_{r}=\lambda_{c}=\lambda\in\{0.01,0.03,0.05,0.07,0.1\}$ , $N=5$ and $\ell=100$ . The gradient, given by ( 32 ), is approximated using the Monte Carlo method with a sample size of 10,000.

Panel of Figure 5 — Figure 5: Convergence of stochastic PGD over 20 runs with random initializations on $\mathbb{S}^{d-1}\times\mathbb{S}^{d-1}$ , optimizing for the penalized empirical risk ( 9 ). The solid curve represents the empirical median across the runs. The hyperparameters of the experiment are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ , $\mu_{1}=10$ , $\mu_{2}=1$ . The inverse temperature schedules are $\lambda_{t}=10^{-1}/(1+\max(t-5\cdot 10^{3},0)\cdot 10^{-4})$ , for the erf predictor and $\lambda_{t}=10^{-1}/(1+\max(t-2\cdot 10^{4},0)\cdot 10^{-4})$ for the softmax predictor. For computing the oracle and model losses we use the initial inverse temperatures.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Fig. 3	Fig. 4 , Fig. 13	Fig. 5	Fig. 6
	( $\mathcal{M}$ initialization)	(Constant $\lambda$ s)	(Penalized)	(Unpenalized)
Grid height ( $M$ )	5	5	5	5
Grid width ( $N$ )	5	5	5	5
Embedding dim. ( $d$ )	20	20	20	20
Shift magnitude ( $\ell$ )	100	100	100	100
Special token variance ( $\gamma^{2}$ )	1	1	1	1

Appendix

Table 1: Summary of hyperparameters used across the numerical experiments.

Panel of Figure 6 — Figure 6: Optimization dynamics from 20 random initializations without the alignment penalties ( $\mu_{1}=\mu_{2}=0$ ). The solid curve represents the empirical median across the runs. The bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals. The parameters used are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ and $\lambda_{t}=10^{-1}/(1+\max(t-5\cdot 10^{4},0)\cdot 10^{-4})$ , for the erf predictor and $\lambda_{t}=10^{-1}/(1+\max(t-2\cdot 10^{4},0)\cdot 10^{-4})$ for the softmax predictor. For computing the oracle loss we use the initial inverse temperatures.

Panel of Figure 7 — Figure 7: Training dynamics of the predictor with a 20k warm-up phase $\lambda_{c}=\lambda_{r}=\lambda=0.1$ followed by constant inverse temperatures $\lambda_{c}=\lambda_{r}=0.01$ . The solid curve represents the empirical median across 20 runs and the bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals (although not always visible). The parameters used are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ .

Panel of Figure 8 — Figure 8: Training dynamics of the predictor with a 20k warm-up phase $\lambda_{c}=\lambda_{r}=\lambda=0.1$ followed by constant inverse temperatures $\lambda_{c}=\lambda_{r}=0.03$ . The solid curve represents the empirical median across 20 runs and the bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals (although not always visible). The parameters used are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ .

Panel of Figure 9 — Figure 9: Training dynamics of the predictor with a 20k warm-up phase $\lambda_{c}=\lambda_{r}=\lambda=0.1$ followed by constant inverse temperatures $\lambda_{c}=\lambda_{r}=0.05$ . The solid curve represents the empirical median across 20 runs and the bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals (although not always visible). The parameters used are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ .

Panel of Figure 10 — Figure 10: Training dynamics of the predictor with a 20k warm-up phase $\lambda_{c}=\lambda_{r}=\lambda=0.1$ followed by constant inverse temperatures $\lambda_{c}=\lambda_{r}=0.07$ . The solid curve represents the empirical median across 20 runs and the bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals (although not always visible). The parameters used are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ .

Panel of Figure 11 — Figure 11: Numerically computed gradients of the theoretical risk $\mathcal{R}^{<}_{\lambda_{r},\lambda_{c}}(\kappa,\nu)$ for $\lambda_{r}\in\{0.1,0.01\}$ and varying values of $\lambda_{c}$ . The parameters used are $M=N=5$ and $\ell=100$ . The gradients, given by Equation ( 32 ), are approximated using the Monte Carlo method with a sample size of 10,000 on a $100\times 100$ discretization grid.

Panel of Figure 12 — Figure 12: Numerically computed gradients of the theoretical risk $\mathcal{R}^{<}_{\lambda_{r},\lambda_{c}}(\kappa,\nu)$ for $\lambda_{c}\in\{0.1,0.01\}$ and varying values of $\lambda_{r}$ . The parameters used are $M=N=5$ and $\ell=100$ . The gradients, given by Equation ( 32 ), are approximated using the Monte Carlo method with a sample size of 10,000 on a $100\times 100$ discretization grid.

Panel of Figure 13 — Figure 13: Training dynamics of the predictor with constant $\lambda_{c}=0.01$ and $\lambda_{r}=0.1$ . The solid curve represents the empirical median across the runs and the bands represent the 0.05-0.95 and, respectively, the 0.25-0.75 quantile intervals (although not always visible). The parameters used are $M=N=5$ , $d=20$ , $\gamma=1$ , $\ell=100$ .

Panel of Figure 14 — Figure 14: Numerically computed gradients of the theoretical risk $\mathcal{R}^{<}_{\lambda_{r},\lambda_{c}}(\kappa,\nu)$ for $\lambda_{r}=\lambda_{c}=0.1$ , $\ell=10$ and varying values of $N$ . The gradients, given by Equation ( 32 ), are approximated using the Monte Carlo method with a sample size of 10,000 on a $100\times 100$ discretization grid.

Panel of Figure 15 — Figure 15: Numerically computed gradients of the theoretical risk $\mathcal{R}^{<}_{\lambda_{r},\lambda_{c}}(\kappa,\nu)$ for $\lambda_{r}=\lambda_{c}=0.1$ , $\ell=100$ and varying values of $N$ . The gradients, given by Equation ( 32 ), are approximated using the Monte Carlo method with a sample size of 10,000 on a $100\times 100$ discretization grid.

Panel of Figure 16 — Figure 16: Numerically computed gradients of the theoretical risk $\mathcal{R}^{<}_{\lambda_{r},\lambda_{c}}(\kappa,\nu)$ for $\lambda_{r}=\lambda_{c}=0.1$ , $\ell=1000$ and varying values of $N$ . The gradients, given by Equation ( 32 ), are approximated using the Monte Carlo method with a sample size of 10,000 on a $100\times 100$ discretization grid.

Attention Mechanisms

Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.

See all 517 papers

Universal Redundancies in Time Series Foundation Models
Feb 2, 2026Anthony Bao, Venkata Hasith Vattikuti, Jeffrey Lai +1Transformer Interpretability Kernel Regression
ECE Department, UT Austin · Department of Physics, UT Austin · Oden Institute, UT Austin
High-Dimensional Learning Dynamics of Attention-Indexed Models
Sep 3, 2026Yizhou Xu, Margarita Sagitova, Lenka Zdeborová +1Implicit Bias Attention Mechanisms
1Information, Learning and Physics Laboratory, École Polytechnique Fédérale de Lausanne (EPFL) · 2Statistical Physics of Computation Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)
Post-Training in Time Series Foundation Models: A Unifying Framework
Jul 22, 2026Shifeng Xie, Ambroise Odonnat, Zehao Xiao +7Model Compression Time-Series Foundation Models
Noah’s Ark Lab · Université Paris Cité · Inria +2

Feb 2, 2026J/K move · Enter open · S save

Universal Redundancies in Time Series Foundation Models

Anthony Bao, Venkata Hasith Vattikuti, Jeffrey Lai, William Gilpin

ECE Department, UT Austin · Department of Physics, UT Austin · Oden Institute, UT Austin

Transformer Interpretability Kernel Regression Mechanistic Interpretability Transformer Pruning Time-Series Foundation Models Neural Network Pruning

Open paper AskPDF arXiv

Time Series Foundation Models (TSFMs) leverage extensive pretraining to accurately predict unseen time series during inference, without the need for task-specific fine-tuning. Through large-scale evaluations on standard benchmarks, we find that leading transformer-based TSFMs exhibit redundant components in their intermediate layers. We introduce a set of tools for mechanistic interpretability of TSFMs, including ablations of specific components and direct logit attribution on the residual stream. Our findings are consistent across several leading TSFMs with diverse architectures, and across a diverse set of real-world and synthetic time-series datasets. We discover that all models in our study are robust to ablations of entire layers. Furthermore, we develop a theoretical framework framing transformers as kernel regressors, motivating a purely intrinsic strategy for ablating heads based on the stable rank of the per-head projection matrices. Using this approach, we uncover the specific heads responsible for degenerate phenomena widely observed in TSFMs, such as parroting of motifs from the context and seasonality bias. Our study sheds light on the universal properties of this emerging class of architectures for continuous-time sequence modeling.

Abstract

Attention Mechanisms

Explore similar work

Universal Redundancies in Time Series Foundation Models

High-Dimensional Learning Dynamics of Attention-Indexed Models

Post-Training in Time Series Foundation Models: A Unifying Framework