Mechanistic Evidence for Spectral Structures in Prior-Data Fitted Networks

Kaustubh Sharma; Srijan Tiwari; Ojasva Nema; Parikshit Pareek

cs.LGJan 29, 2026

Mechanistic Evidence for Spectral Structures in Prior-Data Fitted Networks

Authors: Kaustubh Sharma, Srijan Tiwari, Ojasva Nema, Parikshit Pareek

Organizations: Department of Electrical Engineering, Indian Institute of Technology Roorkee (IIT Roorkee), Uttarakhand, India · Department of Metallurgical and Materials Engineering, Indian Institute of Technology Roorkee (IIT Roorkee), Uttarakhand, India

Abstract

Prior-Data Fitted Networks (PFNs) perform approximate Bayesian inference in a single forward pass, and tabular foundation models (TFMs) built on them are now widely used. To understand what networks infer internally, recent mechanistic studies of TFMs locate where predictions form, but treat these models as tabular predictors rather than as PFNs. It therefore remains unknown whether PFNs represent the spectral content of their context, the quantity that specifies a stationary kernel, and whether this content can be read out as an explicit kernel. We answer both questions. First, across seven PFNs, including four pretrained TFMs and a model trained only on a decision-tree prior, a linear probe on the residual stream recovers the frequency of the context with $R^2 \geq 0.95$ . This structure is led by a single principal direction. Second, activation and subspace patching show that the network uses the structure through a low-dimensional subspace, where a few spectral directions move predictions far more than random ones. This holds even for the decision-tree model, so a spectral training prior is not required. On real datasets with up to 499 features, 64 of the 192 directions of TabPFN, chosen without labels, carry 85 to 95% of the causal effect of the context in all but one pair. Third, we introduce a Filter Bank Decoder that turns frozen PFN representations into an explicit stationary kernel through Bochner's theorem. Without any test-time optimization, the decoded kernel supports Gaussian process regression competitive with deep kernel learning and random Fourier features at about $200\times$ lower cost. PFN latents therefore hold spectral structure that is causally used and recoverable as a portable kernel.

Figures & tables

Model	Prior	$\rho_{\Delta}$	$\|r\|_{\text{PC}_{0}}$	Linear $R^{2}$ ( $f$ )	Linear $R^{2}$ ( $a$ )	$\max\|\Delta\|$
VA-PFN	SM-GP	$0.615\pm 0.038$	$0.812\pm 0.026$	0.967	0.934	0.027
DVA-PFN	SM-GP	$0.852\pm 0.015$	$0.981\pm 0.006$	0.998	0.997	0.002
TabPFN 2.5	tabular	$0.765\pm 0.021$	$0.882\pm 0.010$	0.993	0.982	0.002
TabPFN 3.0	tabular	$0.788\pm 0.009$	$0.926\pm 0.006$	1.000	1.000	$<0.001$
TabICLv2	tabular	$0.938\pm 0.008$	$0.985\pm 0.002$	1.000	0.999	$<0.001$
TabSwift	tabular	$0.861\pm 0.005$	$0.956\pm 0.004$	1.000	0.999	$<0.001$

Table 1: Parameter-free alignment (Sec. 3.1 ) and linear probes for frequency

f

(single sinusoids) and mixing weight

a

(two-component signals Sec. 3.2 ) on

\mkern 1.5mu\overline{\mkern-1.5muH\mkern-1.5mu}\mkern 1.5mu

\max|\Delta|

is the largest

|R^{2}_{\text{MLP}}-R^{2}_{\text{linear}}|

over the two targets. Mean

\pm

std over 5 seeds is shown; SM-GP: spectral-mixture GP prior. All entries use 250 context points per signal.

Figure 1: Causal dose–response under targeted subspace patching across architectures. Each panel sweeps patched directions

k

from

1

d

and reports the causal effect (CE) for spectral (red), non-spectral (blue), or random (green, 20 trials) subspaces.

	Air-5	ETT-L	ETT-D	Energy	En-perm	Wine	Air-20	Gas	ISOLET
$D$	5	6	6	8	8	11	20	127	499
$k_{50}$	24	24	24	32	16	32	32	48	64
B 12 top	0.95	0.88	0.89	0.85	0.93	0.90	0.87	0.87	0.59
B 12 rand.	0.35	0.13	0.16	0.06	0.17	0.20	0.31	0.10	0.11
B 24 top	0.96	0.94	0.90	0.91	0.95	0.93	0.89	0.77	0.77
B 24 rand.	0.34	0.27	0.39	0.32	0.13	0.30	0.34	0.23	0.36

Table 2: Fraction of each pair’s full-replacement effect transferred by patching the top 64 variance directions (top) or 64 random directions (rand.) at block 12 or block 24.

k_{50}

is the smallest

k

that reaches half of this effect at block 12. Means over 24 draws, with standard error at most 0.05.

Figure 2: The proposed Filter Bank Decoder (steps described in the text). Multi-query attention pooling helps once spectral complexity exceeds a single component (Table 6 in Appendix). The PFN is never updated; decoder details are in Appendix D.2 .

True kernel	Decoder (Ours)	PFN (DVA-PFN)	Amor-struct.	DKL	RFF	Oracle GP
RBF	$1.1\times 10^{-3}$	$9.7\times 10^{-4}$	$4.6\times 10^{-4}$	$8.4\times 10^{-4}$	$1.0\times 10^{-3}$	$1.7\times 10^{-4}$
Periodic	$1.5\times 10^{-3}$	$1.4\times 10^{-3}$	$1.5\times 10^{-3}$	$6.6\times 10^{-3}$	$6.4\times 10^{-3}$	$8.1\times 10^{-4}$
$K_{\text{RP}}$	$2.7\times 10^{-3}$	$1.1\times 10^{-3}$	$1.6\times 10^{-3}$	$4.9\times 10^{-3}$	$4.6\times 10^{-3}$	$9.5\times 10^{-4}$
SM ( $Q=1$ )	$1.3\times 10^{-3}$	$4.6\times 10^{-4}$	$4.3\times 10^{-4}$	$6.4\times 10^{-4}$	$8.0\times 10^{-4}$	$2.9\times 10^{-4}$
SM ( $Q=4$ )	$1.5\times 10^{-3}$	$6.2\times 10^{-4}$	$4.4\times 10^{-4}$	$1.1\times 10^{-3}$	$1.4\times 10^{-3}$	$1.9\times 10^{-4}$
Avg. Time (s)	0.0036	0.0020	0.0288	0.9180	0.7580	3.9460

Table 3: GP regression MSE on five kernel families with 50 context points (

K_{\text{RP}}

denotes RBF

\times

Periodic). The Decoder (single-realization) and PFN columns use DVA-PFN. The decoder serves as a diagnostic extraction, not a general predictive replacement.

Appendix figures & tables40 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Manifold structures colored by frequency obtained via frequecy varying experiment described in Section 3 . Note the tight spectral clusters in DVA-PFN compared to the smoother manifolds in VA and TabPFN due to diffusion of information inside attention mechanism.

Target	$H$	$V$	$H+V$
Single Frequency	0.98	0.21	0.99
Dual Frequencies	0.96	0.00	0.96
Full Spectral $(f_{1},f_{2},a_{1},a_{2})$	0.50	0.00	0.50

Appendix

Table 4: Probing R 2 scores for spectral parameter extraction using linear probes on DVA-PFN.

H

consistently dominates

V

across all targets.

Figure 4: Frequency Probing: Linear and MLP Relational Scatter Plots across PFN architectures.

Figure 5: Weight Probing: Linear and MLP Relational Scatter Plots across PFN architectures.

Figure 6: Control experiment: probing on trained vs. randomly initialized DVA-PFN. Linear and MLP probes are trained to recover frequency (left) and mixing weight (right) from frozen

\mkern 1.5mu\overline{\mkern-1.5muH\mkern-1.5mu}\mkern 1.5mu

. On the trained network both probes succeed (

R^{2}\geq 0.99

). On the randomly initialized network, the MLP probe still recovers both targets—consistent with its capacity to learn the mapping itself [ Hewitt and Liang, 2019 ] —but the linear probe fails (

R^{2}=0.18

for frequency,

0.64

for weight), confirming that linear accessibility is a consequence of learned representational structure, not of trivial input geometry.

	1D sinusoids		5D RBF-GP
$k$	top- $k$	PC 0 removed	top- $k$	PC 0 removed
1 (PC 0 alone)	0.714	—	0.422	—
4	0.868	0.170	0.863	0.402
8	0.869	0.171	0.875	0.406
16	0.879	0.184	0.880	0.413
64	0.889	0.196	0.883	0.416

Appendix

Table 5: Probe

R^{2}

from the top-

k

principal components of

\overline{H}

on TabPFN, with and without PC 0 , at 120 context points per signal (Table 1 uses 250).

Figure 10: Probing performance on multi-component signals shows successful frequency extraction (

f_{1},f_{2}

) but poor relative amplitude prediction (

a_{1},a_{2}

) on TabPFN.

Task	Mean( $H$ )	Mean( $H\!+\!V$ )	Attn( $H$ )	Attn( $H\!+\!V$ )	$\Delta$
Easy (1 param)	0.999	1.000	1.000	1.000	$+0.000$
Medium (2 params)	0.982	0.983	0.991	0.991	$+0.008$
Hard (4 params)	0.560	0.564	0.610	0.602	$+0.038$
Very Hard (6 params)	0.333	0.342	0.404	0.399	$+0.057$

Appendix

Table 6: Pooling ablation on

\mkern 1.5mu\overline{\mkern-1.5muH\mkern-1.5mu}\mkern 1.5mu

(MLP probe,

R^{2}

) on DVA-PFN. Mean pooling degrades sharply as spectral complexity grows, while multi-query attention pooling recovers an increasing fraction of the lost signal.

V

-only probes yield

R^{2}\!=\!0

at every difficulty and are omitted.

\Delta

reports the gain of attention pooling over mean pooling on

H\!+\!V

	RBF		Matérn-3/2		Matérn-5/2
$D$	$\rho_{\Delta}$	$\|r\|_{\mathrm{PC0}}$	$\rho_{\Delta}$	$\|r\|_{\mathrm{PC0}}$	$\rho_{\Delta}$	$\|r\|_{\mathrm{PC0}}$
1	0.42	0.71	0.33	0.67	0.48	0.75
2	0.66	0.79	0.48	0.68	0.55	0.72
3	0.81	0.93	0.83	0.90	0.78	0.91
4	0.81	0.90	0.88	0.95	0.82	0.92
5	0.82	0.89	0.88	0.94	0.83	0.91

Appendix

Table 7: Parameter-free probing metrics for GP functions on frozen TabPFN across input dimensions

D

(500 functions of 500 points per cell,

\ell

log-uniform on

[0.05,10]

). At

D=5

the values agree with an earlier independent run (RBF 0.815 / 0.900; Matérn-3/2 0.861 / 0.928).

Figure 11: t-SNE projections of mean-pooled TabPFN embeddings $\mkern 1.5mu\overline{\mkern-1.5muH\mkern-1.5mu}\mkern 1.5mu$ for 500 functions drawn from 5D GPs with RBF (left) and Matérn-3/2 (right) kernels, colored by characteristic frequency $f_{\mathrm{char}}=1/(2\pi\ell)$ . Despite the absence of any sinusoidal structure in the generating process, embeddings organize smoothly by spectral scale, consistent with the quantitative metrics in Table 7 .

Figure 12: Layer-wise correlation

\rho_{\Delta}

between embedding distances and generating-frequency distances. DVA-PFN (left) peaks at L2 (

\rho_{\Delta}=0.95

) and TabPFN (right) plateaus near

0.93

by B12. Both architectures show a mild decline in the final layers, consistent with a shift toward posterior formatting.

Layer	H-patch CE	V-patch CE	K-patch CE	$p$ (H vs. K)
L1	$0.000\pm 0.000$	$0.999\pm 0.002$	$0.000\pm 0.000$	n.a.
L2	$0.999\pm 0.002$	$0.999\pm 0.002$	$0.000\pm 0.000$	$<10^{-133}$
L3–L6	$0.999\pm 0.002$	$0.999\pm 0.002$	$0.000\pm 0.000$	$<10^{-133}$

Appendix

Table 8: Layer-wise Causal Effect (CE) for activation patching. Mean

\pm

s.d. over

n=50

pairs;

p

-values from paired

t

-tests (H vs. K).

Figure 13: Representative RBF-GP signal pairs used for activation patching. Small $\ell$ (red) produces high-frequency wiggly functions; large $\ell$ (blue) produces smooth slowly-varying functions. The large visual separation ensures a strong baseline $\mathrm{MSE}(\hat{y}_{A},\hat{y}_{B})$ , making the causal effect measurement well-conditioned.

Patch type	L1	L2	L3	L4	L5	L6
H (causal)	0.00	1.00	1.00	1.00	1.00	1.00
V (+control)	1.00	1.00	1.00	1.00	1.00	1.00
K ( $-$ control)	0.00	0.00	0.00	0.00	0.00	0.00

Appendix

Table 9: Layer-wise CE for DVA-PFN on RBF-GP pairs (

\ell_{A}\!=\!0.05,\,\ell_{B}\!=\!2.0

). H-patch reaches CE

\approx

1 by layer 2; K-patch gives exactly 0.

$k$	% of dims	Spectral CE	Non-spectral CE	Random CE	Selectivity
1	0.78%	$0.026$	$0.000$	$0.007$	—
2	1.56%	$\mathbf{0.592}$	$0.008$	$0.018$	$\mathbf{33.9\times}$
4	3.13%	$\mathbf{0.676}$	$0.017$	$0.032$	$21.3\times$
8	6.25%	$0.676$	$0.023$	$0.066$	$10.3\times$
16	12.5%	$0.689$	$0.115$	$0.149$	$4.6\times$
32	25.0%	$\mathbf{0.977}$	$0.135$	$0.326$	$3.0\times$

Appendix

Table 10: Targeted subspace patching on DVA-PFN: causal effect as a function of subspace dimensionality

k

(out of

d=128

). Selectivity

=

Spectral CE / Random CE.

Parameter	Value
Input grid size $N$	200
Lengthscale range $[\ell_{\min},\ell_{\max}]$	$[0.05,\,2.0]$
Probe set size	300 signals (log-uniform $\ell$ )
Patch pairs	30 (min gap $\|\ell_{A}-\ell_{B}\|\geq 0.8$ )
Subspace dims $k$	$\{1,2,4,8,16,32,64,128\}$
DVA-PFN intervention layer	Cross-attention block 2 (of 6)

Appendix

Table 11: Hyperparameters for non-sinusoidal patching experiments.

Figure 14: Targeted subspace patching on RBF-GP functions. Causal Effect (CE) as a function of patching dimensionality

k

for DVA-PFN ( left ) and TabPFN ( right ). Red : spectral subspace (PCs most correlated with

\ell

). Blue : non-spectral subspace (bottom-

k

PCs). Grey triangles: random

k

-dimensional subspace. Dashed grey: uniform

k/d

baseline. Dashed green: full-

H

replacement ceiling. Shaded bands: 95% confidence intervals over 30 signal pairs. The spectral subspace dominates at every

k

for both architectures, demonstrating that causal structural information is compactly organised beyond the sinusoidal training distribution.

	3-block regressor			6-block regressor
$k$	top	random	bottom	top	random	bottom
1	$0.14\pm 0.11$	$0.11\pm 0.07$	$0.00\pm 0.00$	$0.04\pm 0.02$	$0.02\pm 0.02$	$0.00\pm 0.00$
4	$0.34\pm 0.23$	$0.01\pm 0.00$	$0.01\pm 0.01$	$0.14\pm 0.08$	$0.00\pm 0.00$	$0.01\pm 0.01$
16	$0.67\pm 0.22$	$0.10\pm 0.06$	$0.02\pm 0.02$	$0.68\pm 0.23$	$0.06\pm 0.04$	$0.01\pm 0.01$
64	$0.94\pm 0.02$	$0.64\pm 0.18$	$0.07\pm 0.03$	$0.88\pm 0.03$	$0.09\pm 0.02$	$0.05\pm 0.03$

Appendix

Table 12: Targeted subspace patching of tree-prior nanoTabPFN at the last block after 10,000 steps: CE for the top-

k

k

random and bottom-

k

frequency-ranked directions, mean

\pm

std over five trained models.

d=96

for the 3-block model and

192

for the 6-block model.

Pair	$D$	$k_{50}$	$k_{80}$	$k_{50}$	$k_{80}$	$\mathrm{CE}_{\mathrm{full}}$
		Block 12		Block 24
Airline/Milk, 5 lags	5	24	32	24	32	0.89
ETTh1/ETTh2, levels	6	24	64	24	48	0.74
ETTh1/ETTh2, differenced	6	24	48	32	48	0.75
Energy heating/cooling	8	32	64	32	48	0.58
Energy, permuted target	8	16	24	32	48	0.94

Appendix

Table 13: Smallest

k

(of 192) whose top-variance patch reaches 50% (

k_{50}

) or 80% (

k_{80}

) of the full-replacement effect, read from the sweep grid, and the full-replacement effect

\mathrm{CE}_{\mathrm{full}}

of each pair.

Pair	$D$	top-32	rand-32	bot-32	top-64	rand-64	bot-64
Airline/Milk, 5 lags	5	$0.864\pm 0.010$	0.192	0.003	$0.948\pm 0.010$	0.347	0.003
ETTh1/ETTh2, levels	6	$0.694\pm 0.019$	0.077	0.002	$0.875\pm 0.022$	0.128	0.005
ETTh1/ETTh2, diff.	6	$0.756\pm 0.025$	0.106	0.004	$0.886\pm 0.025$	0.159	0.005
Energy heating/cooling	8	$0.641\pm 0.043$	0.047	0.002	$0.849\pm 0.048$	0.057	0.007
Energy, permuted	8	$0.862\pm 0.006$	0.038	0.001	$0.928\pm 0.007$	0.171	0.003
Wine red/white	11	$0.693\pm 0.024$	0.080	0.002	$0.896\pm 0.019$	0.201	0.003

Appendix

Table 14: Block 12: fraction of the full-replacement effect for the top-, random- and bottom-

k

subspaces. Mean over 24 draws,

\pm

standard error for the top subspace.

Pair	$D$	top-32	rand-32	bot-32	top-64	rand-64	bot-64
Airline/Milk, 5 lags	5	$0.860\pm 0.007$	0.109	0.001	$0.959\pm 0.010$	0.344	0.003
ETTh1/ETTh2, levels	6	$0.715\pm 0.019$	0.148	0.001	$0.943\pm 0.023$	0.268	0.002
ETTh1/ETTh2, diff.	6	$0.614\pm 0.019$	0.208	0.002	$0.897\pm 0.024$	0.385	0.007
Energy heating/cooling	8	$0.631\pm 0.036$	0.174	0.001	$0.910\pm 0.052$	0.321	0.002
Energy, permuted	8	$0.738\pm 0.006$	0.029	0.000	$0.949\pm 0.007$	0.126	0.000
Wine red/white	11	$0.670\pm 0.019$	0.155	0.002	$0.925\pm 0.018$	0.296	0.001

Appendix

Table 15: Block 24 (final block), as in Table 14 .

Parameter	Multi-Realization	Single-Realization
$n_{\text{samples}}$	100,000	300,000
$n_{\text{points}}$	200	200
$n_{\text{bins}}$	50	50
$d_{\text{model}}$	128	128
$d_{\text{ff}}$	256	256
$n_{\text{queries}}$	4	4

Appendix

Table 16: Decoder training hyperparameters.

Variant	RBF	Periodic	LP	SM
Full decoder	6.4	3.0	1.3	0.9
Raw input	6.1	3.1	1.2	0.9
Random PFN	518.6	346.0	229.6	146.7
$H$ only	6.3	3.0	1.3	0.9
$V$ only	572.2	349.6	256.6	187.6

Appendix

Table 17: Decoder controls: Ker-MSE (

\times 10^{-3}

\downarrow

), mean over 5 seeds. LP: locally periodic.

Kernel	Decoder	DKL	RFF	Oracle GP
RBF	$-2.64\pm 0.58$	$-1.63\pm 0.19$	$-1.62\pm 0.17$	$-2.86\pm 0.45$
Periodic	$-1.66\pm 0.84$	$-1.45\pm 0.37$	$-0.79\pm 0.94$	$-1.94\pm 1.50$
$K_{\text{RP}}$	$-1.53\pm 0.96$	$-1.52\pm 0.21$	$-0.71\pm 0.96$	$-2.86\pm 0.34$
SM ( $Q=2$ )	$-2.67\pm 0.67$	$-1.64\pm 0.16$	$-1.65\pm 0.20$	$-3.08\pm 0.54$
SM ( $Q=4$ )	$-2.27\pm 1.38$	$-1.63\pm 0.06$	$-1.12\pm 1.11$	$-3.00\pm 0.46$

Appendix

Table 18: NLPD (

\downarrow

) on 150 held-out targets with 50 context points; mean

\pm

std over 100 realizations. The decoder uses DVA-PFN.

Figure 15: GP-MSE versus context length on (left) in-distribution sinusoids and (centre) out-of-distribution triangle waves with random frequencies

\sim\mathcal{U}(1.0,3.0)

. Decoded kernels match DKL/RFF on both, without per-task optimization. Computational time (right) is reported over 100 tasks.

Figure 16: Kernel reconstruction from single function observations . Global amplitude is ambiguous (Theorem 1 ); dominant periodic structure and lengthscales are recovered.

Figure 17: Single-realization decoder performance. (a) Qualitative GP regression from a decoded kernel with

N=20

context points. (b) GP-MSE versus context size

N_{\text{context}}

. The gap to the oracle remains roughly constant in

N

, consistent with the single-realization weight ambiguity (Theorem 1 ).

Kernel family	Method	Ker-MSE ( $\downarrow$ )	Ker-CKA ( $\uparrow$ )	GP-MSE ( $\downarrow$ )
RBF	Periodogram	0.0060	0.7316	0.0859
	Lomb–Scargle	0.0031	0.7177	0.0788
	PFN Decoder (Ours)	0.0062	0.3648	0.0011
Periodic	Periodogram	0.0028	0.8768	0.0964
	Lomb–Scargle	0.0023	0.7191	0.0468
	PFN Decoder (Ours)	0.0031	0.5791	0.0015

Appendix

Table 19: Spectral recovery benchmark (50 context, 150 target points). Classical methods overfit the context set (low Ker-MSE) but fail at GP-MSE on unseen targets. The PFN-based decoder uses pretraining-induced regularization to generalize. We report Mean MSE over 100 samples. Bold marks best per column within each kernel family.

Figure 18: GP predictive posterior on a complex spectral mixture from a singlerealization. Periodogram and Lomb–Scargle interpolate the 50 context points exactly but oscillate wildly between observations; the decoded kernel recovers the underlying structure and tracks the held-out target.

Kernel type	Decoded Kernel + GP	Oracle GP
SM ( $Q\!=\!1$ )	$0.001\pm 0.002$	$0.005\pm 0.014$
SM ( $Q\!=\!2$ )	$0.002\pm 0.009$	$0.004\pm 0.015$
SM ( $Q\!=\!4$ )	$0.006\pm 0.029$	$0.004\pm 0.028$

Appendix

Table 20: Bayesian optimization with the decoded kernel. Final simple regret (mean

\pm

std,

\downarrow

better, 100 functions each).

	PFN (per step)	Decoded Kernel + GP
Parameters	859,108	9
Memory	3.4 MB	72 B
Time / iteration (CPU)	24.8 ms	4.8 ms
GPU required	yes	no

Appendix

Table 21: Resource comparison for iterative inference on CPU.

Figure 19: Wasserstein distance between true and decoded spectral densities as a function of the number of independent realizations

M

, for three bandwidths

\sigma\in\{0.01,0.03,0.05\}

. The overall decrease is consistent with Theorem 2 .

Figure 20: Decoded vs. ground-truth spectral densities for 1–4 component mixtures.

Out-of-distribution
Kernel family	Oracle GP MSE	Decoder MSE
SM ( $Q\!=\!1$ )	$1.13\!\times\!10^{-4}$	$6.15\!\times\!10^{-4}$
SM ( $Q\!=\!2$ )	$1.45\!\times\!10^{-4}$	$1.14\!\times\!10^{-4}$
SM ( $Q\!=\!4$ )	$1.17\!\times\!10^{-4}$	$2.99\!\times\!10^{-4}$
RBF	$1.34\!\times\!10^{-4}$	$4.76\!\times\!10^{-3}$
Matérn-1/2	$9.50\!\times\!10^{-2}$	$1.78\!\times\!10^{-1}$

Appendix

Table 22: Decoder GP-MSE on a fixed support of 16 functions per kernel, evaluated on 20 unseen test functions from the same prior. RBF and Matérn families lie outside the sparse-spectral-mixture model class; the decoder shows graceful degradation rather than failure.

Kernel (additive)	Oracle GP MSE	Decoder MSE
RBF (5D)	$1.30\!\times\!10^{-4}$	$2.90\!\times\!10^{-4}$
SM ( $Q\!=\!1$ , 5D)	$1.10\!\times\!10^{-4}$	$3.00\!\times\!10^{-4}$
SM ( $Q\!=\!2$ , 5D)	$1.20\!\times\!10^{-4}$	$2.40\!\times\!10^{-4}$
SM ( $Q\!=\!4$ , 5D)	$1.30\!\times\!10^{-4}$	$2.90\!\times\!10^{-4}$
Periodic (5D)	$2.60\!\times\!10^{-4}$	$9.50\!\times\!10^{-3}$
RBF (10D)	$1.30\!\times\!10^{-4}$	$3.25\!\times\!10^{-4}$

Appendix

Table 23: Multi-realization decoder GP-MSE in 5D and 10D additive-kernel settings,

M=16

functions per task. Predictions remain within a small constant factor of the oracle except on Periodic, which is hardest under additive structure.

$d$	$L$	attn	params	MSE (mean)	MSE (std)	lat (ms)
16	2	MQA	8.6K	$1.09\!\times\!10^{-2}$	$1.58\!\times\!10^{-2}$	0.89
32	2	MQA	28.8K	$6.26\!\times\!10^{-4}$	$2.17\!\times\!10^{-4}$	0.80
48	2	MQA	61.6K	$9.57\!\times\!10^{-5}$	$6.54\!\times\!10^{-5}$	0.82
64	2	MQA	104.6K	$7.74\!\times\!10^{-5}$	$3.83\!\times\!10^{-5}$	1.15
96	2	MQA	229.1K	$5.65\!\times\!10^{-5}$	$2.47\!\times\!10^{-5}$	1.06
128	2	MQA	401.6K	$6.71\!\times\!10^{-5}$	$1.37\!\times\!10^{-5}$	1.20

Appendix

Table 24: Full architecture grid (Appendix E ). Six widths

\times

three depths

\times

two attention variants, three seeds each. Mean and standard deviation of test MSE across seeds; latency reported in ms per batch of 32 on a single GPU.

Figure 21: Architecture ablation supporting the mechanistic predictions of Secs. 3.2 – 4 . Test MSE on the held-out 2000-function set across a

6\!\times\!3\!\times\!2

grid of widths

d

, depths

L

, and attention variants (3 seeds each). (a) MSE vs.

d_{\mathrm{model}}

, with shaded

\pm 1

std bands across seeds. Width gains plateau by

d\!\approx\!48

–

64

, and the spread between depths shrinks as

d

shrinks — consistent with the

L_{1}\!\to\!L_{2}

emergence of spectral coding (Sec. 4 ): once the first cross-attention step has constructed the spectrally-organized latent, additional layers contribute mainly posterior formatting, with diminishing returns. (b) Parameter–MSE Pareto frontier (dashed). The annotated configuration (

d\!=\!64

L\!=\!2

, MQA) achieves test MSE

7.7\!\times\!10^{-5}

\sim\!105

K parameters and

1.15

ms/batch — within

2.5\times

of the largest configuration (

d\!=\!128

L\!=\!6

, Standard;

1.25

M parameters,

3.66

ms) at

12\times

fewer parameters and

3.2\times

lower latency. 12 of 15 Pareto-optimal points use MQA. Full numerical table in Table 24 .

Explore similar work

Jun 9, 2026cs.LG

CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

Prior-fitted networks (PFNs) are a promising class of tabular foundation models that perform in-context learning, whereby the entire labelled training set is supplied as context, and predictions for test queries are produced in a single forward pass. However, the quadratically scaling self-attention mechanism in many PFN architectures makes inference prohibitive for very large training datasets. We propose CRUMB (Clustered Retrieval Using Minimised-MMD Batching), a three-stage inference wrapper that (i) clusters the test queries, (ii) selects a small, distributionally matched training subset for each cluster by greedily minimising the maximum mean discrepancy (MMD), and (iii) runs exact PFN inference on each reduced-context batch. CRUMB is architecture-agnostic and requires no retraining. On the 51-dataset TabArena benchmark, evaluated across three PFN architectures (TabPFNv2, TabICLv1, TabICLv2), we show that CRUMB outperforms similar state-of-the-art context selection strategies. We also show that CRUMB is resilient to covariate drift, as the MMD-minimisation step naturally helps align the training context distribution to match the current test batch distributions.

Jamie Heredge, Mattia J. Villani, Pranav Deshpande +2

Global Technology Applied Research, JPMorganChase, New York, NY 10001, USA

In-Context Learning Prior-Data Fitted Networks Tabular ICL

Ask PDF

May 7, 2026stat.ML

Decoupled PFNs: Identifiable Epistemic-Aleatoric Decomposition via Structured Synthetic Priors

Prior-Fitted Networks (PFNs) amortize Bayesian prediction by meta-learning over a synthetic task prior, but their standard output is a posterior predictive distribution over noisy observations. For sequential decision-making, such as active learning and Bayesian optimization, acquisition should prioritize epistemic uncertainty about the latent signal rather than irreducible aleatoric observation noise. We show that this epistemic--aleatoric split is not identifiable in general from the posterior predictive distribution alone, even when that distribution is known exactly. We then exploit a distinctive advantage of PFNs: because the synthetic data-generating process is under our control, each task can contain an explicit latent signal and noise function, and the generator can provide query-level labels for both the noiseless target and the observation-noise variance. We use these labels to train a decoupled PFN with separate latent-signal and aleatoric heads. The observation-level predictive is induced by convolving the latent signal distribution with the learned noise model. Empirically, epistemic-only acquisition mitigates the failure mode of total-variance exploration in noisy and heteroscedastic settings. In matched comparisons, decoupled models usually improve over tuned observation-level baselines, with the clearest gains in HPO; in broader sweeps, a decoupled model obtains the best average rank in both HPO and synthetic BO.

Richard Bergna, Stefan Depeweg, José Miguel Hernández-Lobato

University of Cambridge · 2Siemens AG

Uncertainty Quantification Bayesian Optimization Active Learning Prior-Data Fitted Networks

Ask PDF

May 11, 2026stat.ML

PFN-TS: Thompson Sampling for Contextual Bandits via Prior-Data Fitted Networks

Thompson sampling is a widely used strategy for contextual bandits: at each round, it samples a reward function from a Bayesian posterior and acts greedily under that sample. Prior-data fitted networks (PFNs), such as TabPFN v2+ and TabICL v2, are attractive candidates for this purpose because they approximate Bayesian posterior predictive distributions in a single forward pass. However, PFNs predict noisy future rewards, while Thompson sampling requires uncertainty over the latent mean reward function. We propose PFN-TS, a Thompson sampling algorithm that converts PFN posterior predictives into mean-reward samples using a subsampled predictive central limit theorem. The method estimates posterior variance from a geometric grid of

O(\log n)

dataset prefixes rather than the full

O(n)

predictive sequence used in previous predictive-sequence approaches, and reuses TabICL's cached representations across rounds. We prove consistency of the subsampled variance estimator and give a Bayesian regret bound that decomposes PFN-TS regret into exact posterior-sampling regret under the PFN prior plus approximation terms. Empirically, PFN-TS achieves the best average rank across nonlinear synthetic and OpenML classification-to-bandit benchmarks, remains competitive on linear and BART-generated rewards, and attains the highest estimated policy value in an offline mobile-health evaluation. Code is available at https://anonymous.4open.science/r/PFN_TS-36ED/.

Yan Shuo Tan, Kenyon Ng, Ruizhe Deng +3

National University of Singapore · Monash University · Renmin University of China +1

Thompson Sampling Prior-Data Fitted Networks Contextual Bandits

Ask PDF