Amortized Spectral Kernel Discovery via Prior-Data Fitted Network

Prior-Data Fitted Networks (PFNs) enable efficient amortized inference but lack transparent access to their learned priors and kernels. This opacity hinders their use in downstream tasks, such as surrogate-based optimization, that require explicit covariance models. We introduce an interpretability-driven framework for amortized spectral discovery from pre-trained PFNs with decoupled attention. We perform a mechanistic analysis on a trained PFN that identifies attention latent output as the key intermediary, linking observed function data to spectral structure. Building on this insight, we propose decoder architectures that map PFN latents to explicit spectral density estimates and corresponding stationary kernels via Bochner’s theorem. We study this pipeline in both single-realization and multi-realization regimes, contextualizing theoretical limits on spectral identifiability and proving consistency when multiple function samples are available. Empirically, the proposed decoders recover complex multi-peak spectral mixtures and produce explicit kernels that support Gaussian process regression with accuracy comparable to PFNs and optimization-based baselines, while requiring only a single forward pass. This yields orders-of-magnitude reductions in inference time compared to optimization-based baselines.

💡 Research Summary

The paper tackles the opacity problem of Prior‑Data Fitted Networks (PFNs), which provide fast amortized Bayesian inference but hide the underlying prior’s covariance structure. By focusing on PFNs that employ Decoupled‑Value Attention (DVA), the authors reveal that the attention output latent H encodes the spectral density of the data‑generating Gaussian Process. Through extensive mechanistic experiments—including t‑SNE visualizations, linear probing, and ablation studies—they demonstrate that H correlates strongly with underlying frequencies, while the value embeddings V carry only amplitude information.

Building on this insight, the authors design a decoder that maps the frozen latent H (or multiple pooled versions of H) to an explicit, non‑negative spectral density estimate Ŝ(ω) on a predefined frequency grid. Bochner’s theorem is then applied to transform Ŝ(ω) into a stationary kernel k̂(τ) = ∫Ŝ(ω) e^{iωτ} dω, which can be plugged directly into a Gaussian Process regression pipeline. The decoder operates in a zero‑shot fashion: PFN parameters remain fixed, and a single forward pass yields both the spectral estimate and the kernel.

The paper also provides a theoretical analysis of spectral identifiability. With a single function realization, only the locations of spectral peaks are statistically identifiable; the overall scale (weights) remains non‑identifiable. When multiple independent realizations from the same prior are available, the authors prove that the full spectral density—including weights—can be consistently recovered using an unbiased estimator based on the empirical second moment across realizations.

Empirically, the method is evaluated on synthetic datasets featuring single‑ and multi‑peak spectral mixtures, as well as on real‑world time‑series from physical simulations and sensor measurements. The decoded kernels achieve regression RMSE and log‑likelihood comparable to Deep Kernel Learning, Random Fourier Features, and iterative marginal‑likelihood optimization baselines. Crucially, inference time is reduced by one to two orders of magnitude because no test‑time optimization is required. Multi‑Query Attention Pooling further improves performance on complex spectra by providing diverse summaries of H, mitigating the loss of amplitude information that occurs with simple mean pooling.

Overall, the work reframes PFNs from black‑box predictors to “zero‑shot spectral inference engines.” It shows that PFNs already learn rich kernel information, and that this information can be extracted efficiently for downstream tasks that require explicit, interpretable covariance models, such as surrogate‑based optimization or scientific discovery. Limitations include the focus on stationary kernels and the fact that the recovered kernel approximates the PFN’s internal representation rather than the true data‑generating kernel. Future work may extend the approach to non‑stationary settings and explore tighter guarantees on kernel fidelity.