Search arXiv⌕ Search
Home/Archive/arXiv

arXiv · 2610.40188

Understanding axial attention in TSFMs through single location regression

Andrei-Tiberiu Pantea·Aymeric Dieuleveut
View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

stat.MLcs.AIcs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Andrei-Tiberiu Pantea, Aymeric Dieuleveut. 2026-10-29. Understanding axial attention in TSFMs through single location regression. https://arxiv.org/abs/2610.40188

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Robust performance metrics for imbalanced classification problems

We show that established performance metrics in binary classification, such as Matthews' correlation coefficient (MCC), Cohen's $κ$, the F-score or the Jaccard similarity coefficient are not robust to class imbalance in the sense that if the proportion of the minority class tends to $0$, the true positive rate (TPR) of the Bayes classifier under these metrics {becomes very small or even tends to $0$}. Thus, in imbalanced classification problems, these metrics favour classifiers which ignore the minority class. To alleviate this issue we introduce robustified modifications of the MCC, of Cohen's $κ$ and of the F-score with an additional tuning parameter which allows to adapt the amount of robustness against class imbalance. As theoretical guarantee we show that the Bayes-optimal classifier for these robustified performance metrics, when expressed in terms of the density ratio $f_1/f_0$ of the class-conditional densities $f_i$, has a threshold parameter which is upper-bounded in terms of the tuning parameters. Therefore, even in strongly imbalanced settings, the TPR associated to this classifier will be bounded away from $0${, provided that $f_1/f_0$ exceeds this bound with positive probability under the positive class}. We numerically illustrate the behaviour of the various performance metrics and the effect of the tuning parameters in simulations as well as on a credit default data set. We also discuss connections to the receiver operating characteristic and precision-recall curves, which provide an alternative perspective on the proposed notion of robustness, and give recommendations on how to combine their usage with performance metrics.

stat.ML↗
arXiv2026

Conformalized Credal Regions for Classification with Ambiguous Ground Truth

An open question in \emph{Imprecise Probabilistic Machine Learning} is how to empirically derive a credal region (i.e., a closed and convex family of probabilities on the output space) from the available data, without any prior knowledge or assumption. In classification problems, credal regions are a tool that is able to provide provable guarantees under realistic assumptions by characterizing the uncertainty about the distribution of the labels. Building on previous work, we show that credal regions can be directly constructed using conformal methods. This allows us to provide a novel extension of classical conformal prediction to problems with ambiguous ground truth, that is, when the exact labels for given inputs are not exactly known. The resulting construction enjoys desirable practical and theoretical properties: (i) conformal coverage guarantees, (ii) smaller prediction sets (compared to classical conformal prediction regions) and (iii) disentanglement of uncertainty sources (epistemic, aleatoric). We empirically verify our findings on both synthetic and real datasets.

stat.ML↗
arXiv2026

On the minimax optimality of Flow Matching through the connection to kernel density estimation

Flow Matching has recently gained attention in generative modeling as a simple and flexible alternative to diffusion models. While existing statistical guarantees adapt tools from the analysis of diffusion models, we take a different perspective by connecting Flow Matching to kernel density estimation. We first verify that the kernel density estimator matches the optimal rate of convergence in Wasserstein distance up to logarithmic factors, improving existing bounds for the Gaussian kernel. Based on this result, we prove that for sufficiently large networks, Flow Matching achieves the optimal rate up to logarithmic factors. If the target distribution lies on a lower-dimensional manifold, we show that the kernel density estimator profits from the smaller intrinsic dimension on a small tube around the manifold. The faster rate also applies to Flow Matching, providing a theoretical foundation for its empirical success in high-dimensional settings.

stat.ML↗