PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention Kernels for Learning Maps Between Heavy-Tailed Measures

Kailen Hargenrader, Edoardo Calvello, Bohan Chen

arXiv:2610.00564v1cs.LG

TL;DR

Softmax’s exponential weighting can make measure-level attention integrals diverge for heavy-tailed measures. The paper builds closed-form operator benchmarks and uses them to study kernel growth and preprocessing; slower-growing kernels avoid collapse across both heavy-tailed tasks without preprocessing, while symlog rescues softmax on only one.

  • Problem

    For polynomial-tailed measures, softmax’s exponential weighting can make measure-level attention integrals diverge, motivating alternatives for operator learning on heavy-tailed ensembles.

    Source: §1 · pp. 1–2

  • Method

    The paper constructs two closed-form operator benchmarks and tests post-norm transformers with different attention kernels and symlog preprocessing, alongside a test of tail sensitivity in empirical distribution distances.

    Source: §5 · p. 8

  • Results

    Without preprocessing, softmax collapses on both heavy-tailed tasks while three slower-growing kernels avoid collapse; symlog rescues softmax only on matrix inverse, and larger budgets reveal tail differences with energy distance.

    Source: §4.2 · pp. 7–8, §5 · p. 8

  • Takeaways & Limitations

    These experiments support slower-growing attention kernels as an effective design choice for post-norm transformers trained on heavy-tailed ensembles.

    Source: §5 · p. 8

    The experiments use two synthetic operators on two-dimensional measures and one training seed per configuration, with different attention implementations and precision across kernels.

    Source: §5 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

The introduction motivates measure-operator learning for heavy-tailed ensembles, where softmax’s exponential weighting can make attention integrals diverge. The paper relates this setup to prior attention and measure-learning work, then presents benchmark, ablation, and tail-sensitivity contributions.

1 Introduction

  • For the mirrored Pareto family, any nonzero linear score makes softmax’s normalizing integral diverge, motivating study of slower-growing attention kernels.The introduction frames heavy-tailed measures as relevant to applications including weather forecasting, finance, and turbulent flow, where extremes can represent important events.Source: §1 · pp. 1–2

1.1 Related work

1.2 Contributions

  • The paper constructs two identifiable operator benchmarks with closed-form targets, enabling an i.i.d.-sampling reference for evaluation.Source: §1.2 · p. 2
  • The ablations test how attention-kernel growth and data transformation affect post-norm transformer learning on heavy-tailed measures.Source: §1.2 · p. 2
  • A tail-sensitivity test measures how set divergences respond to distribution-tail differences and informs the choice of energy-distance sample estimator.Source: §1.2 · p. 2

2 Background: Attention on measures and tail integrability

Attention can be formulated as an operator on probability measures because each output is a normalized weighted average and sample sets enter through empirical measures. For heavy-tailed inputs, softmax’s population integrals may fail to be well-defined, motivating slower-growing kernels.

2 Background: Attention on measures and tail integrability

  • Each attention output is a normalized weighted average of values, so a sample set enters through its empirical measure and attention induces a map on query measures.The operator uses a query measure and a key/value measure; its measure-level definition recovers sequence attention when the measures are empirical.Source: §2 · pp. 2–3
  • For the attention operator to be well-defined, its numerator must be absolutely integrable and its denominator positive and finite.Source: §2 · p. 3
  • Softmax requires sufficient exponential-moment integrability, but the mirrored Pareto family has an infinite normalizing integral for any nonzero linear score.Finite-token attention sums remain finite for finite scores, so population integrability and finite-ensemble training are distinct questions.Source: §2 · p. 3
  • Replacing the exponential with a slower-growing positive function weakens the tail-integrability requirement and is studied for learning from finite ensembles.Source: §2 · p. 3

3 Methods

The methods formulate operator learning from finite ensembles sampled from parameterized input–output measures, then define benchmark families, transformer variants, preprocessing, and an energy-distance training objective.

3.1 Setting: operator learning on measures in 2D

  • The model learns an operator across parameter instances by mapping fixed-size input ensembles to output ensembles and minimizing a set-divergence objective.Each ensemble contains N i.i.d. samples from the corresponding input or output measure.Source: §3.1 · p. 3

3.2.1 Heavy-tailed and light-tailed marginals η

  • Benchmark measures are linear transformations of products of independent marginals, with mirrored Pareto and Gaussian families representing heavy and light tails.The Pareto parameter range gives finite first moments but infinite variance; the Gaussian control uses exchanged parameter ranges so the second axis has greater spread.Source: §3.2 · p. 3, §3.2.1 · p. 4

3.2.2 Operators and parameter distributions

  • The two benchmark operators are a sheared swap and a matrix inverse, defined through linear transformations of the input measure.The sheared swap uses a product of two shear matrices and an exchange matrix; the matrix inverse maps through B and B^-1.Source: §3.2.2 · p. 4

3.3.1 Attention kernel suite

  • A six-layer non-causal transformer approximates the operator, and its attention replaces softmax’s exponential with one of four strictly positive scalar maps.The suite spans exponential to logarithmic growth and includes three ELU-style maps.Source: §3.3 · p. 4, §3.3.1 · p. 5, Table 1 · p. 5

3.3.2 Post-norm blocks

  • The main experiments use post-norm blocks, applying layer normalization after each residual addition, with residual scaling inspired by Wang et al. [2024].The feed-forward sublayer uses the same residual structure.Source: §3.3.2 · p. 5

3.4.1 Symlog data transformation

  • Symlog preprocessing transforms input samples before the transformer and applies its inverse to the output; without preprocessing, the transform is the identity.The predictor is composed as σ^-1 ◦ fθ ◦ σ.Source: §3.4.1 · p. 5

3.4.2 Energy distance

  • Training uses the squared energy distance with p = 1, estimated from sample sets by a V-statistic.The population formula is stated for p ∈ (0, 2) and measures with finite pth absolute moments.Source: §3.4.2 · p. 5
  • The finite-sample V-statistic gives independent size-N ensembles from the same nondegenerate target a positive expected distance, so their mean or median supplies an operator- and family-specific reference.The reference averages or takes the median over trials drawing a parameter instance and two independent target ensembles.Source: §3.4.2 · p. 6

4 Experiments

The experiments assess tail sensitivity in empirical distribution comparisons and compare post-norm attention kernels with and without symlog preprocessing. On heavy-tailed benchmarks, slower-growing kernels avoid the collapse seen with untransformed softmax; symlog changes that outcome on only one task.

4.1 Tail sensitivity of the energy distance

  • The experiments compare energy distance on a testbed designed to distinguish measures with different tail decay despite equivalent high-density regions.The population energy distance is zero exactly when the measures agree.Source: §4.1 · p. 6
  • The tail-sensitivity experiment compares empirical divergences on paired one-dimensional measures that coincide below a seam and differ in the tail above it.The experiment examines how sample budget affects energy-distance sensitivity to tail differences and informs the sample budget used in the ablations.Source: §4.1 · pp. 6–7, §4 · p. 6
  • At larger sample budgets, energy distance becomes more sensitive to tail-index differences and its minimum occurs at the reference index a*=1.5 for all seam locations.The empirical 2-Wasserstein curve remains comparatively flat for a*>1.5; N=16384 is used in Section 4.2 because training at N=65536 is computationally infeasible.Source: Figure 1 · p. 6, §4.1 · p. 7

4.2 Ablation grid

  • Without symlog, post-norm softmax collapses on both heavy-tailed benchmarks, while all three slower-growing kernels avoid collapse and achieve lower energy distances.The figures cover the sheared swap and matrix inverse tasks.Source: §4.2 · p. 7, Figure 2 · p. 7, Figure 3 · p. 8
  • With symlog, softmax avoids collapse on the matrix inverse task but still collapses on the sheared swap task.For the three slower-growing kernels, symlog reduces median distance ratios to 0.767–0.863 on sheared swap and 1.135–1.167 on matrix inverse; the softmax matrix-inverse ratio is 1.120.Source: Figure 2 · p. 7, §4.2 · p. 8, Figure 3 · p. 8
  • On the Gaussian sheared-swap control, all four kernels avoid collapse and achieve similar energy distances.The separation between softmax and slower-growing kernels is specific to the heavy-tailed configurations in these experiments.Source: Figure 2 · p. 7, §4.2 · p. 8

5 Conclusion

The paper introduces two closed-form operator-learning benchmarks and reports that slower-growing kernels avoid the collapse observed with post-norm softmax on heavy-tailed tasks. It also identifies limits on how broadly the experimental kernel comparisons can be interpreted.

5 Conclusion

  • The two closed-form benchmarks show post-norm softmax collapse on both heavy-tailed tasks without preprocessing, while all three slower-growing kernels avoid collapse.Symlog enables softmax to avoid collapse on one of the two tasks.Source: §5 · p. 8
  • At larger sample budgets, energy distance distinguishes tail-index differences that are less apparent in empirical 2-Wasserstein curves.The paper presents this as a finding of its tail-sensitivity testbed.Source: §5 · p. 8
  • The experiments use two synthetic operators on two-dimensional measures and one training seed per configuration.Bootstrap intervals vary across evaluation instances at fixed trained weights; the attention implementations also differ in precision, limiting reproducibility assessment.Source: §5 · p. 8
  • Future work includes broader kernel families, higher-dimensional tasks, nonlinear sample transformations, and applications in data assimilation and inverse problems.The paper also proposes tracking attention weights and gradients around collapse and evaluating errors against known tail indices across thresholds.Source: §5 · p. 8

A Tail decay and attention integrability

Softmax attention requires strong tail conditions for its measure-level integrals to be defined, and polynomial-tail measures do not satisfy those conditions. For slower-growing positive kernels, a finite moment condition can suffice.

A Tail decay and attention integrability

  • A sufficient condition for softmax attention to be defined at every query and finite choice of weights is a faster-than-exponential radial tail bound.The passage states that this condition is stronger than light-tailedness.Source: A Tail decay and attention integrability · p. 11
  • Polynomial-tailed measures do not satisfy the stated softmax tail condition.For the mirrored Pareto family, the exponential moment is infinite for every nonzero rate.Source: A Tail decay and attention integrability · p. 11
  • Population softmax attention on the mirrored Pareto benchmark measures is undefined for any nonzero linear score.A zero score instead gives uniform weighting and requires only integrability of the values.Source: A Tail decay and attention integrability · p. 11
  • For a strictly positive kernel with growth bounded by Cg(1 + |t|^q), a finite (q + 1)st absolute moment suffices for both integrals at fixed query and finite weights.The passage gives this as a sufficient condition for the slower-growing kernel case.Source: A Tail decay and attention integrability · p. 11

B Expectation over the energy distance on sample sets of fixed size

This material formalizes how fixed-size energy-distance estimates relate to pooled estimates, defines the study’s marginal families and operator benchmarks, and states moment and preprocessing limits relevant to heavy tails.

B Expectation over the energy distance on sample sets of fixed size

  • Averaging blockwise energy distances can differ from computing the distance on pooled ensembles, as quantified by the stated identity.The identity applies to fixed ensembles, including model outputs with dependent points; when there are no off-block pairs, the distances coincide.Source: B Expectation over the energy distance… · p. 12

C The marginal families

  • For a Pareto scale λ and index a, the one-sided Pareto is supported on [λ, ∞) and has survival function (z/λ)^-a.The mirrored construction reflects and translates the Pareto branches to meet at z = λ.Source: §C.1 · p. 12

C.1 The mirrored Pareto family and the light-tailed control

  • The mirrored Pareto construction joins reflected and translated Pareto branches continuously at the seam and supports closed-form quantile transport.The study also compares Pareto and mirrored Pareto densities with a shifted Gaussian control.Source: §C.1 · pp. 12–13

C.2 The admissible range of tail indices

  • For mirrored Pareto index a, the absolute p-th moment is finite exactly when p < a; the mean exists for a > 1 and variance for a > 2.Thus indices in [1 + ϵ, 2] give a finite first moment and infinite variance.Source: §C.2 · p. 13

C.3 The tail of a Pareto under the signed-log transform

  • Symlog compresses polynomial Pareto tails to exponential tails, but does not guarantee finite population softmax integrals at unbounded learned score scales.The transform preserves sign, is near-identity around zero, and compresses large magnitudes logarithmically.Source: §C.3 · p. 14

D The operators and their parameter distributions

  • The benchmarks represent measures as linear transformations of products of independent one-dimensional marginals, with the transformation shared by each input–output pair.A fresh transformation is drawn for each trial in a training step.Source: §D.1 · p. 14

D.1 The mixing matrices and the parameter distributions ΠT

  • The sheared swap exchanges coordinates through a conjugated exchange map, while the matrix inverse applies the inverse matrix twice to the latent measure.The input and output are formed from the same transformed latent measure in each operator.Source: §D.1 · pp. 14–15

D.2 Choice of operators

  • The sheared swap depends nonlinearly on its mixing matrix, as does the matrix-inverse operator; nonlinear dependence on the latent measure is left for future work.The benchmarks are two-dimensional for direct visualization, and the Gaussian matrix-inverse case is excluded because it is unidentifiable.Source: §D.2 · p. 15

D.3 Identifiability

The section establishes which benchmark operators can be determined from their input measures, then describes the transformer, attention-kernel suite, and shared experiment settings.

D.3 Identifiability

  • The sheared swap is identifiable for both marginal families, while the matrix inverse is identifiable for mirrored Pareto but unidentifiable for shifted Gaussian marginals.For the inverse operator, the Gaussian counterexample constructs distinct outputs from the same input measure.Source: §D.3 · pp. 17–19

E.1 The model in full

  • The model is a six-layer non-causal transformer with two-dimensional states and 256-dimensional embeddings, using affine input and output maps around six blocks.Each block has attention and feed-forward sublayers with residual connections; post-norm applies layer normalization after the residual addition.Source: §E.1 · pp. 19–20

E.2 The attention kernel suite

  • The kernel suite replaces the exponential with four maps; all are C1, strictly positive, and equal 1 at score zero.Their positive-score growth ordering is exponential, quadratic, linear, then logarithmic.Source: §E.2 · p. 20, Figure 6 · p. 20

E.3 Hyperparameters

  • The main experiments vary attention kernel and data transform, while supplementary experiments also vary normalization placement.The shared model and optimizer settings are listed in Table 3.Source: Table 3 · p. 21

E.4 Compute and resource requirements

This material reports compute requirements and details a tail-sensitivity testbed that splices candidate Pareto tails onto a shared reference body and normalizes statistics across repeated trials.

E.4 Compute and resource requirements

  • The custom attention implementation takes 11.84 seconds per training step, reduced to 2.51 seconds with torch.compile; sub-exponential kernels cost more because they lack a fused kernel.Training uses query tiling and checkpointing rather than materializing the full score tensor.Source: §E.4 · p. 21, Table 4 · p. 21

F The tail-sensitivity testbed

  • The tail-sensitivity testbed compares empirical statistics across five sample budgets, with reference index 1.5 and candidates that differ in tail exponent.When the candidate exponent differs from 1.5, the corresponding population W2 is infinite even though the displayed empirical Wasserstein distances are finite.Source: F The tail-sensitivity testbed · pp. 21–22

F.1 Splicing distributions to isolate the tail

  • Splicing makes each candidate match the reference through quantile q0 and changes only the tail above the seam, whose mass is 1 − q0.The splice matches the densities at the seam; Figure 7 shows both densities and survival functions.Source: §F.1 · p. 22, Figure 7 · p. 22

F.2 Noise units

  • Each plotted statistic is shifted and scaled using repeated-trial means and standard deviations, with the reference-index candidate identified on the sweep.The sweep varies candidate index, sample budget, and seam quantile.Source: §F.2 · p. 22, §F.1 · p. 22

G Pre-norm and other results

Across the tested pre-norm kernels and transformations, the experiments report no ensemble collapse. The appendix also explains how to interpret energy-distance comparisons and the uncertainty intervals, and illustrates collapsed post-norm softmax outputs.

G Pre-norm and other results

H.1 Why a reported score can fall below the sampling reference

  • A score below the matched i.i.d. sampling reference does not mean a model output is better than the target.For deterministic outputs, the reported score can fall below that reference when the output's population energy distance is below δ/N.Source: §H.1 · p. 26

H.2 The bootstrap interval

  • The reported bootstrap intervals describe task-instance variability at fixed trained weights, not variation across training runs; nominal 95% coverage is unestablished for heavy-tailed mean ratios.Intervals containing or lying below the sampling reference correspond to bold table entries.Source: §H.2 · p. 27

I Ensemble collapse

  • For post-norm sheared swap on mirrored Pareto marginals, the collapsed softmax output concentrates near one location, unlike the target and log-ELU output.The figure shows one evaluation instance with 16,384 points.Source: Figure 12 · p. 27, I Ensemble collapse · p. 27

I.1 Training curves

On the sheared swap task with mirrored Pareto marginals, post-norm softmax eventually collapses during training even though its loss initially falls. The paper also defines a Hill-error diagnostic for comparing upper-tail estimates with target variability.

I.1 Training curves

  • Softmax loss falls to 1.8 × 10−2 without symlog and 6.3 × 10−4 with symlog, then abruptly diverges and collapses; the other three kernels continue descending.Collapse begins at approximately step 3500 without symlog and step 4500 with it, ruling out stopping training early as its explanation.Source: Figure 13 · p. 28

J Hill-based tail diagnostics

  • The Hill-based diagnostic compares each model's discrepancy from a target Hill estimate with the discrepancy between two independent target Hill estimates.It is used only for mirrored Pareto marginals; the tail fractions are 0.1 at N = 16384 and 0.01 at N = 65536.Source: J Hill-based tail diagnostics · p. 28, §J.1 · p. 28

J.1 The statistic

  • The energy distance between the distributions of candidate and reference Hill errors is descriptive: its sign alone is not a significance test and does not identify which group has smaller errors.The estimator is unbiased when trial pools are independent and identically distributed and the errors have finite first moments; finite-sample estimates may be negative.Source: §J.1 · p. 29

J.2 Results

The appendix presents Hill-error distribution discrepancies for both operators and shows their dependence on normalization and symlog transformation. The captions identify the statistic and conditions, but do not report the plotted outcomes.

J.2 Results

  • Table 7 reports E(e, ˆe) for sheared swap at N = 16384; its expectation is zero when candidate and reference Hill-error distributions agree.The two post-norm softmax entries are the collapsed runs reported in Table 5.Source: Table 7 · p. 29
  • Table 8 gives the corresponding Hill-error statistic for the matrix inverse operator.The caption does not state the table's numerical results.Source: Table 8 · p. 30
  • Figures 14–17 compare the statistic across normalization and symlog conditions for the two operators; larger population values mean greater differences between Hill-error distributions.The finite-trial estimate can be negative, and the figure axes are scaled separately.Source: Figure 14 · p. 30, Figure 15 · p. 31, Figure 16 · p. 31

Related papers

Loading 2610.00564v1…