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

Stability of Measure-to-Measure Transformers on Sub-Gaussian Data
Oct 6, 2026Frank Cole, Nicholas H. Nelsen, Takashi FuruyaTransformer Statistical Learning Theory
UCLA · UT Austin · Doshisha University, RIKEN AIP
Attention Kernels for Learning Maps Between Heavy-Tailed Measures
Sep 30, 2026Kailen Hargenrader, Edoardo Calvello, Bohan ChenSoftmax Attention Attention Mechanisms
Computing and Mathematical Sciences California Institute of Technology · Lawrence Berkeley National Laboratory University of California, Berkeley ICSI
Universal Approximation of Measure-to-Measure Operators by Pushforwards
Sep 28, 2026Takashi Furuya, Nicholas H. Nelsen, Frank ColeTransformer Universal Approximation
Doshisha University, RIKEN AIP · UT Austin · UCLA

Oct 6, 2026J/K move · Enter open · S save

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

Frank Cole, Nicholas H. Nelsen, Takashi Furuya

UCLA · UT Austin · Doshisha University, RIKEN AIP

Transformer Statistical Learning Theory

Open paper AskPDF arXiv

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.

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