Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Amortized Spectral Kernel Discovery via Prior-Data Fitted Network

Kaustubh Sharma, Srijan Tiwari, Ojasva Nema, Parikshit Pareek

TL;DR

This work addresses the opacity of Prior-Data Fitted Networks (PFNs) by deriving an interpretability-driven pipeline that decodes explicit spectral densities $S(\omega)$ and stationary kernels $k(\tau)$ from PFN latents using Bochner's theorem. Through mechanistic analysis of Decoupled-Value Attention (DVA), the authors show that the hidden state $H$ encodes spectral structure, enabling two decoders: a Multi-Realization Decoder for full spectral mixtures and a Single-Realization Decoder for dominant components, with an analytical scaling step to match observed energy. They establish identifiability limits: spectral weights are not identifiable from a single realization but become identifiable with multiple realizations, and they provide consistency guarantees and computational savings by amortizing spectral inference. Empirically, decoded spectral kernels achieve GP regression accuracy comparable to iterative baselines while delivering orders-of-magnitude speedups, demonstrating practical value for surrogate-based optimization and scientific modeling. The framework emphasizes interpretability and reuse of PFN priors, offering a path toward transparent, zero-shot kernel construction from amortized inference.

Abstract

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.

Amortized Spectral Kernel Discovery via Prior-Data Fitted Network

TL;DR

This work addresses the opacity of Prior-Data Fitted Networks (PFNs) by deriving an interpretability-driven pipeline that decodes explicit spectral densities and stationary kernels from PFN latents using Bochner's theorem. Through mechanistic analysis of Decoupled-Value Attention (DVA), the authors show that the hidden state encodes spectral structure, enabling two decoders: a Multi-Realization Decoder for full spectral mixtures and a Single-Realization Decoder for dominant components, with an analytical scaling step to match observed energy. They establish identifiability limits: spectral weights are not identifiable from a single realization but become identifiable with multiple realizations, and they provide consistency guarantees and computational savings by amortizing spectral inference. Empirically, decoded spectral kernels achieve GP regression accuracy comparable to iterative baselines while delivering orders-of-magnitude speedups, demonstrating practical value for surrogate-based optimization and scientific modeling. The framework emphasizes interpretability and reuse of PFN priors, offering a path toward transparent, zero-shot kernel construction from amortized inference.

Abstract

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.
Paper Structure (47 sections, 3 theorems, 39 equations, 10 figures, 6 tables)

This paper contains 47 sections, 3 theorems, 39 equations, 10 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Positioning and Contributions:
  3. Background and Related Works
  4. Prior-Data Fitted Networks with DVA
  5. Kernel and Prior Learning
  6. Mechanistic Analysis of PFNs
  7. Statistical Identifiability of Spectral Density
  8. The Single-Realization Limit
  9. The Multi-Realization Guarantee
  10. Proposed Filter Bank Decoder
  11. Experimental Results and Discussions
  12. Multi-Realization Setting
  13. Single-Realization Setting
  14. Scaling Multi-Realization Decoder
  15. Conclusions
  16. ...and 32 more sections

Key Result

Theorem 4.1

Let $f : \mathbb{R} \to \mathbb{R}$ be a zero-mean stationary Gaussian process $f \sim \mathcal{GP}(0, k),$ where $k(\tau)$ is a continuous stationary kernel with spectral density constructed using $w_q > 0$ as $S(\omega) = \sum_{q=1}^Q w_q \, \mathcal{N}(\omega \mid \mu_q, \sigma_q^2)$. If $\{f(x_i

Figures (10)

  • Figure 1: t-SNE visualization of PFN embeddings. H (left) forms a smooth manifold ordered by frequency, while V (right) shows weaker spectral structure.
  • Figure 2: Architectural overview of the proposed Filter Bank Decoder: The framework operates on either single or multiple function realizations processed by a frozen PFN. Latent representations ($H, V$) are aggregated via Multi-Query Attention (MQA) pooling. A multi-head decoder identifies active spectral bins and regresses precise filter parameters. Finally, the explicit stationary kernel is reconstructed via Bochner’s theorem, utilizing the analytical scaling factor $\alpha$ to ensure energy consistency with the observations.
  • Figure 3: Wasserstein distance between true and predicted spectral densities as a function of the number of GP samples. Similar trends across different bandwidths $\sigma$ reflect the trade-off between peak sharpness and frequency localization under the Wasserstein metric.
  • Figure 4: Spectral density recovery. The horizontal axis represents frequency $\omega$ (Hz) and the vertical axis represents the spectral density magnitude $S(\omega)$. Blue solid curves denote the ground truth density, while red dashed curves indicate the explicit density recovered by Multi realization decoder. The model accurately identifies peak locations and bandwidths across varying mixture complexities (1--4 components).
  • Figure 5: Comparison of MSE (Log scale) on in-distribution (left) and out-of-distribution (right) tasks. In both the cases our framework gives comparable MSE to DKL and RFF baselines. As DKL and RFF require iterative optimization (approximately 1.44--2.17 seconds per task), our decoder produces a kernel in under $10^{-3}$ seconds via a single forward pass leading to $\approx1000\times$ reduction in inference time. Out-of-Distribution data is non-Gaussian triangle waves with random frequencies $\sim U(1.0, 3.0)$.
  • ...and 5 more figures

Theorems & Definitions (6)

  • Theorem 4.1: Single-function non-identifiability of spectral weights
  • Proposition 4.2: Unbiased estimation of the kernel scale
  • Theorem 4.3: Identifiability of spectral weights from multiple realizations
  • proof
  • proof
  • proof