Measure Neural Mappings

Ricardo Baptista; Edoardo Calvello; Bohan Chen; Andrew Stuart; Urbain Vaes

cs.LGOct 12, 2026

Measure Neural Mappings

Authors: Ricardo Baptista, Edoardo Calvello, Bohan Chen, Andrew Stuart, Urbain Vaes

Organizations: University of Toronto · Lawrence Berkeley National Laboratory · California Institute of Technology · Inria Paris, ´Ecole des Ponts

Abstract

We introduce and define measure neural mappings (MNMs), which combine measure-dependent neural transport with ensemble consistency for learning operators on probability measures. Ensemble consistency requires outputs from growing ensembles sampled from an input law to converge to the specified population output in probability. We establish quantitative continuum limits for attention on noncompact laws under strong sub-Gaussian and tail-integrability assumptions. These laws accommodate unbounded state spaces in scientific machine learning, including Bayesian inference and uncertainty quantification. For fixed self-attention transformers, the expected output Wasserstein error is bounded by the empirical input Wasserstein error plus an additional $N^{-1/2}$ attention term. These self-attention transformers are MNMs under the stated assumptions: pre-norm architectures follow from compact stability after normalization, while post-norm architectures are covered under the stated tail conditions. We prove corresponding cross-attention estimates and extend the results to multihead attention and grouped transformer evaluation, with convergence as the group size grows. For Set Transformers viewed as measure-to-vector mappings, we obtain an $N^{-1/2}$ bound on the expected Euclidean output error under the stated assumptions. Experiments on seven operators in dimensions $d=1,2,3,5,10$ support the predicted particle-count dependence, using one-dimensional Wasserstein computations and Sinkhorn approximations in higher dimensions. The observed Wasserstein and sliced Wasserstein decay rates are close to their respective theoretical sampling benchmarks.

Figures & tables

Panel of Figure 1 — Figure 1: Post-norm self-attention stack ( $L\geq 2$ ). Both rows use the same fixed blocks. The first step uses Theorem 12 and its particle-coupling proof. Layer normalization gives a common deterministic compact support (Lemma 33 ), so later steps use compact stability (Proposition 36 ) without independence. Dashed arrows denote convergence in expected $W_{p}$ as $N\to\infty$ .

Figure 2: Convergence to the learned reference for all seven operators. Each panel shows

W_{2}

d=1

and the Sinkhorn approximation in

d=2,3,5,10

; torus doubling uses periodic ground distance. Curves use input-measure seed

0

and average

R_{d}

sampling replicates as specified in Section 4.1 , with one sampling standard error conditional on the shared reference shown by shading. Dashed guides have shapes

r_{d}(N)

with fitted multiplicative constants. The horizontal axis is shared; vertical limits and ticks are chosen separately for each panel.

Operator	$d=1$	$d=2$	$d=3$	$d=5$	$d=10$
Torus doubling	$0.4806\pm 0.0013$	$0.4502\pm 0.0017$	$0.3543\pm 0.0018$	$0.1941\pm 0.0010$	$0.1035\pm 0.0002$
Gaussian rescaling	$0.4904\pm 0.0080$	$0.4325\pm 0.0023$	$0.3167\pm 0.0008$	$0.1904\pm 0.0004$	$0.1049\pm 0.0001$
Gaussian inversion	$0.4609\pm 0.0171$	$0.4465\pm 0.0019$	$0.3216\pm 0.0016$	$0.1948\pm 0.0007$	$0.1046\pm 0.0012$
Mixture rescaling	$0.4950\pm 0.0057$	$0.4421\pm 0.0015$	$0.3202\pm 0.0005$	$0.1914\pm 0.0004$	$0.1051\pm 0.0010$
Mixture inversion	$0.4630\pm 0.0183$	$0.4560\pm 0.0083$	$0.3228\pm 0.0012$	$0.2038\pm 0.0012$	$0.1056\pm 0.0002$
Gaussian tilting	$0.4911\pm 0.0089$	$0.4410\pm 0.0009$	$0.3158\pm 0.0002$	$0.1966\pm 0.0003$	$0.1084\pm 0.0004$

Table 1: Numerical decay exponents for

W_{2}

(

d=1

) and the Sinkhorn approximation (

d=2,3,5,10

). Each measure-level rate averages

R_{d}

separately fitted sampling exponents; entries report mean

\pm

SE across the three input measures. The last row is the

W_{2}

comparison on the finite grid, not a theorem for fixed-regularization Sinkhorn divergence.

Figure 3: Sliced comparison with the learned reference in dimensions

d=1,2,3,5,10

. The

d=1

curves retain

W_{2}

(periodic

W_{2}

for torus doubling). Higher-dimensional torus curves use the sine–cosine embedding; the others use Euclidean projections. Curves show input-measure seed

0

, averaged over

R_{d}

sampling replicates, with one sampling standard error conditional on the shared reference. Dashed guides have slope

-1/2

. Each panel has its own vertical limits and ticks.

Operator	$d=1$	$d=2$	$d=3$	$d=5$	$d=10$
Torus doubling	$0.4806\pm 0.0013$	$0.4806\pm 0.0233$	$0.4964\pm 0.0082$	$0.4998\pm 0.0025$	$0.4766\pm 0.0039$
Gaussian rescaling	$0.4904\pm 0.0080$	$0.4916\pm 0.0051$	$0.5084\pm 0.0012$	$0.5137\pm 0.0044$	$0.4874\pm 0.0047$
Gaussian inversion	$0.4609\pm 0.0171$	$0.4885\pm 0.0034$	$0.4985\pm 0.0057$	$0.4996\pm 0.0074$	$0.4962\pm 0.0047$
Mixture rescaling	$0.4950\pm 0.0057$	$0.5068\pm 0.0079$	$0.5029\pm 0.0130$	$0.4824\pm 0.0069$	$0.4866\pm 0.0035$
Mixture inversion	$0.4630\pm 0.0183$	$0.4925\pm 0.0128$	$0.4999\pm 0.0061$	$0.4728\pm 0.0129$	$0.4986\pm 0.0120$
Gaussian tilting	$0.4911\pm 0.0089$	$0.4936\pm 0.0056$	$0.5012\pm 0.0045$	$0.4913\pm 0.0061$	$0.4901\pm 0.0081$

Table 2: Sliced-comparison decay exponents, reported as mean

\pm

SE across three measure-level rates, each averaging

R_{d}

per-sampling fits. The

d=1

entries use

W_{2}

and coincide with Table 1 .

Figure 4: Approximation to the true operator using

W_{2}

d=1

and Sinkhorn in

d=2,3,5,10

. All curves use input-measure seed

0

and the mean over

R_{d}

sampling replicates, with one sampling standard error. The true-output law is fixed; independent target particles are drawn in each sampling replicate. Panels have individual vertical scales.

Figure 5: Approximation to the true operator using sliced Wasserstein in

d=2,3,5,10

, retaining

W_{2}

d=1

. Seed selection, averaging and bands match Figure 4 .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Family	$d$	Width	Blocks	Heads	Steps	Batch size
Torus	1	256	6	8	10,000	8
Torus	2, 3	384	8	8	20,000	4
Euclidean	1	192	6	6	15,000	8
Euclidean	2	256	8	8	25,000	6
Euclidean	3	384	10	12	35,000	4

Appendix

Table 3: Architecture and training parameters for

d=1,2,3

. The Euclidean rows apply to all six Euclidean operators. Batch size counts independent particle ensembles.

Operator	Width	Blocks	Steps	Peak learning rate
Torus doubling	384 / 512	8 / 8	30,000 / 50,000	$2\!\times\!10^{-4}$ / $1.5\!\times\!10^{-4}$
Gaussian rescaling	512 / 512	10 / 12	45,000 / 65,000	$1.5\!\times\!10^{-4}$ / $10^{-4}$
Gaussian inversion	512 / 512	10 / 12	50,000 / 80,000	$10^{-4}$ / $7.5\!\times\!10^{-5}$
Mixture rescaling	512 / 512	10 / 12	55,000 / 85,000	$10^{-4}$ / $7.5\!\times\!10^{-5}$
Mixture inversion	512 / 512	10 / 12	70,000 / 110,000	$10^{-4}$ / $5\!\times\!10^{-5}$
Gaussian tilting	512 / 512	10 / 12	50,000 / 80,000	$5\!\times\!10^{-5}$ / $3\!\times\!10^{-5}$

Appendix

Table 4: Architecture and training parameters for the H100 experiments. Every paired entry lists

d=5

d=10

. All models use eight attention heads, batch size four, and feed-forward width four times the model width.

	Training (min)			$N=65{,}536$ (s)			$M=500{,}000$ (s)
Operator	$d=1$	$d=2$	$d=3$	$d=1$	$d=2$	$d=3$	$d=1$	$d=2$	$d=3$
Torus doubling	2.6	6.0	6.0	2.95	4.89	4.89	170.1	283.9	283.3
Gaussian rescaling	2.9	6.3	12.7	2.55	3.92	7.36	127.6	227.1	426.0
Gaussian inversion	2.9	6.5	12.8	2.56	3.92	7.37	127.3	227.2	425.9
Mixture rescaling	3.2	7.0	13.5	2.53	3.92	7.36	127.5	227.1	425.4
Mixture inversion	3.2	7.0	13.6	2.53	3.91	7.36	127.4	226.4	425.3

Appendix

Table 5: RTX 4080 SUPER training and inference times for

d=1,2,3

. Training times are minutes; median forward times are seconds. For the first five operators, training times are measured in independent full-budget reruns with the original configurations and seed

7

; the two tilting operators retain their original training timers. All inference times use the original evaluated checkpoints. References are computed once per input law and reused across all sampling replicates.

	Training (min)		$N=65{,}536$ (s)		$M=500{,}000$ (s)
Operator	$d=5$	$d=10$	$d=5$	$d=10$	$d=5$	$d=10$
Torus doubling	8.7	16.7	2.08	2.13	127.3	128.8
Gaussian rescaling	15.8	29.3	2.66	3.20	161.8	194.8
Gaussian inversion	17.2	36.6	2.67	3.21	162.3	193.9
Mixture rescaling	19.5	39.6	2.65	3.19	160.0	192.3
Mixture inversion	26.2	53.3	2.68	3.17	162.6	191.4

Appendix

Table 6: Measured H100 training and inference times. Each pair of columns lists

d=5

and

d=10

. Training times are minutes; median forward times are seconds. References are computed once per input law and reused across its five sampling replicates.

Setting	Blur	Scaling	Truncate	Exponent	Max. error change
Baseline	$0.01$	$0.8$	$5$	$0.3075\pm 0.0084$	$0.000\%$
Smaller blur	$0.005$	$0.8$	$5$	$0.3069\pm 0.0083$	$0.397\%$
Larger blur	$0.02$	$0.8$	$5$	$0.3106\pm 0.0086$	$1.672\%$
Largest blur	$0.05$	$0.8$	$5$	$0.3334\pm 0.0098$	$12.080\%$
Slower annealing	$0.01$	$0.9$	$5$	$0.3049\pm 0.0084$	$1.210\%$
Shorter truncation	$0.01$	$0.8$	$3$	$0.3075\pm 0.0084$	$<0.001\%$

Appendix

Table 7: Paired Sinkhorn ablation for cubic tilting in

d=3

and input-measure seed

0

. Exponents are mean

\pm

sample standard deviation of three per-realization fits. The last column gives the maximum absolute relative change from the baseline over all

15

paired errors. All settings use the multiscale backend and the same debiased distance.

Figure 6: Mean reference errors with one sampling standard error for

R=5,10,25,50,100

in panels (a)–(c),

R=5,10,25

in panel (d), and

R=5,10

in panel (e). Dashed

N^{-1/2}

guides are fitted vertically to the curve with the largest

R

Operator Learning

Operator learning studies models that map functions or fields to functions, enabling fast prediction across scientific and dynamical systems.

See all 99 papers

Oct 6, 2026

Stability of Measure-to-Measure Transformers on Sub-Gaussian Data

Transformers have exhibited impressive empirical success across various domains, but their theoretical foundations remain less developed. This work constitutes a mathematical study of the measure-to-measure operators defined by transformers. We show that transformers map sub-Gaussian inputs to sub-Gaussian outputs; this ensures that taking arbitrary-length compositions of the softmax operator is well-defined. We then show that transformers are Hölder continuous with respect to the 1-Wasserstein distance on appropriate spaces of sub-Gaussian inputs. This allows us to establish estimates on the error propagation along a transformer between a sub-Gaussian input and its empirical approximation. We also study a mean-field analog of the cross-attention mechanism, which is an operator from a pair of probability measures to a single probability measure. We show that cross-attention exhibits different Hölder regularity and sample-complexity in its two input arguments. Last, we apply our results to deduce approximation guarantees for measure-to-measure transformers. Together, these results provide a firm stability and finite-sample theory for transformers on sub-Gaussian data.

Frank Cole, Nicholas H. Nelsen, Takashi Furuya

UCLA · UT Austin · Doshisha University, RIKEN AIP

Transformer Statistical Learning Theory

Ask PDF

Sep 30, 2026

Attention Kernels for Learning Maps Between Heavy-Tailed Measures

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.

Kailen Hargenrader, Edoardo Calvello, Bohan Chen

Computing and Mathematical Sciences California Institute of Technology · Lawrence Berkeley National Laboratory University of California, Berkeley ICSI

Softmax Attention Attention Mechanisms Heavy-Tailed Distributions Operator Learning

Ask PDF

Sep 28, 2026

Universal Approximation of Measure-to-Measure Operators by Pushforwards

Many learning tasks map an input distribution to an output distribution. A natural way to model such an operator is to transform each input sample using a continuous function that may depend on the entire input distribution, and then take the distribution of the transformed samples. This defines a measure-dependent pushforward model and includes measure-theoretic formulations of transformers. We ask when such models can approximate arbitrary continuous operators between spaces of probability measures. We first show that universal approximation fails when atomic inputs are allowed: some continuous measure-to-measure operators that split or redistribute atomic mass cannot be approximated arbitrarily well by deterministic pushforward models. We then introduce the uniform level set condition, which requires a continuous measure-dependent scalarization whose shrinking level set neighborhoods carry uniformly vanishing mass over the input family. This condition is satisfied, in particular, by compact families of absolutely continuous measures. On every compact family satisfying this condition, we prove that any continuous measure-to-measure operator with outputs of finite

p

-th moment can be uniformly approximated, in the

p

-Wasserstein distance, by continuous measure-dependent pushforwards. Combining our theorem with existing approximation results for measure-dependent in-context maps yields universal approximation by measure-theoretic transformers. We also extend the framework to continuously-varying source measures, yielding a corresponding universality result for a class of pushforward models that are closely aligned with cross-attention architectures.

Takashi Furuya, Nicholas H. Nelsen, Frank Cole

Doshisha University, RIKEN AIP · UT Austin · UCLA

Transformer Universal Approximation Operator Learning

Ask PDF

Abstract

Operator Learning

Explore similar work

Stability of Measure-to-Measure Transformers on Sub-Gaussian Data

Attention Kernels for Learning Maps Between Heavy-Tailed Measures

Universal Approximation of Measure-to-Measure Operators by Pushforwards