Attention Kernels for Learning Maps Between Heavy-Tailed Measures
Not on the map. Why?
Abstract (arXiv)
Operator learning on probability measures can be accomplished with transformers. For measures with polynomial tails, the exponential weighting in softmax can make the corresponding measure-level attention integrals diverge. This motivates replacing the exponential with slower-growing functions. We construct two benchmarks for operator learning on measures with closed-form targets. We use these benchmarks to study attention kernel growth and data transformation in post-norm transformers. Without data transformation, the softmax models exhibit ensemble collapse on both heavy-tailed benchmarks, while the three slower-growing kernels avoid collapse. Symlog preprocessing allows softmax to avoid collapse on the matrix inverse task but not on the sheared swap task. On the Gaussian control, all four kernels perform similarly. We also examine how sample size affects the sensitivity of empirical energy and Wasserstein distances to tail differences. These results support slower-growing attention kernels as an effective design choice for post-norm transformers learning from heavy-tailed ensembles.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 22 of its 32 references that are in the map. First posted on arXiv 30 Sep 2026.
Builds on
22 of 32 references.
References · 22
- 2017Attention Is All You Needcited by 59,932
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2020Fourier Neural Operator for Parametric Partial Differential Equationscited by 2,806
- 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
- 2019DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operatorscited by 1,965
- 2017Deep Setscited by 1,394
- 2022Transformer Quality in Linear Timecited by 189
- 2022DeepNet: Scaling Transformers to 1,000 Layerscited by 176
- 2023A mathematical perspective on Transformerscited by 160
- 2023The emergence of clusters in self-attention dynamicscited by 133
- 2024Theory, Analysis, and Best Practices for Sigmoid Self-Attentioncited by 62
- 2023A Study on ReLU and Softmax in Transformercited by 53
10 more
- 2020A Mathematical Theory of Attentioncited by 43
- 2023Replacing softmax with ReLU in Vision Transformerscited by 43
- 2024Continuum Attention for Neural Operatorscited by 39
- 2025Learning Enhanced Ensemble Filterscited by 11
- 2024Rethinking Attention: Polynomial Alternatives to Softmax in Transformerscited by 8
- 2023Postprocessing of Ensemble Weather Forecasts Using Permutation-invariant Neural Networkscited by 7
- 2026In-Context Operator Learning on the Space of Probability Measurescited by 7
- 2026Learning Probabilistic Filters with Strictly Proper Scoring Rulescited by 4
- 2026One Operator for Many Densities: Amortized Approximation of Conditioning by Neural Operatorscited by 2
- 2026Measure-to-measure Regression with Transformerscited by 1
Built on by
Not cited yet.