Source-linked AI summary
Understanding axial attention in TSFMs through single location regression
Andrei-Tiberiu Pantea, Aymeric Dieuleveut
TL;DR
Axial attention is widely used in time-series foundation models, yet its theoretical properties in sparse-information settings remain largely unexplored. The paper extends single-location regression to 2D data and analyzes a tractable axial-attention-type predictor. It proves asymptotic Bayes optimality with oracle parameters and global convergence of projected gradient descent under specified conditions.
Problem
Theoretical understanding of axial attention remains limited, including whether its restricted information flow can support optimal prediction in token-wise sparse settings.
Source: §1 · pp. 2–3Method
Results
The predictor is asymptotically Bayes optimal with oracle parameters, semi-linear variants without the inner non-linearity are strictly suboptimal, and Projected Gradient Descent converges globally under specified temperature scaling.
Source: §5 · p. 9Takeaways & Limitations
The results establish theoretical properties for a simplified axial-attention predictor, while experiments highlight inverse-temperature scheduling in convergence.
Source: §5 · p. 9The theory concerns a simplified erf predictor with shared keys, so implications for full TSFMs with realistic projections, residuals, normalization, and multiple blocks are qualitative.
Source: §5 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
Axial attention is increasingly used in time-series foundation models, but its theoretical behavior in sparse-information settings remains unclear. This paper extends single-location regression to 2D data and analyzes an axial-attention-type predictor and its training dynamics.
1 Introduction
- Axial attention is used in several time-series foundation models, while flattened global attention is a computational bottleneck and does not explicitly use the time-by-channel grid.The paper names Chronos-2, TabPFN-TS, and Toto as examples.Source: §1 · p. 1
- The standard layer performs row-wise attention followed by column-wise attention, with the activation applied row-wise.The paper notes that reversing the order does not affect its results.Source: §1 · pp. 1–2, Figure 1 · p. 2
- The paper asks whether axial attention can match global attention's statistical and optimization performance when the target depends on a small number of tokens.Axial attention restricts each token to attending only to tokens sharing an axis.Source: §1 · p. 2
- The authors extend the Single Location Regression framework to 2D structured data and use it to assess axial attention's theoretical capabilities.The motivating applications include multivariate time series and tables.Source: §1 · p. 2
- Their analysis finds an asymptotic regime where the predictor is statistically optimal with oracle parameters, while removing its inner non-linearity makes it suboptimal.They also analyze optimization and report numerical studies of a complex optimization landscape.Source: §1 · p. 3
- Axial attention's sequential, axis-wise information flow breaks the symmetry of global attention and creates cascading nonlinear dependencies that complicate analysis.The authors say this requires new analytical techniques rather than standard closed-form approaches.Source: §1 · p. 3
2 SLR model and simplified axial attention predictor
The paper models a target that depends on one identifiable token in a 2D input, with multivariate time-series anomaly analysis as a motivating application. It then simplifies axial attention to make statistical analysis tractable.
Single location regression model
- The SLR model places a single target-relevant token at random coordinates, with its mean shifted along a localization direction while the other tokens are standard normal.The target depends on the localized token, making its location part of the inference problem.Source: §2 · p. 3
Connection to multivariate time series anomaly analysis
- For anomaly analysis, the model treats a system-wide event as localized in the embedding of one primary driving sensor at one time step.This assumes the embedding layer has already captured correlated dynamics across sensors.Source: §2 · p. 3, Figure 2 · p. 4
An analytically tractable simplification of axial attention for SLR
- The proposed predictor simplifies axial attention to make its statistical analysis tractable, using learnable unit vectors k and v and an intermediate half-step representation.The paper identifies X(1/2) as the half step in both the predictor and axial attention.Source: §2 · p. 4
- The predictor replaces softmax with the component-wise erf nonlinearity because softmax couples tokens globally and obstructs rigorous statistical analysis.The paper presents this substitution as retaining the network's core behavioral properties while enabling analysis.Source: §2 · p. 4
3 Statistical analysis
The analysis reduces the axial predictor’s risk to a lower-dimensional alignment representation and identifies a regime where the oracle predictor is asymptotically optimal. A semi-linear simplification without the inner non-linearity has strictly higher risk under the stated condition.
3.1 Asymptotic optimality of the oracle predictor
- Under a specific asymptotic regime, the oracle axial predictor is optimal; the analysis reduces its risk to five alignment scalars under orthogonality of k∗ and v∗.The reduction makes the risk landscape tractable despite the sequential nature of axial attention.Source: §3.1 · p. 5
3.2 Removing the inner non-linearity is suboptimal
- The semi-linear comparison replaces the inner non-linearity with a fixed linear-combination layer, identifying that non-linearity as necessary for asymptotic optimality.The paper contrasts this simplification with the axial predictor’s inner non-linearity.Source: §3.2 · pp. 5–6
- For any nonzero β, the semi-linear predictor has risk strictly greater than the noise variance when no token location has probability one.This result applies across the stated choices of M, N, ℓ, and λc.Source: §3.2 · p. 6
4 Gradient descent dynamics
The paper analyzes projected gradient descent on a constrained manifold and proves recovery under a sufficiently small inverse-temperature regime. Numerical experiments show that temperature scheduling matters for convergence, including from generic initializations with a manifold penalty.
4.1 Convergence of gradient descent
- Projected Riemannian gradient descent keeps parameters on the manifold, where the reduced risk formulation is valid and the analysis becomes tractable.The paper motivates the restriction because the full five-dimensional dynamics are intractable.Source: §4.1 · p. 6
- For generic initialization on M, Theorem 5 guarantees convergence to the oracle parameters up to joint sign when λr is below a threshold, λc=τλr, and the step size is sufficiently small.The analysis restricts dynamics to M, which is invariant under the stated assumptions when initialized there.Source: §4.1 · pp. 6–7
- Training and inference use separate inverse-temperature regimes: optimization recovers parameters, then inference temperatures yield an asymptotically Bayes-optimal predictor.The paper describes this as a sequential two-stage guarantee.Source: §4.1 · p. 7
4.2 Understanding the dynamics - numerical experiments
- κ reaches one for λ=0.01, while larger tested temperatures can prevent its convergence; the experiments also identify warm-up as necessary at small λ.The alignment and risk-gradient plots vary the inverse temperature after a 20k-step warm-up.Source: Figure 4 · p. 8, §4.2 · pp. 8–9
- From generic random starts, penalties guide both erf and softmax predictors toward M; convergence is slower, and their empirical-loss transitions accompany near-perfect parameter alignments.The erf transition is associated with ν alignment, while the softmax transition occurs when both ν and κ are nearly aligned.Source: §4.2 · pp. 8–9, Figure 5 · p. 9
5 Conclusion and discussion
The paper extends single location regression to 2D grids to analyze axial attention, proving oracle Bayes optimality and global PGD convergence under specified conditions. Its conclusions remain bounded by simplifying assumptions and a single informative token.
5 Conclusion and discussion
- The 2D-grid framework proves asymptotic Bayes optimality for the oracle predictor, while semi-linear variants without inner non-linearity are strictly suboptimal.The paper also establishes global PGD convergence to the optimal oracle parameters for sufficiently small inverse temperatures.Source: §5 · p. 9
- The analysis assumes fixed dimensions, independent background tokens, and a target depending on one localized token; relaxing these assumptions is left for future work.The paper identifies autoregressive dynamics, dense cross-channel correlations, growing dimensions, and multiple informative tokens as possible extensions.Source: §5 · p. 9
- The theoretical results rely on rank-1 query/key projections, erf instead of softmax, and shared keys; implications for full TSFMs are qualitative.Experiments with the corresponding softmax predictor provide encouraging evidence of similar optimization behavior, but full TSFMs also include projection matrices, residual connections, normalization, and multiple axial-attention blocks.Source: §5 · p. 10
- Multi-token regression is proposed as a possible multi-head extension, but its statistical and optimization guarantees remain future work.The authors note that jointly trained heads could have interchangeable roles or compete for the same informative location.Source: §5 · p. 10
A Proof of Lemma 1 and further discussion about architectural constraints
The appendix relates the simplified predictor to an axial-attention block using aggregation tokens and architectural constraints. It also supplies risk calculations supporting the oracle result under a specified asymptotic scaling.
A Proof of Lemma 1 and further discussion about architectural constraints
- The predictor corresponds to an axial-attention block with [REG_ROW] tokens and a [REG_COL] global aggregator, under rank-1 projections, shared keys, and specified projection settings.The equivalence assumes pr = pc = 1, oc = 1, Oc = 1, Kr = Kc, and Vr equal to the identity.Source: A Proof of Lemma 1 and further… · pp. 11–12
- The discussion treats REG tokens as learned global aggregation rather than pooling or positional decoding; shared keys are a strong practical constraint, and the proofs use erf.The rank-1 restriction is motivated by the single informative direction in SLR; extending the analysis to other activations is left as future work.Source: A Proof of Lemma 1 and further… · p. 12
B.1 Proofs of Lemma 2 and of Proposition 3
- For oracle parameters, the risk converges to ϵ^2 when both inverse temperatures times ℓ diverge under the stated scaling.The proof considers κ = ν = 1 for the oracle parameters.Source: §B.1 · p. 15
B.2 Proof of Proposition 4
The proof compares the information retained by row summaries with the information in the input. Conditional on the active location, only its row summary carries information about the target.
B.2 Proof of Proposition 4
- The expected conditional variance after row summaries is at least the variance conditional on the active row summary and location, with equality only under the stated single-column condition.Equality requires some j0 with pj0 = 1 and βj = 0 for every j other than j0.Source: §B.2 · pp. 16–17
- Given the active location, row summaries from rows other than the active row are independent of the target and the active row summary.The conditional variance given all row summaries therefore equals the variance given the active row summary and location.Source: §B.2 · p. 16
- The proof evaluates the active row summary’s relationship to the target using their variances and covariance under the stated Gaussian construction.The construction specifies the active token and row summary in terms of Gaussian variables and reports Var(Y) = γ^2 + ϵ^2 and Cov(Si0, Y) = γ^2βj0v∗.Source: §B.2 · p. 16
C Proof of Theorem 5
The proof reduces PGD on an invariant manifold to dynamics in two alignment variables. For sufficiently small steps and inverse temperatures, it uses risk decrease and fixed-point analysis to establish convergence to oracle minima.
C Proof of Theorem 5
- The manifold M is invariant under PGD, allowing the dynamics to be represented by an autonomous system in the alignment variables κt and νt.The reduction is stated for initialization on M.Source: C Proof of Theorem 5 · p. 17
- For a sufficiently small step size, risk strictly decreases along PGD iterates unless they are at a fixed point, and accumulation points must be stationary.This narrows the convergence analysis to the fixed points of the reduced dynamics.Source: C Proof of Theorem 5 · p. 17
- For sufficiently small inverse temperatures, the reduced dynamics have fixed points (±1, ±1) and (0, 0), with (0, 0) classified as a saddle.The proof further characterizes the saddle’s local geometry to support escape from random initializations.Source: C Proof of Theorem 5 · pp. 17–18
- The local map near (0, 0) has one Jacobian eigenvalue below one and one above one, and is injective on [−1, 1]^2.Together with exclusion of the saddle and maxima, these results yield global convergence to the oracle.Source: C Proof of Theorem 5 · p. 18
C.1 Proof of Lemma 6
The proof establishes that the manifold M is invariant under PGD and that the risk decreases along the reduced dynamics. It then shows that the successive iterates converge to accumulation points.
C.1 Proof of Lemma 6
- On M, cross-term risk derivatives vanish, so the gradients simplify and the PGD updates preserve the manifold's orthogonality constraints.The proof verifies the three constraints for each update and concludes that the updated parameters remain in M.Source: §C.1 · pp. 18–19
C.2 Proof of Lemma 7
- The normalized PGD update on M yields the reduced dynamics g1(κt, νt) and g2(κt, νt).The update normalization follows from orthogonality and the Pythagorean theorem.Source: §C.2 · p. 19
C.3 Proof of Proposition 8
- For a sufficiently small step size, the risk decreases along PGD iterates, and the successive parameter changes tend to zero.The proof uses smoothness on the compact square and a telescoping sum to establish convergence of the relevant gradient terms.Source: §C.3 · pp. 19–20
C.4 Proof of Proposition 9
The appendix analyzes critical points and PGD dynamics, then gives supplementary optimization findings and experimental details. On M, generic random initialization converges to the oracle structure.
C.4 Proof of Proposition 9
- The origin is the only interior critical point of the reduced risk, and it is a saddle point.The proof also analyzes behavior on the boundaries of the parameter square.Source: §C.4 · p. 22
C.5 Proof of Proposition 10
- For a sufficiently small step size, the update map is injective on [−1, 1]2.The proof bounds the update's Lipschitz component and uses αBL < 1 to rule out two distinct inputs mapping to the same output.Source: §C.5 · p. 24
C.6 Final details.
- For a generic random initialization on M, PGD almost surely bypasses the saddle point and maxima, converging to the true oracle structure (±1, ±1).The initializations that converge to the saddle point have Lebesgue measure zero.Source: §C.6 · pp. 24–25
E Additional Experimental Details and Results
- The appendix lists supplementary results on unpenalized training, full Figure 4 results, decoupled inverse temperatures, and variations in ℓ and N.It also states that the experiments were run on a MacBook Pro M1 CPU and that each 20-run experiment took less than an hour.Source: E Additional Experimental Details and… · p. 26, §E.1 · p. 26
E.1 Implementation details and hyperparameters
- Table 1 summarizes the hyperparameters and schedules used across the numerical experiments; µ1 and µ2 weight the structural-constraint penalties.The table is identified as a summary, but the supplied passage does not give its individual values.Source: §E.1 · p. 26, Table 1 · p. 26
E.2 Optimization without Manifold Penalties
- Without alignment penalties, erf runs can become trapped in high-loss regions, while some softmax runs become trapped far from M.Most softmax runs converge similarly to the penalized version, but some remain at high positive ν and high negative κ.Source: §E.2 · pp. 26–27, Figure 6 · p. 27
E.3 Full results for the experiment in Figure 4
The experiments vary constant inverse temperatures after a shared warm-up and show that convergence depends on the chosen value. Lower values leave more runs near the origin, while a decreasing schedule supports escape followed by precise convergence.
E.3 Full results for the experiment in Figure 4
- After a 20,000-step warm-up at λc = λr = 0.1, not all runs converged within 100,000 steps at constant λ ∈ {0.01, 0.03, 0.05, 0.07}.Figures 7–10 show training dynamics for these constant values after the warm-up.Source: §E.3 · p. 27, Figure 7 · p. 28, Figure 8 · p. 28
- Lower constant λ values left more runs stuck around the origin than larger values, which produced more consistent convergence.The runs were evaluated over the simulated 100,000 steps.Source: §E.3 · p. 27
- A decreasing inverse-temperature schedule first pushes parameters away from the saddle and then permits fine-tuned convergence to the exact oracle parameters.A static λ must compromise between escaping the origin quickly and converging precisely.Source: §E.3 · p. 27
E.4 Gradient analysis for decoupled inverse temperatures
The appendix examines how decoupling the row and column inverse temperatures changes the reduced-risk gradient field. It also shows that their ratio provides another way to control training dynamics.
E.4 Gradient analysis for decoupled inverse temperatures
- Breaking symmetry with mixed λc and λr shifts gradient magnitudes and directions, while the 0-level curve of ∂κR< varies substantially across temperature pairs.Figures 11 and 12 present numerically computed gradients for mixed configurations.Source: §E.4 · p. 29, Figure 11 · p. 30, Figure 12 · p. 31
- The ratio τ = λc/λr offers another control on training dynamics; with λc = 0.01 and λr = 0.1, |κ| grows faster, but 100,000 steps do not reach the oracle.In this setting, the runs consistently move toward convergence, and |κ| growth more closely mirrors |ν| growth.Source: §E.4 · p. 29, Figure 13 · p. 32
E.5 Gradient analysis for other hyperparameters N and ℓ
This section concerns gradient analyses for the hyperparameters N and ℓ, though the supplied leads differ on whether it reports findings: one describes their effects on escape from the origin and ∂νR, while the other reports no findings on these analyses.
E.5 Gradient analysis for other hyperparameters N and ℓ
- The tested (N, ℓ) values do not follow the asymptotic regime required for Bayes optimality, and full training for each pair was omitted as computationally expensive.The authors describe the values as more reflective of practical settings with tens of channels and hundreds of time patches, where normalization limits an individual token’s mean shift.Source: §E.5 · pp. 32–33
- The effect of ℓ on the gradients overshadows that of N; for small ℓ, higher N appears to eliminate κ’s escaping-the-origin problem, but this effect disappears at larger ℓ values.Source: §E.5 · p. 33
- As ℓ increases, ∂νR< takes higher values and its sign change around κ = 0 becomes steeper.The gradients were approximated by Monte Carlo sampling with 10,000 samples on a 100 × 100 discretization grid.Source: §E.5 · pp. 32–34
NeurIPS Paper Checklist
- LLMs were used only for writing, editing, or formatting, so the paper marks their use as N/A for core method development.The checklist says the core method does not involve LLMs as important, original, or non-standard components.Source: NeurIPS Paper Checklist · p. 40