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

arXiv:2610.08578 · cs.LG · Submitted 2026-10-06 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Random Feature Gaussian Process Attention".

Jane: The gist The proposed Random Feature Gaussian Process Attention (RFF-GPA) module introduces an attention mechanism that models attention as a Gaussian process posterior approximated by random Fourier features,

Tom: First, who's behind it and why it matters.

Paper summary: Jane: So, wrapping up this discussion on "Random Feature Gaussian Process Attention: Linear-Time Probabilistic Attention with Calibrated Uncertainty", what is the actual practical meaning of all this RFF-GPA stuff for someone who isn't deep in the math?

Tom: Basically, they took a complex attention mechanism and gave it a way to report how confident it is about its predictions by treating that attention as a Gaussian process >

Lu: They solved the problem where standard methods were too slow because of the cubic complexity, and they did it by approximating the kernel with random Fourier features so they only need linear time for approximation >

Jane: So, instead of getting this expensive cubic time for a long sequence, you get linear time complexity while also getting a calibrated uncertainty estimate that works on real data >

Tom: That's the core contribution. It allows us to have a fully Bayesian model where we can actually handle long sequences efficiently and understand the uncertainty behind those predictions >

Meng: For deployment, it means we can build Transformers that are more reliable in safety-critical settings because they give us actual measures of how much to trust a prediction when things get messy >

Lu: The authors show this works well on image and text data sets, demonstrating that the RFF-GPA module is effective across different types of data >

Jane: It moves us toward a system where we can have attention that isn't just deterministic; it carries its own measure of uncertainty that actually changes when the input changes >

Tom: We've seen how this approach scales linearly, and it provides calibrated uncertainty metrics that hold up under distribution shifts, which is exactly what we need for robust AI systems >

Conclusion: Tom: So, we're at the end of this paper on Random Feature Gaussian Process Attention. The authors call it RFF-GPA—Random Fourier Feature Gaussian Process Attention—and they're saying that by using random Fourier features to approximate the kernel, they can get a fully Bayesian model for attention that scales linearly with sequence length.

Jane: So, if I’m hearing you right, it means instead of having this super slow cubic time complexity when the sequence gets long, we can actually run it in linear time when we want to predict stuff? That's the big deal for making these models practical.

Lu: Exactly. The core idea is that they take a complex attention output and interpret it as a Gaussian process posterior, but they use this low-rank approximation with random Fourier features so they don't have to do the heavy lifting of inverting that huge matrix every single time.

Meng: I'm thinking about deployment here. If the complexity drops from cubic to linear, that means we can finally train these on much longer sequences without needing a super powerful cluster just for the attention part. What about calibration? Does it actually give us better uncertainty measures than what we have now?

Tom: It does. The paper shows that this module improves calibration while keeping the accuracy pretty high on real-world image and text datasets, like Fashion-MNIST and CIFAR-ten. We're looking at predictive mean and variance estimates, which means we get a token-wise measure of uncertainty for every single word or pixel we process.

Jane: That uncertainty part is what really matters for safety. It tells us not just the answer, but how much to trust that answer when the input data is slightly off or shifted in some way. This makes it much more reliable than just getting a single prediction.

Lu: And they show it holds up under distribution shift too, which is crucial. The authors are suggesting this approach is a solid way to move towards more robust AI because it's not just about the prediction itself, but the confidence alongside it.

Meng: So, for someone actually shipping models right now, if we can get linear scaling and better uncertainty without needing massive hardware upgrades or spending forever on hyperparameter tuning for every single sequence length, that’s a win.

Tom: Right. The authors conclude that this RFF-GPA module is a way to build scalable Transformers using this GP framework with random Fourier features as the kernel approximation. It’s moving Bayesian attention out of the purely theoretical realm and into something we can actually use on larger data.

Jane: It feels like they've taken a really heavy math idea, Gaussian processes, and made it practical enough for real-world AI models to actually use in a way that makes sense for uncertainty.

Tom: That’s the gist of it. Next up, we’re going to look at how this works on correlated attention—that extension they introduced called RFFCGP—and see if that adds anything new to the picture.

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

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

cs.LG

Submitted: 2026-10-06

Updated: 2026-10-06

Comments: 14 pages, 3 figures, 3 tables

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 83/100

The gist: The gist The proposed Random Feature Gaussian Process Attention (RFF-GPA) module introduces an attention mechanism that models attention as a Gaussian process posterior approximated by random Fourier

Key concepts

Gaussian Process (GP) Posterior: A GP models the uncertainty around a function by treating it as a probability distribution rather than a single point estimate. In this context, the attention mechanism's output is treated as the mean of this posterior distribution, allowing the model to quantify its own confidence in its predictions.
Random Fourier Features (RFF): RFF is a low-rank approximation technique used to efficiently represent complex kernels. Instead of using a full, expensive kernel matrix, RFF approximates the stationary attention kernel with a set of random features. This allows the model to capture the essential structure of the attention relationship without needing massive computations.
Linear-Time Complexity: This refers to an algorithm whose runtime grows linearly with respect to the sequence length (L) rather than quadratically or cubically. By using RFF and matrix identities, this method can compute the necessary posterior mean and variance estimates in time proportional only to L, making it highly scalable for long sequences.
Calibration: Calibration is the process of ensuring that a model's predicted uncertainty accurately reflects its actual error rate. Poor calibration means the model might be overconfident or underconfident; this method improves calibration by providing token-wise variance estimates derived from the GP posterior, ensuring reliable uncertainty reporting.

Summary

The gist The proposed Random Feature Gaussian Process Attention (RFF-GPA) module introduces an attention mechanism that models attention as a Gaussian process posterior approximated by random Fourier features, achieving linear-time complexity for approximating the posterior mean and variance while improving calibration in Transformers.

Background on the Problem

Transformers are state-of-the-art modeling frameworks, yet poor calibration limits their reliability in safety-critical applications. Recent work has bridged the standard attention output to the mean of a Gaussian process (GP) posterior. Existing approaches remain computationally expensive, typically requiring inversion of an L×Ldimensional kernel matrix, resulting in cubicO(L 3)time and quadratic memory complexity. While decoupled GP variants reduced the cost to quadratic, the computation remains prohibitive in practice.

Proposed Solution: RFF-GPA

The paper proposes the plug-and-play random Fourier feature Gaussian process attention (RFF-GPA) module, which represents 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. The main contributions include introducing RFF-GPA, which approximates the attention kernel with random Fourier features (RFF) and performs GP posterior inference in the resulting feature space using the Woodbury matrix identity to avoid direct inversion of the full L × L kernel matrix. This avoids direct inversion of the full L × L kernel matrix and computes predictive means and covariance estimates in linear time with respect to the sequence length.

Methodology Details

The RFF-GPA module interprets keys as “training inputs” and values as (noisy) observations, where the d-th column of the kernel attention output is precisely the GP predictive posterior mean. The overall output of the kernel attention module Y can thus be interpreted as the predictive mean of D independent GPs, with the corresponding predictive covariance capturing associated uncertainty. By exploiting the low-rank structure of ΦK, a Woodbury identity is applied to obtain an M × M inverse, thereby reducing the overall cost to O(LM2), i.e., linear in L. The predictive uncertainty is reported as diag(Σ), which represents the token-wise variance.

Correlated Attention Extension

The model is further extended to a correlated random-feature GP attention variant, RFFCGP, which models query and key representations through a shared latent GP structure. This variant removes the requirement that WQ = WK, allowing for asymmetric kernels and enhancing representation capacity. The predictive mean in RFFCGP can be expressed as yd ≈ ΦQΦ⊤O ΦOΦ⊤O +σ2I−1ΦOΦ⊤K ΦKΦ⊤K +σ2I−1 vd. The token-wise predictive variance is given by diag(Σ) = σ2 diag(ΦQΦ⊤OΦO + σ2I−1 Φ⊤Q).

Empirical Results and Scalability

Empirical results on multiple real-world datasets show that the attention module improves calibration while maintaining predictive accuracy, and simultaneously reduces computational complexity to linear in the sequence length. RFF-GPA is strongest on the image benchmarks, attaining the best accuracy on Fashion-MNIST (0.915), CIFAR-10 (0.759), and SVHN (0.915). Furthermore, RFF-GPA substantially improves over SGPA on the image datasets, indicating that the random-feature GP formulation not only reduces cost but can also yield a more effective attention mechanism than inducing-point sparse GPs. The methods scale linearly with L, thanks to the RFF approximation that only inverts an M×M matrix, in contrast to the quadratic scaling of inducing-point variants like SGPA and SCGP.

Out-of-Distribution Robustness

The CIFAR-10-C evaluation confirms that the random-feature GP framework yields reliable uncertainty under distribution shift while preserving predictive performance. RFF-GPA is the most accurate method at every severity level, from the clean set (0.759) down to the most severe corruptions (0.514 at s=5), consistently ahead of the strong SNGP baseline and improving on the inducing-point SGPA by a large margin. The central finding is that RFF-GPA is the best-calibrated method across the corruption range in terms of NLL and ECE, attaining the lowest NLL at every severity and the lowest ECE at every non-clean severity (s≥1). This distinction matters because calibration under shift is meaningful only when accuracy collapses.

Conclusion

In this paper, we introduced a Gaussian process Transformer where we model the attention kernel using a random Fourier feature approximation. By taking advantage of the low-rank factorization of the attention module, we can obtain a fully Bayesian model that is scalable with respect to the sequence length. We have demonstrated the effectiveness of the RFFGPA on image and text data sets, especially with respect to uncertainty-aware metrics, compared to the vanilla Transformer and other competing GP-based methods. In future work, we hope to further enrich the approximate kernel structure by modeling non-stationarity across each head, or by assuming a deep generative model as a prior on the random frequencies.

References

Bochner, S. 1959. Lectures on Fourier Integrals, volume 42. Princeton University Press.

Brown, T.; Mann, B.; Ryder, N.; Subbiah, M.; Kaplan, J. D.; Dhariwal, P.; Neelakantan, A.; Shyam, P.; Sastry, G.; Askell, A., et al. 2020. Language models are few-shot learners. Advances in Neural Information Processing Systems.

Bui, L. M.; Huu, T. T.; Dinh, D.; Nguyen, T. M.; and Hoang, T. N., 2024. Revisiting Kernel Attention with Correlated Gaussian Process Representation.

Chen, W.; and Li, Y., 2023. Calibrating Transformers via sparse Gaussian processes.

Choromanski, K.; Likhosherstov, V.; Dohan, D.; Song, X.; Gane, A.; Sarlos, T.; Hawkins, P.; Davis, J., et al. 2021. Rethinking attention with performers.

Damianou, A., and Lawrence, N. D., 2013. Deep Gaussian processes.

Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M., et al. 2021. An image is worth 16x16 words: Transformers for image recognition at scale.

Fan, X.; Zhang, S.; Chen, B.; and Zhou, M., 2020. Bayesian attention modules.

Gal, Y.; and Ghahramani, Z., 2016. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning.

Guo, C.; Pleiss, G.; Sun, Y.; and Weinberger, K. Q., 2017. On calibration of modern neural networks <ref:2610.

Improvements for AI systems

Bold header: RFF-GPA for Scalable Uncertainty Awareness

The proposed RFF-GPA module allows for linear time complexity in approximating posterior mean and variance by approximat[ing] the attention kernel with random Fourier features. This enables models to achieve calibrated uncertainty while maintaining predictive accuracy in a way that is far more scalable compared to previous work.

Bold header: Correlated Attention for Asymmetric Modeling

The extension to RFF-CGP allows for modeling asymmetric query-key interactions by using a structure where the attention output is modeled as the cross-covariance between two correlated Gaussian processes, which naturally removes the requirement that WQ = WK.

Bold header: Robust Out-of-Distribution Generalization

The RFF-GPA model demonstrates superior robustness under distribution shift, as it is shown to be the most accurate method at every severity level and attains the lowest NLL at every severity and the lowest ECE at every non-clean severity, while avoid[ing] this failure mode [of deterministic MLE] by reducing overconfidence.

Bold header: Efficient Inference for Long Sequences

The framework achieves a computational advantage by replacing the expensive inversion of a L×L kernel matrix with an inversion of an M × M matrix, which reduces the cost to linear in the sequence length.