Decoupled-Value Attention for Prior-Data Fitted Networks: GP-Inference for Physical Equations

ICLR 2026 Conference Submission15541 Authors

Keywords: Gaussian Process, Meta-Learning, Prior-data Fitted Networks, Learning of Physics

TL;DR: Decoupled-Value Attention (DVA) separates input similarity from label propagation, mirroring Gaussian process updates and enabling scalable, kernel-free PFNs. This achieves architecture-agnostic and scalable PFNs.

Abstract: Prior-data fitted networks (PFNs) are a promising alternative to time-consuming Gaussian process (GP) inference for creating fast surrogates of physical systems. PFN reduces the computational burden of GP-training by replacing Bayesian inference in GP with a single forward pass of a learned prediction model. However, with standard Transformer attention, PFNs show limited effectiveness on high-dimensional regression tasks. We introduce Decoupled-Value Attention (DVA)-- motivated by the GP property that the function space is fully characterized by the kernel over inputs and the predictive mean is a weighted sum of training targets. DVA computes similarities from inputs only and propagates labels solely through values. Thus, the proposed DVA mirrors the GP update while remaining kernel-free. We demonstrate that the crucial factor for scaling PFNs is the attention rule rather than the architecture itself. Specifically, our results demonstrate that (a) localized attention consistently reduces out-of-sample validation loss in PFNs across different dimensional settings, with validation loss reduced by more than 50\% in five- and ten-dimensional cases, and (b) the role of attention is more decisive than the choice of backbone architecture, showing that CNN-based PFNs can perform at par with their Transformer-based counterparts. The proposed PFNs provide 64-dimensional power flow equation approximations with a mean absolute error of the order of $10^{-3}$, while being over $80\times$ faster than exact GP inference.

Primary Area: transfer learning, meta learning, and lifelong learning

Submission Number: 15541