Random Feature Gaussian Process Attention: Linear-Time Probabilistic Attention with Calibrated Uncertainty

Amir Mohammad Mahfoozi; Zi Yang; Ying Li; Michael Minyi Zhang

cs.LGOct 6, 2026

Random Feature Gaussian Process Attention: Linear-Time Probabilistic Attention with Calibrated Uncertainty

Authors: Amir Mohammad Mahfoozi, Zi Yang, Ying Li, Michael Minyi Zhang

Organizations: Department of Computer Engineering, Sharif University of Technology · School of Computing and Data Science, The University of Hong Kong

Abstract

Transformers provide a state-of-the-art modeling framework, yet poor calibration limits their reliability in safety-critical applications. A promising direction addresses this issue by interpreting attention as a Gaussian process (GP) posterior, which enables principled uncertainty calibration but incurs cubic complexity in sequence length due to the inversion of the kernel; although decoupled GP variants reduced the cost to quadratic, the computation remains prohibitive in practice. In this paper, we propose the plug-and-play random Fourier feature Gaussian process attention (RFF-GPA) module, which represents the attention as a GP with a stationary kernel approximated by random Fourier features. This low-rank approximation results in linear-time complexity for approximating the posterior mean and variance, making it far more scalable compared to previous work. Empirical results on multiple real-world datasets show that our attention module improves calibration while maintaining predictive accuracy, and simultaneously reduces computational complexity to linear in the sequence length.

Figures & tables

Figure 1: Overview of RFF-GPA. (A) Standard attention produces deterministic token representations without explicit uncertainty. (B,D) A GP formulation provides posterior mean and variance, but exact GP attention requires an expensive token-space inverse. (C) Random Fourier features provide a compact approximation of the stationary kernel. (E) RFF-GPA performs GP inference in the resulting feature space, replacing the $L\times L$ inverse with an $M\times M$ inverse while retaining the posterior mean and token-wise predictive variance.

METHOD	IMAGE			TEXT
METHOD	F-MNIST	CIFAR-10	SVHN	20NG	Hyperp.	SST-2
METRIC	ACC $\uparrow$
MLE	0.889 $\pm$ 0.003	0.715 $\pm$ 0.008	0.909 $\pm$ 0.007	0.654 $\pm$ 0.015	0.744 $\pm$ 0.015	0.775 $\pm$ 0.009
MLE+Temp	0.889 $\pm$ 0.003	0.715 $\pm$ 0.008	0.909 $\pm$ 0.007	0.654 $\pm$ 0.015	0.744 $\pm$ 0.015	0.775 $\pm$ 0.009
MCD	0.885 $\pm$ 0.001	0.713 $\pm$ 0.002	0.902 $\pm$ 0.001	0.682 $\pm$ 0.004	0.785 $\pm$ 0.010	0.781 $\pm$ 0.006
SNGP	0.903 $\pm$ 0.002	0.727 $\pm$ 0.007	0.913 $\pm$ 0.004	0.680 $\pm$ 0.009	0.774 $\pm$ 0.015	0.789 $\pm$ 0.004

Table 1: In-distribution accuracy and calibration on three image and three text classification benchmarks. Each column reports one dataset; row blocks report ACC

\uparrow

, NLL

\downarrow

, and ECE

\downarrow

. Mean and standard deviation are computed over five runs. The best result per column is bolded and the second-best is underlined .

Panel of Figure 2 — Figure 2: Training Time (Left): Training time per step as a function of sequence length $L$ across all methods (log–log axes; shaded 95% CI). Throughput (Right): Throughput (tokens/sec) as a function of $L$ . Curves for the inducing-point baselines (SGPA, SCGP) are continued past their largest feasible length along their fitted scaling trend.

METHOD	$s{=}0$ (clean)	$s{=}1$	$s{=}2$	$s{=}3$	$s{=}4$	$s{=}5$
METRIC	ACC $\uparrow$
MLE	0.715 $\pm$ 0.008	0.681 $\pm$ 0.010	0.637 $\pm$ 0.009	0.605 $\pm$ 0.009	0.564 $\pm$ 0.010	0.503 $\pm$ 0.008
SNGP	0.727 $\pm$ 0.007	0.693 $\pm$ 0.006	0.647 $\pm$ 0.001	0.613 $\pm$ 0.001	0.572 $\pm$ 0.002	0.508 $\pm$ 0.003
SGPA	0.665 $\pm$ 0.023	0.613 $\pm$ 0.023	0.556 $\pm$ 0.024	0.530 $\pm$ 0.018	0.495 $\pm$ 0.014	0.441 $\pm$ 0.012
RFF-CGP (ours)	0.737 $\pm$ 0.003	0.664 $\pm$ 0.003	0.616 $\pm$ 0.004	0.588 $\pm$ 0.004	0.548 $\pm$ 0.001	0.490 $\pm$ 0.002
RFF-GPA (ours)	0.759 $\pm$ 0.001	0.702 $\pm$ 0.001	0.653 $\pm$ 0.002	0.620 $\pm$ 0.002	0.576 $\pm$ 0.003	0.514 $\pm$ 0.002

Table 2: Out-of-distribution robustness on CIFAR-10-C. Each metric is averaged over the 19 corruption types at severity

s\!\in\!\{1,\dots,5\}

;

s\!=\!0

is the clean CIFAR-10 test set. Row blocks report ACC

\uparrow

, NLL

\downarrow

, and ECE

\downarrow

. Values are mean

\pm

95% CI over five runs; all models are evaluated without re-training. The best result per column is bolded and the second-best is underlined .

Panel of Figure 3 — Figure 3: Localized token-wise posterior uncertainty. Object perturbation (top) and matched background perturbation (bottom). Red grids mark the $12$ perturbed tokens; clean and corrupted uncertainty maps use a shared scale.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Name	Description	Size	Citation/source
Fashion-MNIST	Grayscale $28{\times}28$ clothing-image classification with $10$ classes; patch size $4{\times}4$ gives sequence length $49$ .	$60$ k train / $10$ k test	Xiao, Rasul, and Vollgraf (2017)
CIFAR-10	Color $32{\times}32$ natural-image classification with $10$ classes; patch size $2{\times}2$ gives sequence length $256$ .	$50$ k train / $10$ k test	Krizhevsky and Hinton (2009)
SVHN	Real-world $32{\times}32$ street-view house-number digit classification with $10$ classes; sequence length $256$ .	$\sim$ 73k train / $26$ k test	Netzer et al. (2011)
CIFAR-10-C	Distribution-shift benchmark built from CIFAR-10 with common image corruptions at multiple severities.	$19$ corruptions $\times$ $5$ severities $\times$ $10$ k images	Hendrycks and Dietterich (2019)
20 Newsgroups	Topic classification over newsgroup posts with $20$ classes; text is truncated/padded to $512$ tokens.	$\sim$ 18k documents	Lang (1995)
Hyperpartisan	Binary news-bias classification using the by-article split; text is truncated/padded to $512$ tokens.	Small by-article split	Kiesel et al. (2019)

Appendix

Table 3: Datasets used in the in-distribution and out-of-distribution experiments. Images are patchified inside the Vision-Transformer backbone and standardized per channel; text uses a shared WordPiece tokenizer.

Explore similar work

Projection-Free Transformers via Gaussian Kernel Attention
May 4, 2026Debarshi Kundu, Archisman Ghosh, Swaroop Ghosh +1Transformer Attention Transformer Architectures
The Pennsylvania State University, University Park, USA
Robust Filter Attention: Self-Attention as Precision-Weighted State Estimation
Sep 4, 2025Peter RacioppoLinear Attention State Estimation
Independent Researcher, Los Angeles, CA, USA.
Exact Linear Attention
May 13, 2026Weinuo OuKimi Delta Attention Transformer Attention
Wuyi University

May 4, 2026, cs.LGJ/K move · Enter open · S save

Projection-Free Transformers via Gaussian Kernel Attention

Debarshi Kundu, Archisman Ghosh, Swaroop Ghosh, Vasant Honavar

The Pennsylvania State University, University Park, USA

Transformer Attention Transformer Architectures Kernel Method Autoregressive Language Models

Open paper AskPDF arXiv

Self-attention in Transformers is typically implemented as

\mathrm{softmax}(QK^\top/\sqrt{d})V

, where

Q=XW_Q

K=XW_K

, and

V=XW_V

are learned linear projections of the input

X

. We ask whether these learned projections are necessary, or whether they can be replaced by a simpler similarity-based diffusion operator. We introduce \textbf{Gaussian Kernel Attention} (GKA), a drop-in replacement for dot-product attention that computes token affinities directly using a Gaussian radial basis function (RBF) kernel applied to per-head token features. Each head learns only a bandwidth parameter

σ_h

, while a single output projection

W_O

preserves compatibility with the standard Transformer interface. GKA can be interpreted as normalized kernel regression over tokens, linking modern Transformer architectures to classical non-local filtering and kernel smoothing methods. We evaluate GKA in both vision and language modeling settings. For autoregressive language modeling within the \texttt{nanochat} framework, we implement causal masking and sliding-window constraints by masking and renormalizing the Gaussian kernel. At depth 20, a GKA model with

0.42\times

the parameters and

0.49\times

the total training FLOPs of a standard attention baseline trains stably, exhibits a near-zero train-validation gap, and demonstrates competitive behavior on standard benchmarks, albeit with higher bits-per-byte (BPB) at this compute scale. Overall, GKA provides a minimal, interpretable attention mechanism with an explicit locality scale, offering a dimension in the accuracy-efficiency trade-off for Transformer design.

Abstract

Explore similar work

Projection-Free Transformers via Gaussian Kernel Attention

Robust Filter Attention: Self-Attention as Precision-Weighted State Estimation

Exact Linear Attention