A Theory of Diversity for Random Matrices with Applications to In-Context Learning of Schrödinger Equations

TL;DR

This work introduces a notion of diversity for random matrix distributions via the triviality of a matrix centralizer, linking this property to the transferability of transformer-based in-context learning for discretized Schrödinger operators. It proves two main results: (i) an augmented-sample centralizer-triviality bound for matrices of the form $\mathbf{A}=\mathbf{K}+\mathbf{V}$, and (ii) a vanilla-sample bound under diagonal randomness, both yielding exponential-in-$N$ probability guarantees. The results are specialized to finite difference and finite element discretizations of random Schrödinger operators, enabling rigorous generalization guarantees for transformers solving linear PDEs and providing empirical validation through 1D and 2D experiments that show robust in-domain and out-of-domain generalization. The work highlights a blessing of dimensionality phenomenon in the finite difference setting and outlines open directions for extending the theory to broader operators and improving dependence on dimension, with potential impact on data-efficient scientific ML models. All mathematical statements are presented with explicit $\,p$- and $d$-dependent bounds, and numerical evidence supports the theoretical claims about diversity and generalization across discretizations and potential distributions.

Abstract

We address the following question: given a collection $\{\mathbf{A}^{(1)}, \dots, \mathbf{A}^{(N)}\}$ of independent $d \times d$ random matrices drawn from a common distribution $\mathbb{P}$, what is the probability that the centralizer of $\{\mathbf{A}^{(1)}, \dots, \mathbf{A}^{(N)}\}$ is trivial? We provide lower bounds on this probability in terms of the sample size $N$ and the dimension $d$ for several families of random matrices which arise from the discretization of linear Schrödinger operators with random potentials. When combined with recent work on machine learning theory, our results provide guarantees on the generalization ability of transformer-based neural networks for in-context learning of Schrödinger equations.

Figures (6)

• Figure 1: The probability that the centralizer of discrete random 1D Schrödinger matrices—whose potential entries are sampled from ($\mathrm{Ber}_p\{1,2\}$)—is trivial, shown as a function of the sample size $N$. Different colors correspond to different values of the Bernoulli parameter $p$.
• Figure 2: In-domain generalization MSE with respect to inference prompt length $m$ with $D=1, d=10, N=40000$ training tasks, $n=40000$ training prompt length, $M=1000$ testing tasks. Both training and testing are performed with the finite difference method. The four plots are overlaid with various distributions for potential function $V(x)$: piecewise constant, and lognormal with parameters $(\alpha,\beta) = (0,1),(1,2),(2,3)$. The solid black line corresponds to the $O(m^{-1})$ scaling rate.
• Figure 3: Out-of-domain generalization shifted relative error with respect to inference prompt length $m$ with $d=10$, $N=40000$ training tasks, $n=40000$ training prompt length, $M=10000$ testing tasks. The transparent black lines correspond to the $O(m^{-1})$ scaling rate. Top left: we train with FD and test with FEM. Bottom left: we train with FEM and test with FD. Top middle: we train with $V$ as a piecewise constant, and test with $V$ from a lognormal random field with $(\alpha,\beta)=(2,2)$. Bottom middle: we train with $V$ from a lognormal random field with $(\alpha,\beta)=(2,2)$, and test with $V$ as a piecewise constant. Top right: we train with $V$ from a lognormal random field with $(\alpha,\beta)=(2,2)$ and test with $V$ from a lognormal random field with varied $\beta$. Bottom right: we train with $V$ from a lognormal random field with $(\alpha,\beta)=(2,2)$ and test with $V$ from a lognormal random field with varied $\alpha$.
• Figure 4: The probability that the centralizer of discrete random 2D Schrödinger matrices—whose potential entries are sampled from ($\mathrm{Ber}_p\{1,2\}$)—is trivial, shown as a function of the sample size $N$. Different colors correspond to different values of the Bernoulli parameter $p$. The dimension of the linear systems here is $d = 9.$
• Figure 5: In-domain generalization MSE with respect to inference prompt length $m$ with $D=2, d=25, N=40000$ training tasks, $n=40000$ training prompt length, $M=1000$ testing tasks. Both training and testing is performed with the finite difference method. 4 plots are overlaid with various distributions for potential function $V(x)$: piecewise constant, and lognormal with parameters $(\alpha,\beta) = (0,1),(1,2),(2,3)$. The solid black line corresponds to the $O(-1)$ scaling rate.

Theorems & Definitions (15)

• Theorem 1
• Theorem 2
• proof
• Remark 1
• Theorem 3
• proof
• Remark 2
• Theorem 4
• Theorem 5
• Lemma 1