Understanding axial attention in TSFMs through single location regression
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.
Related
- Patch Attention Wisely: Understanding Patch Length in Attention-Based Forecasting
- Convergent Stochastic Training of Multi-Headed Attention and Understanding LoRA
- Diagonalized Attention for Individualized Regression: Latent-Row Localization and Prediction
- Sharp Generalization for Nonparametric Regression in Interpolation Space by Shallow Over-Parameterized Neural Networks with Channel Attention
- A Solution Method for Single-Source Capacitated Facility Location Problems Under Decision-Dependent Demand Uncertainty
BibTeX
@misc{pantea2026understanding,
title = {{Understanding axial attention in TSFMs through single location regression}},
author = {Andrei-Tiberiu Pantea and Aymeric Dieuleveut},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.40188},
eprint = {2610.40188},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.40188},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Understanding axial attention in TSFMs through single location regression
AU - Andrei-Tiberiu Pantea
AU - Aymeric Dieuleveut
PY - 2026
DA - 2026/10/29
PB - arXiv
DO - 10.48550/arxiv.2610.40188
UR - https://arxiv.org/abs/2610.40188
AB - 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.
ER -