PaperlayerView exact revision on arXiv

Source-linked AI summary

Softmax Attention on Gaussian Mixtures: Linear When It Can, Selective When It Must

Simon Gabet, Etienne Boursier, Claire Boyer

arXiv:2610.11798v1stat.MLcs.LG

TL;DR

The paper asks what query-dependent selection adds beyond the affine behavior of softmax attention on Gaussian prompts. It analyzes the infinite-prompt operator on Gaussian mixtures and shows complementary results: attention handles linear classification and uses selection for optimal denoising. Its optimization guarantees have stated scope boundaries, including noncoercivity and unresolved convergence behavior in some settings.

  • Problem

    Prior tractable analyses of softmax attention use linear kernels or Gaussian prompts, where the population operator is linear and query-dependent selection disappears.

    Source: §1 · p. 1

  • Method

    The paper studies the infinite-prompt attention operator on Gaussian mixtures and characterizes softmax attention as a query-dependent gated combination of affine experts.

    Source: §1 · pp. 1–2

  • Results

    Softmax attention learns the linear LDA solution and can represent optimal Gaussian-mixture denoising, for which linear-attention predictors cannot approach optimal error.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Within these tasks, Gaussian mixtures expose when attention's gate is redundant for a linear classifier and essential for selecting components in denoising.

    Source: §1 · p. 2, §5 · p. 11

    The landscape analysis does not establish bounded training parameters, and the degeneracy discussed for another regime determines neither convergence nor its rate.

    Source: §3.2 · p. 5, §3.3 · p. 6

Abstract

from arXiv · show

Softmax attention, at the heart of Transformers, has demonstrated remarkable capabilities. Yet its underlying mechanisms remain only partially understood. Recent theoretical work studies Gaussian prompts, where the infinite-prompt limit reduces softmax attention to a linear map, but also removes the query-dependent selection that distinguishes it from linear attention. This work studies the infinite-prompt limit of softmax attention on Gaussian mixtures, which retain the tractability of Gaussian data while introducing latent structure, multimodality, and nonlinear dependencies. We show that softmax attention can represent and learn, via gradient-based methods, optimal solutions to a range of statistical tasks, including supervised classification and denoising. Our results highlight two complementary capabilities of softmax attention: it can recover linear tasks as effectively as its simpler linear counterpart, while also exploiting query-dependent context selection to solve nonlinear tasks beyond the reach of linear attention.

1 Introduction

The paper studies infinite-prompt softmax attention on Gaussian mixtures, which preserve tractability while restoring query-dependent context selection. It develops a gated-expert account and analyzes complementary classification and denoising tasks.

1 Introduction

  • In the large-prompt limit, attention acts on the prompt distribution, and Gaussian mixtures add latent structure and query-dependent selection beyond the affine Gaussian case.The authors use this setting to remove finite-sample fluctuations while retaining closed-form computations.Source: §1 · pp. 1–2
  • On Gaussian mixtures, a softmax head combines affine experts through a query-dependent gate, whereas linear attention uses fixed component proportions.With shared covariances, the mixture map splits into a global linear map and a softmax gate that selects relevant components.Source: §1 · p. 2
  • For LDA, a single head represents the linear Bayes classifier, and gradient flow converges exponentially near regular optimal parameters; the balanced symmetric case also has a global convergence result.The effective risk has no spurious critical points, while in the original attention parameters nonoptimal critical points are strict saddles when class means are not symmetric.Source: §1 · p. 2
  • For denoising, two softmax heads reproduce the optimal Gaussian-mixture denoiser, while linear-attention predictors cannot approach its optimal error.With residual connections, one parameter set denoises optimally across cluster locations and proportions, using prompt attention to infer task-specific information.Source: §1 · p. 2
  • The paper positions its results as extending Gaussian-prompt theory to mixtures and learning the Bayes boundary directly with a softmax layer.It contrasts its noisy-context denoising results with earlier clean-context associative-memory work and notes a strict separation from linear attention.Source: §1 · p. 2

2 Population Attention on Mixtures

The section represents prompts as probability measures and studies the resulting population attention operator. On Gaussian mixtures, component-wise attention maps compete through normalization, enabling query-dependent selection absent from linear attention.

2 Population Attention on Mixtures

  • For a finite prompt, attention on its empirical measure equals softmax dot-product attention, and the population operator is its almost-sure limit for i.i.d. sub-Gaussian tokens.The results in this paper concern the limiting operator.Source: §2 · p. 3
  • On a single Gaussian, the population softmax operator is affine in the query, while on a mixture the components compete through normalization.The stated setup assumes positive-definite Gaussian covariances unless a degenerate case is specified.Source: §2 · p. 3
  • Attention on a mixture is a query-dependent convex combination of component-wise attention maps, with gates determined by alignment to the query.When a gate approaches one, the operator essentially selects that component's context.Source: §2 · p. 3
  • Linear attention combines component contexts using fixed mixture proportions, independently of the query, so it cannot separate their contributions.For shared covariances, the softmax mixture map separates into a global linear term and a query-dependent component-selection term.Source: §2 · p. 3

3 Attention-based Linear Discriminant Analysis

Under LDA, the Bayes score is affine, and one softmax head can represent it exactly; the section studies whether gradient-based training can recover it despite redundant, nonconvex parameters.

3.1 From attention to an affine–sigmoid predictor

  • One softmax head reduces to an affine predictor plus a sigmoidal correction and can realize the Bayes score; either the sigmoid vanishes or becomes constant and is absorbed into the bias.These mechanisms have different parameter-space preimages, and both need not be available for every choice of class means and priors.Source: §3.1 · p. 4

3.2 A benign but noncoercive optimization landscape

  • In effective coordinates, every critical point of the logistic risk is a global minimizer, although the model is not jointly convex.The affine predictor alone contains the Bayes score.Source: §3.2 · p. 5
  • With asymmetric class means, every attention-parameter critical point is either globally optimal or a strict saddle, so all local minima are global.The reparameterization is regular when the means are asymmetric and v21 is nonzero; the exceptional set requires a separate negative-curvature argument.Source: §3.2 · p. 5
  • Diverging parameter sequences can approach the Bayes score and optimal risk, so the critical-point landscape alone does not establish boundedness or gradient-flow convergence.As the sigmoid amplitude grows and its direction shrinks, the affine part cancels its constant and linear contributions.Source: §3.2 · p. 5

3.3 Gradient-flow convergence

  • Near the regular Bayes manifold, gradient flow converges exponentially to an optimal parameter when the means are asymmetric and the initialization is sufficiently close.The theorem also requires s(θ⋆) ≠ 0 at the starting point on that manifold.Source: §3.3 · p. 6
  • The constant-gate Bayes component has transverse flat directions, so the local Morse–Bott argument does not extend; Figure 1 shows slower decay there, but does not determine its rate.The slower decay is an empirical observation, not a convergence-rate result.Source: §3.3 · pp. 6–7
  • In the symmetric balanced case, gradient flow converges exponentially to the Bayes predictor on a Krylov-structured invariant subspace under the stated initialization conditions.The Krylov space makes the constraints invariant under the dynamics.Source: §3.3 · p. 6

4 Denoising: Selection Becomes Essential

Gaussian-mixture denoising requires both within-component shrinkage and query-dependent selection, making it a setting where the nonlinear part of softmax attention matters. The section shows exact two-head representation, a separation-dependent clustering transition under iteration, and in-context adaptation across mixture tasks.

4 Denoising: Selection Becomes Essential

  • The denoising task requires both shrinking observations within components and using query-dependent selection to infer which component generated them.This selection is where softmax attention differs from linear attention.Source: §4 · p. 7

4.1 The Bayes denoiser as a gated affine map

  • The Bayes denoiser combines within-component Gaussian shrinkage with a posterior-weighted barycenter of component means.The posterior weights are β_k(z) = P(C = k | Z = z); at high noise the barycenter approaches the prior mean, while at low noise the denoiser approaches the identity.Source: §4.1 · p. 7

4.2 Two heads are sufficient, and softmax is necessary

  • Two softmax heads exactly represent the Bayes denoiser, while no finite multi-head linear-attention predictor can approximate its optimal risk.A single head suffices if and only if γ^2 = ξ^2; numerical optimization in d = 2 also brings excess risk close to zero.Source: §4.2 · pp. 8–9

4.3 Iterating the denoiser produces a clustering transition

  • At λ ≤ 1, iterated denoising collapses inputs to the midpoint; above 1, almost every initialization converges to two stable fixed points approaching the component means as separation grows.The fixed-point equation reduces to u = tanh(λ^2u), with the transition at λ = 1.Source: §4.3 · pp. 9–10

4.4 Residual attention denoises in-context

  • A single-layer residual softmax architecture exactly adapts to every admissible mixture task without retraining as component means or weights change.The context supplies task-specific centroids, while residual heads let the output adjust the linear contraction separately from the barycentric mean.Source: §4.4 · pp. 10–11

5 Conclusion

The paper connects features of attention architectures to the statistical operations they perform across classification and denoising. It identifies finite prompts and how training discovers adaptive solutions as natural next steps.

5 Conclusion

  • Gaussian mixtures reveal complementary roles for attention: affine prediction suffices for classification, while denoising uses nonlinear context selection and residual connections to handle different mixtures.The conclusion frames these results as links between architectural features and statistical operations.Source: §5 · p. 11
  • Extending the analysis to finite prompts and explaining how training discovers adaptive solutions are identified as natural next steps.Source: §5 · p. 11

AI use statement

  • The authors report that they reviewed and verified AI-assisted work, including calculations and proofs, and take responsibility for the final content.Source: AI use statement · p. 11

A Proofs and Complements for Section 2

This appendix develops proofs and extensions for Gaussian-mixture attention, including LDA and logistic-risk facts used in the Section 3 landscape and convergence analysis. It characterizes Bayes-minimizing predictors, curvature at Bayes realizations, and when stationarity in attention or effective coordinates implies optimality.

A.1 Gaussian-mixture attention (proof of Proposition 2.1)

  • The Gaussian-mixture attention proof invokes the Gaussian moment-generating function and differentiation.The supplied passages do not include the resulting expressions.Source: §A.1 · p. 14

A.2 General mixtures (extension of Proposition 2.1)

  • The general-mixture extension defines normalized attention for a positive alignment function and states Proposition A.1 on attention over a general mixture.The supplied proposition passage does not include its formula or conclusion.Source: §A.2 · p. 14

B Proofs and Complements for Section 3

  • The Section 3 analysis separates identifying risk minimizers from proving convergence: it establishes local convergence near the regular Bayes manifold and global convergence for a prescribed initialization in a symmetric invariant subspace.The distinction matters because finite critical-point structure alone does not control trajectories at infinity.Source: B Proofs and Complements for Section 3 · p. 14

B.1 Effective LDA predictor (derivation of Equations (6)–(7))

  • Lemma B.1 considers the joint law µLDA defined by Equation (4) and states a property for every (x, y) ∈ Rd × {−1, 0, 1}; σ is called the sigmoid function.Source: §B.1 · p. 15

B.2 Bayes realizability (preparation for Theorems 3.1 and 3.3)

  • Attention parameters can realize the LDA Bayes score exactly, so the attention class attains Bayes logistic risk.The Bayes classifier thresholds the affine Bayes score; the attention realization matches that score for every x.Source: §B.2 · p. 16

B.3 Population logistic-risk minimizer (preparation for Theorem 3.1)

  • The population logistic risk over measurable scores is uniquely minimized almost surely by the Bayes predictor, whereas the quadratic-risk target need not belong to the attention class.Although thresholding the quadratic target yields the same Bayes classifier, minimizing quadratic risk within the attention class need not recover its global minimizer or optimal classifier.Source: §B.3 · pp. 17–18

B.4 Differential identities (tools for the LDA landscape and convergence proofs)

  • At a Bayes realization, the Hessian ignores perturbations in s; its kernel also depends on whether s⋆ is zero.When s⋆ ≠ 0, the kernel forces δw, δb, and δr to zero; when s⋆ = 0, it includes vectors of the form (0, −δr σ(s0), δr, δs).Source: §B.4 · pp. 21–22
  • At v21 = 0, the differential of Φ loses rank, so original-risk stationarity may not transfer; under m1 + m−1 ≠ 0, the reduced w-gradient is non-vanishing at the specified critical points.This non-vanishing direction is used to construct strict saddles.Source: §B.4 · p. 23

B.5 Optimization landscape (proofs of Theorems 3.1 and 3.3, Lemma 3.2, and Proposition 3.4)

  • The optimization-landscape proof proceeds through effective minimizers and stationarity, then attention-parameter preimages and regular-point or strict-saddle arguments; loss of compactness means these results alone do not establish global convergence.The appendix separately proves the stated landscape theorems, lemma, and proposition.Source: §B.5 · p. 24

B.5.1 Bayes minimizers (Theorems 3.1 and 3.3, minimizer claims)

  • The effective Bayes minimizers arise either by removing the sigmoid with r = 0 or making it constant with s = 0, and attention minimizers are their preimages under Φ.The preimage description distinguishes optimal effective parameters from those that are actually attainable in attention coordinates.Source: §B.5.1 · pp. 24–25
  • Although the reduced risk is not jointly convex, fixing s yields a convex problem in (w, b, r) that contains the Bayes predictor, so stationarity in those coefficients and amplitude forces global optimality.The Bayes predictor is contained in the restricted problem by choosing r = 0.Source: §B.5.1 · p. 25

B.5.2 Effective stationary points (completion of Theorem 3.1)

The effective risk has no nonglobal critical points, and the parameter map is regular precisely under stated asymmetry and nonzero-factor conditions. This regularity lets the effective-risk result rule out nonglobal critical points in the corresponding attention-parameter region.

B.5.2 Effective stationary points (completion of Theorem 3.1)

  • Every critical point of the effective risk is a global minimizer.The proof uses convexity of the risk with one coordinate fixed.Source: §B.5.2 · pp. 25–26

B.5.3 Regularity of the parameter map (proof of Lemma 3.2)

  • The adjoint differential of the parameter map is injective exactly when the class means are asymmetric and v21 is nonzero.Here, asymmetry means m1 ≠ −m−1.Source: §B.5.3 · p. 26
  • The proof establishes regularity by showing that arbitrary effective-coordinate variations can be obtained through the differential of the parameter map.This makes the differential surjective under the stated conditions.Source: §B.5.3 · pp. 26–27
  • Away from v21 = 0, this regularity transfers stationarity to the effective risk and rules out nonglobal critical points.The transfer argument applies in the asymmetric-mean case.Source: §B.5.3 · pp. 26–27

B.5.4 Critical points and strict saddles (completion of Theorem 3.3)

With asymmetric class means, every critical point is either globally minimizing or a strict saddle. The argument separates regular points from the singular set, where paired curvature directions establish the strict-saddle property.

B.5.4 Critical points and strict saddles (completion of Theorem 3.3)

  • At critical points with v21 ≠ 0, the parameter-map transfer argument implies global optimality.This result assumes m1 ≠ −m−1.Source: §B.5.4 · pp. 27–28
  • At the exceptional set v21 = 0, coupling factor variations while keeping the gate fixed to first order produces directions of opposite curvature.These critical points are strict saddles, not additional local minima.Source: §B.5.4 · p. 28
  • With asymmetric class means, every critical point is globally minimizing or a strict saddle, so all local minima are global.The proof establishes global optimality outside v21 = 0 and negative curvature at nonoptimal critical points.Source: §B.5.4 · pp. 27–29
  • The classification of finite critical points does not bound gradient-flow trajectories because near-optimal directions can escape to infinity.The escaping sequence is not itself a gradient-flow trajectory.Source: §B.5.4 · p. 29

B.5.5 Minimizing sequences at infinity (proof of Proposition 3.4)

The paper constructs effective parameters that diverge while their predictors approach the optimal predictor and their risk approaches the optimal risk. This escaping sequence does not establish that gradient-flow trajectories diverge.

B.5.5 Minimizing sequences at infinity (proof of Proposition 3.4)

  • As R → ∞, the constructed predictors converge to f⋆ in L1(µX) and their effective risk converges to R⋆.The construction yields near-optimal predictors along parameters that escape to infinity.Source: §B.5.5 · p. 29
  • The escaping near-optimal sequence does not bound gradient-flow trajectories or prove that they diverge.The paper distinguishes this construction from a gradient-flow trajectory.Source: §B.5.5 · p. 30
  • The sequence shows why the absence of nonglobal local minima alone is insufficient for a global convergence argument.The missing ingredient is a bound on the gradient-flow trajectory.Source: §B.5.5 · p. 30

B.6 Normal curvature (preparation for Theorem 3.5)

The risk has a critical manifold of equivalent optimal parameters, so its Hessian is flat along that manifold. The section establishes positive curvature in normal directions, the local geometry used to study gradient-flow convergence.

B.6 Normal curvature (preparation for Theorem 3.5)

  • The optimal parameters are not isolated, making positive definiteness of the full Hessian unavailable and unnecessary.The relevant geometric question is whether any directions beyond the manifold’s tangent directions are flat.Source: §B.6 · p. 30
  • At regular optimal parameters, the risk has the Morse–Bott property: the critical manifold’s tangent space coincides with the Hessian kernel.The section assumes asymmetric class means and considers points with s(θ⋆) ≠ 0.Source: §B.6 · p. 30
  • The Hessian is positive definite normal to the critical manifold, while the risk is flat along the manifold of minimizers.Thus the only flat directions are tangent to the manifold.Source: §B.6 · p. 30
  • The Morse–Bott local normal form is a main ingredient in proving local gradient-flow convergence toward the manifold of global minimizers.The risk is locally equivalent, up to a diffeomorphism, to a quadratic function in normal directions.Source: §B.6 · p. 30

B.7 Local gradient-flow convergence (proof of Theorem 3.5)

Near a regular minimizer component, the proof uses Morse–Bott coordinates to separate tangent and normal directions and establishes local exponential gradient-flow convergence to a single Bayes-optimal parameter.

B.7 Local gradient-flow convergence (proof of Theorem 3.5)

  • Morse–Bott coordinates split parameters into tangent directions, where risk is constant, and normal directions, where excess risk is quadratic.The coordinate change induces a metric that must be accounted for to analyze the original Euclidean flow.Source: §B.7 · pp. 33–34
  • From a sufficiently small neighborhood of a regular critical component, Euclidean gradient flow remains local and converges exponentially to a single point in the minimizer set.The theorem assumes asymmetric class means, m1 ≠ −m−1.Source: §B.7 · pp. 33, 36
  • On a compact coordinate neighborhood, normal displacement decays exponentially while bounded tangential displacement prevents the trajectory from leaving the neighborhood.The normal estimate is dV(yt)/dt ≤ −4λV(yt); the tangential curve has finite length and converges.Source: §B.7 · pp. 35–36

B.8 Constant-gate geometry (complement to Theorem 3.5)

These results characterize when softmax attention realizes Bayes denoisers and how representation and initialization shape learning, including constant-gate Hessian-flat directions and conditions for convergence from prescribed Krylov initializations.

B.8 Constant-gate geometry (complement to Theorem 3.5)

  • On E2 with r⋆≠0, the Hessian kernel has dimension d + 1 while the tangent space is one-dimensional, so the constant-gate component is not Morse–Bott.Under m1 ≠ −m−1, the excess flat directions also transfer to attention coordinates; at r⋆=0, the effective branches meet and require separate care.Source: §B.8 · pp. 37–38

B.9 Reduced-risk estimates and conserved balance (preparation for Theorem 3.6)

  • In the symmetric balanced setting, the reduced linear logistic risk is coercive and strongly convex along the effective trajectory, while conserved factor balance prevents the effective dynamics from stalling.These controls supplement convexity in the effective predictor, since training still evolves two factors.Source: §B.9 · pp. 38–40

B.10 Feasibility of the Krylov initialization (complement to Theorem 3.6)

  • For symmetric positive-definite Γ, K = span{m, Γm, . . . , Γ^(d−1)m} equals R^d if and only if Γ has d distinct eigenvalues and m projects nontrivially onto every eigenspace.Source: §B.10 · p. 40

B.11 Invariance of the Krylov subspace (preparation for Theorem 3.6)

  • In the symmetric balanced setting, gradient flow initialized in IK remains there for all t ≥ 0, preserving a linear predictor throughout training.The Krylov constraint is preserved because K and K⊥ are Γ-invariant, together with U11PK = 0.Source: §B.11 · p. 42

B.12 Global convergence on the invariant subspace (proof of Theorem 3.6)

  • Under the prescribed nonzero Krylov initialization, the invariant-subspace flow yields exponential convergence of risk and predictor, followed by convergence of the factors and original attention parameters.The proof uses a conserved balance to bound the effective-flow preconditioner and coercivity to keep the predictor in a region of uniform positive curvature.Source: §B.12 · pp. 43–45

B.13 LDA experiments (illustrations of Theorems 3.5 and 3.6)

  • With balanced classes and opposite means, the specified initialization rapidly decreased excess risk toward a small value, unlike Gaussian initialization in the displayed runs.The initialization does not require m or Γ and does not enforce v21(0) ∈ K⊥, suggesting the favorable behavior can extend beyond that invariant subspace.Source: §B.13 · p. 45, Figure 4 · p. 46, Figure 5 · p. 46

C.1 Bayes denoiser (derivation of Equation (10))

  • The Bayes denoiser is the conditional mean, formed by posterior-weighting each component’s Gaussian conditional mean.As noise grows without bound, the weights approach the priors and the denoiser approaches the mixture mean; as noise vanishes for positive-definite Γ, it approaches z.Source: §C.1 · pp. 47–48

C.2 Expressivity and risk separation (complete proof of Theorem 4.1)

  • Two heads can realize the Bayes denoiser, while a single head requires γ² = ξ²; the theorem also establishes a strict excess-risk gap.The proof constructs two heads, including for opposite means, and shows the gap is strict rather than merely a failure of exact linear representation.Source: §C.2 · p. 48

C.3 Degenerate mixtures (complement to Theorem 4.1)

  • For degenerate two-component mixtures with equal-norm distinct means, two heads realize the Bayes denoiser; in d = 2, one head fails for linearly independent means.When the means are opposite, one head suffices.Source: §C.3 · pp. 48–49

C.4 Denoising-to-clustering transition (proof of Proposition 4.3)

The proof reduces denoiser iteration to a one-dimensional map whose behavior changes at a separation threshold. It also establishes when one attention head cannot represent the optimal denoiser and constructs an exact two-head representation.

C.4 Denoising-to-clustering transition (proof of Proposition 4.3)

  • For λ ≤ 1, iteration converges to the midpoint c; for λ > 1, it converges to one of two outer fixed points according to the sign of the initial projection.As λ grows without bound, the outer fixed points approach the component means.Source: §C.4 · p. 49

C.5 Single-head necessity (Theorem 4.1, single-head claim)

  • When γ² ≠ ξ², no single attention head can represent the optimal denoiser in the stated equal-norm, isotropic two-component setting.The single-head constraints require γ² = ξ².Source: §C.5 · pp. 50–51

C.6 Two-head construction (Theorem 4.1, exact representation)

  • With d ≥ 2, positive mixture weights, distinct equal-norm means, and isotropic covariance, two heads exactly represent the optimal denoiser.The construction's sufficient conditions ensure exact representation.Source: §C.6 · p. 51

C.7 Strict risk gap (Theorem 4.1, separation claim)

In the nontrivial two-component setting, a two-head softmax network reaches Bayes risk while linear attention cannot match it. For residual in-context denoising, the construction is optimal for every task under the stated finite-risk condition.

C.7 Strict risk gap (Theorem 4.1, separation claim)

  • A two-head softmax network attains Bayes risk, while linear attention cannot achieve the same risk in the nontrivial two-component setting.This establishes a strict risk gap between the attention classes in that setting.Source: §C.7 · p. 53
  • Even with arbitrary heads, a network composed of linear attention layers remains linear in z.Linear combinations and compositions of linear maps preserve this constraint.Source: §C.7 · p. 53

C.8 Residual in-context denoising (proof of Proposition 4.4)

  • For residual in-context denoising, the constructed predictor is optimal for any task distribution with finite risk.Its matrices depend only on the shared variances; prompt-specific means and weights enter through the attention operator.Source: §C.8 · p. 54
Loading 2610.11798v1…