What would Bayes do for next-gen neural network design?

What Would Bayes Do
for Next-Gen Neural Network Design?

Yingzhen Li

yingzhen.li@imperial.ac.uk, yingzhen.li@ntu.edu.sg

Images by me (2019) & by Nano Banana 2 (gemini-3.1-flash-image, 2026)

pre-ChatGPT (- 2022)

AGI (when?)

What’s the Future of Bayesian Deep Learning?

Build network architectures

with Bayesian principles

Build Bayesian procedures

on frontier models

Many of you have contributed wonderful works in these developments ☺

\[(t_{2i-1},t_{2i})=(\mathbf{x}_i,\mathbf{y}_i)\]

\[(t_{2n+1},t_{2n+2})=(\mathbf{x}^{*},\mathbf{y}^{*})\]

\[{\color[RGB]{32,76,154}p(\mathbf{y}_i\mid\mathbf{x}_i,\mathbf{x}_{<i},\mathbf{y}_{<i})=p_{\phi}^{2i}(t_{2i}\mid t_{1:2i-1})}\]

\[{\color[RGB]{180,86,21}p(\mathbf{y}^{*}\mid\mathbf{x}^{*},\mathcal{D})=p_{\phi}^{2n+2}(t_{2n+2}\mid t_{1:2n+1})}\]

In-context learning & predictive resampling

Agentic AI: Bayesian pipeline

Bayesian NNs: weight/function-space posterior inference

NN layer design with Bayesian principles

MLSys work needed

(mostly absent now!)

”Modernization” via ”Bayesian LoRA” (?)

What Can Bayes Say about NN Arch Design?

Entrée: Uncertainty quantification

Main: Memory updates & maintenance

Naoki Kiyohara

Harrison Bo Hua Zhu

Riccardo El Hassanin

Wenlong Chen

Work led by my amazing current/former PhD students

Idea 1: Leverage probabilistic models to improve the “reliability” of attention-based Transformers (e.g., ”reliable uncertainty”)

Sparse Gaussian Process Attention

- a Deep GP tailored to Transformer architectures

Chen and Li, ICLR 2023,

Calibrating Transformers via Sparse Gaussian Processes.

Wenlong Chen

Attention in Transformers

Vaswani et al. Attention is all you need. NeurIPS 2017

Tsai et al. Transformer dissection: A unified understanding for transformer’s attention via the lens of kernel. EMNLP 2019

Single head attention

Attention matrix

Replace attention matrix with kernel matrix:

\[Attention(\mathbf{q},\mathbf{k},\mathbf{v})={\color[rgb]{0,0,1}activation(\mathbf{q}\mathbf{k}^{\top})}\mathbf{v}\]

\[KernelAttention(\mathbf{q},\mathbf{k},\mathbf{v})={\color[rgb]{0,0,1}\mathbf{K}_{\mathbf{qk}}}\mathbf{v}\]

Kernel Attention as the Mean of an SVGP

Kernel Attention:

\[\mathbf{F}={\color[rgb]{0,0,1}\mathbf{K}_{\mathbf{qk}}}{\color[rgb]{1,0,0}\mathbf{v}}\]

Recall posterior mean of SVGP:

Equivalent by identifying:

q (queries) = x (queried input locations)

k (keys) = z (inducing locations)

v (values) = $γ$ (variational parameters)

\[[{\color[rgb]{0,0,1}\mathbf{K}_{\mathbf{qk}}}]_{ij}\]

similarity between and

\[\mathbf{q}_i\]

\[\mathbf{k}_j\]

Chen and Li. Calibrating Transformers via Sparse Gaussian Processes. ICLR 2023

$γ = {K_{zz}^{- 1} m}_{z}$

$μ = K_{xz} γ$

$m_{z}$ is a free-parameter, $K_{zz}$ is invertible

=> Can optimize $γ$ as a free-parameter

Adding Covariance Function to Transformer

\[\boldsymbol{\Sigma}^{(post)}=\mathbf{K}_{\mathbf{qq}}+\mathbf{K}_{\mathbf{qk}}(\mathbf{K}_{\mathbf{kk}}^{-1}{\color[rgb]{1,0,0}\mathbf{S}}\mathbf{K}_{\mathbf{kk}}^{-1}-\mathbf{K}_{\mathbf{kk}}^{-1})\mathbf{K}_{\mathbf{kq}}\]

\[\mathbf{m}^{(post)}={\color[rgb]{0,0,1}\mathbf{K}_{\mathbf{qk}}}{\color[rgb]{1,0,0}\mathbf{v}}\]

Chen and Li. Calibrating Transformers via Sparse Gaussian Processes. ICLR 2023

Amortized Inference for Self-Attention

Chen and Li. Calibrating Transformers via Sparse Gaussian Processes. ICLR 2023

Computation Reduction for Self-Attention

Chen and Li. Calibrating Transformers via Sparse Gaussian Processes. ICLR 2023

Cheng and Boots. Variational Inference for Gaussian Process Models with Linear Complexity. NIPS 2017

Chen and Li. Calibrating Transformers via Sparse Gaussian Processes. ICLR 2023

OOD Detection

\[\widehat{Y}=\mathbb{1}\{H[\hat{p}]>\tau\}\]

ViTs trained on CIFAR10

Reflections on SGPA

Kernel attention = mean of SVGP

So can think about how GP methods inspire memory design

SGPA = Deep Sparse GP modernised with Transformer arch.

So still fall in the traditional BNN category

Evaluations on calibration & OOD detection

Frequentist methods are computationally simpler and competitive here

Chen and Li. Calibrating Transformers via Sparse Gaussian Processes. ICLR 2023

Idea 2:

Exploit the kernel interpretation of attention;

Introduce interdomain methods for ”memory maintenance”

Interdomain Attention

– a hybrid attention that implicitly approximate attention with recurrent memory

Kiyohara et al., In submission,

Interdomain Attention:

Beyond Token-Level Key-Value Memory.

Naoki Kiyohara

Harrison Bo Hua Zhu

Riccardo El Hassanin

Wenlong Chen

KV Cache 101

1. Prefill (parallel compute)

Use the last output token $x_{N}$ as input

Read from KV cache for ${k_{n}, v_{n}}_{n = 1}^{N - 1}$

Compute and write $k_{N}, v_{N}$ in KV cache, and compute with query $q_{i}$ , $i = N$ :

The output $o_{i}$ propagates and leads to generated token $x_{N + 1}$

\[o_i=\frac{\sum_{n=1}^{N}\mathcal{K}(q_i,k_n)\,v_n}{\sum_{n^{\prime}=1}^{N}\mathcal{K}(q_i,k_{n^{\prime}})}\]

2. Decode (sequential compute)

LLM “inference” compute in two stages:

User provides prompt token sequence ${x_{n}}_{n = 1}^{N}$

Compute in parallel ${q_{n}, k_{n}, v_{n}}_{n = 1}^{N}$

Cache ${k_{n}, v_{n}}_{n = 1}^{N}$ as memory (KV cache)

https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/

Growing memory size…

A New Type of Memory

Consider approximating kernel as dot-product of features

\[\mathcal{K}(q,k)\approx\xi(q)^{\top}\xi(k)\]

\[o_i=\frac{\sum_{n=1}^{N}\mathcal{K}(q_i,k_n)\,v_n}{\sum_{n^{\prime}=1}^{N}\mathcal{K}(q_i,k_{n^{\prime}})}\]

The $i^{t h}$ output of the attention procedure (with query $q_{i}$ ):

\[o_i\approx\hat{o}_i=\frac{\sum_{n=1}^{N}\xi(q_i)^{\top}\xi(k_n)\,v_n}{\sum_{n^{\prime}=1}^{N}\xi(q_i)^{\top}\xi(k_{n^{\prime}})}\]

$\Rightarrow$

Idea in a nutshell:

Step 1: Construct memory for $ξ (k_{n}), v_{n}$

Step 2: “reconstruct” $ξ (k_{n})$ from the memory

Step 3: Plug-in the “reconstructions” to perform attention

Steps 2 & 3 can be done together!

(no explicit reconstruction & attention)

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

A New Type of Memory

Key idea: function projection as memory

\[o_i\approx\hat{o}_i=\frac{\sum_{n=1}^{N}\xi(q_i)^{\top}\xi(k_n)\,v_n}{\sum_{n^{\prime}=1}^{N}\xi(q_i)^{\top}\xi(k_{n^{\prime}})}\]

Step 1.1. Treat the keys and values as functions of time

$k (t_{n}) : = k_{n}$

$v (t_{n}) : = v_{n}$

Step 1.2. Project key features and values onto a set of basis functions ${ϕ_{m}^{(t_{N})}}_{m = 1}^{M}$ :

\[{\color[RGB]{76,134,37}u_m^{(t_N)}=\int\xi(k(t))\,\phi_m^{(t_N)}(t)\,dt}\]

\[{\color[RGB]{32,76,154}\gamma_m^{(t_N)}=\int v(t)\,\phi_m^{(t_N)}(t)\,dt}\]

Step 2 + 3 (virtual): Reconstruct the key features

+ perform attention (assume we have the values $v_{n}$ )

\[{\color[RGB]{76,134,37}\xi(k_n)\approx\sum_m u_m\,\phi_m(t_n)}\]

$\sum_{n = 1}^{N} ξ {(q_{i})}^{⊤} ξ (k_{n}) v_{n} \approx ξ {(q_{i})}^{⊤} \sum_{n = 1}^{N} \sum_{m}^{} u_{m} ϕ_{m} (t_{n}) v_{n} = ξ {(q_{i})}^{⊤} \sum_{m}^{} u_{m} \sum_{n}^{N} ϕ_{m} (t_{n}) v_{n} \approx ξ {(q_{i})}^{⊤} \sum_{m}^{} u_{m} γ_{m}$

discretized sum (up to const.)

reconstruct

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

A New Type of Memory

Key idea: function projection as memory

\[o_i\approx\hat{o}_i=\frac{\sum_{n=1}^{N}\xi(q_i)^{\top}\xi(k_n)\,v_n}{\sum_{n^{\prime}=1}^{N}\xi(q_i)^{\top}\xi(k_{n^{\prime}})}\]

Step 1.1. Treat the keys and values as functions of time

$k (t_{n}) : = k_{n}$

$v (t_{n}) : = v_{n}$

Step 1.2. Project key features and values onto a set of basis functions ${ϕ_{m}^{(t_{N})}}_{m = 1}^{M}$ :

\[{\color[RGB]{76,134,37}u_m^{(t_N)}=\int\xi(k(t))\,\phi_m^{(t_N)}(t)\,dt}\]

\[{\color[RGB]{32,76,154}\gamma_m^{(t_N)}=\int v(t)\,\phi_m^{(t_N)}(t)\,dt}\]

\[{\color[RGB]{180,86,21}\eta_m^{(t_N)}=\int\phi_m^{(t_N)}(t)\,dt}\]

Step 2 + 3 (virtual): Reconstruct the key features

+ perform attention (denominator, optional)

\[{\color[RGB]{76,134,37}\xi(k_n)\approx\sum_m u_m\,\phi_m(t_n)}\]

$\sum_{n = 1}^{N} ξ {(q_{i})}^{⊤} ξ (k_{n}) \approx ξ {(q_{i})}^{⊤} \sum_{n = 1}^{N} \sum_{m}^{} u_{m} ϕ_{m} (t_{n}) = ξ {(q_{i})}^{⊤} \sum_{m}^{} u_{m} \sum_{n}^{N} ϕ_{m} (t_{n}) \approx ξ {(q_{i})}^{⊤} \sum_{m}^{} u_{m} η_{m}$

discretized sum (up to const.)

reconstruct

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

A New Type of Memory

Key idea: function projection as memory

\[o_i\approx\hat{o}_i=\frac{\sum_{n=1}^{N}\xi(q_i)^{\top}\xi(k_n)\,v_n}{\sum_{n^{\prime}=1}^{N}\xi(q_i)^{\top}\xi(k_{n^{\prime}})}\]

Step 1.1. Treat the keys and values as functions of time

$k (t_{n}) : = k_{n}$

$v (t_{n}) : = v_{n}$

Step 1.2. Project key features and values onto a set of basis functions ${ϕ_{m}^{(t_{N})}}_{m = 1}^{M}$ :

\[{\color[RGB]{76,134,37}u_m^{(t_N)}=\int\xi(k(t))\,\phi_m^{(t_N)}(t)\,dt}\]

\[{\color[RGB]{32,76,154}\gamma_m^{(t_N)}=\int v(t)\,\phi_m^{(t_N)}(t)\,dt}\]

Step 2 + 3: Implement reconstruction + approximate attention in one go!

${\hat{o}}_{i} \approx {\tilde{o}}_{i} = \frac{\sum_{m = 1}^{M} (ξ {(q_{i})}^{⊤} u_{m}) γ_{m}}{\sum_{m = 1}^{M} (ξ {(q_{i})}^{⊤} u_{m}) η_{m}}$

\[{\color[RGB]{180,86,21}\eta_m^{(t_N)}=\int\phi_m^{(t_N)}(t)\,dt}\]

$\tilde{O} = (F_{q} U^{⊤} Γ)$ $⊘ (F_{q} U^{⊤} η)$

\[F_q=[\xi(q_1),\ldots,\xi(q_{N_q})]^{\top}\]

\[{\color[RGB]{76,134,37}U=[u_1,\ldots,u_M]^{\top}}\]

\[{\color[RGB]{49,110,165}\Gamma=[\gamma_1,\ldots,\gamma_M]^{\top}}\]

\[{\color[RGB]{180,86,21}\eta=[\eta_1,\ldots,\eta_M]^{\top}}\]

query features

“KV” memory

Helper terms

$\Rightarrow$

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Connection to Interdomain Inducing Points in SVGPs

Lázaro-Gredilla and Figueiras-Vidal. Inter-domain Gaussian Processes for Sparse Inference using Inducing Features. NIPS 2009

How to construct this function projection in an online manner?

$γ = {K_{ZZ}^{- 1} m}_{z}$

SVGP approximate posterior mean:

$μ = K_{XZ} γ$

$m_{z}$ is a free-parameter, $K_{ZZ}$ is invertible

=> Can optimize $γ$ as a free-parameter

Interdomain inducing points: projections on orthogonal basis functions:

${(K_{XZ})}_{n m} : = Cov (f (x_{n}), \int f (x) ϕ_{m} (x) d x) = \int K (x_{n}, x) ϕ_{m} (x) d x$

$\approx ⟨ ξ (x_{n}), \int ξ (x) ϕ_{m} (x) d x ⟩$

(RFF approx.)

$μ_{n} = K_{x_{n} Z} γ \approx ξ {(x_{n})}^{⊤} \sum_{m = 1}^{M} \int ξ (x) ϕ_{m} (x) d x γ_{m}$

$\approx ξ {(x_{n})}^{⊤} \sum_{m = 1}^{M} \int \sum_{m^{'} = 1}^{M} u_{m^{'}} ϕ_{m^{'}} (x) ϕ_{m} (x) d x γ_{m}$

Plug-in to the mean computation, with $γ : = {[γ_{1}, \dots, γ_{m}]}^{⊤}$ :

$= ξ {(x_{n})}^{⊤} \sum_{m, m'}^{M} u_{m^{'}} δ (m^{'} = m) γ_{m}$

$= ξ {(x_{n})}^{⊤} \sum_{m = 1}^{M} u_{m} γ_{m} = ξ {(x_{n})}^{⊤} U γ$

(reconstruct $ξ (x)$ )

( $\int ϕ_{m^{'}} (x) ϕ_{m} (x) d x = δ (m^{'} = m)$ )

(rewrite matrix-vector product)

HiPPO - An Online Representation of Sequential Data

Legendre polynomial

“Memorising” a function via projection to finite basis:

$Given f (x), x \in [- 1, 1],$

u – coefficients of $f$ projected to $s p a n {P_{n} (x)}_{n = 0}^{N - 1}$

P_{n} (x), x \in [- 1, 1]

f (x) \approx \sum_{n = 0}^{N - 1} u_{n} P_{n} (x), u_{n} = \int_{- 1}^{1} f (x) P_{n} (x) d x

Can be viewed as a finite-dim memory for a function (infinite-dim object)

Fig: wikipedia

Gu et al. Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020

$u_{n} : = \int_{- 1}^{1} f (x) P_{n} (x) d x$

HiPPO - An Online Representation of Sequential Data

“Memorising” a function via projection to finite basis:

$Given f (x), x \in [0, t],$

$u_{n}^{(t)}$ – coefficients of $f$ projected to $s p a n {P_{n}^{(t)} (x)}_{n = 0}^{N - 1}$

Rescaled Legendre polynomial

w^{(t)} (x) = \frac{1}{t} 1_{x \in [0, 1]}

1/t

f (x) \approx \sum_{n = 0}^{N - 1} u_{n}^{(t)} P_{n}^{(t)} (x), u_{n}^{(t)} = \int_{0}^{t} f (x) P_{n}^{(t)} (x) d x

Can be viewed as a finite-dim memory for a function (infinite-dim object)

Memory “evolves” when $t$ increases!

Gu et al. Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020

$u_{n} ≔ \int f (x) P_{n}^{(t)} (x) w^{(t)} (x) d x$

$P_{n}^{(t)} (x) : = P_{n} ((2 x - t) / t)$

$x$

$f (x)$

$t$

$0$

$\frac{1}{t}$

$ω^{(t)} (x)$

$x$

$g_{n}^{(t)} (x)$

$t + Δ t$

$ω^{(t + Δ t)} (x)$

$g_{n}^{(t + Δ t)} (x)$

$u_{n}^{(t)} = \int_{- \infty}^{\infty} f (x) g_{n}^{(t)} (x) ω^{(t)} (x) d x$

$u_{n}$ : $n$ -th coefficient

$f (x)$ : Target function for $x \in [0, + \infty)$

$g_{n}^{(t)} (x)$ : Normalised & scaled $n$ -th basis

$ω^{(t)} (x)$ : Normalised measure (mask)

HiPPO - An Online Representation of Sequential Data

Gu et al. Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020

g_{n}^{(t)} (x) = P_{n} (\frac{2 x}{t} - 1)

The evolution of $u^{(t)} = [u_{0}^{(t)}, \dots, u_{N - 1}^{(t)}]$ over time $t$ follows linear ODE:

$\frac{d}{d t} u^{(t)} = A (t) u^{(t)} + B (t) f (t)$

We can obtain the coefficients $u^{(t)}$
as a summary of the function up to time t in an online manner.

Input sequence to memorize

Specific matrix and vector corresponding to function basis and measure

Gu et al. Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020

HiPPO - An Online Representation of Sequential Data

Gu et al. Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020

HiPPO - An Online Representation of Sequential Data

Interlude: Recurrent Memory for GPs!

Chen et al. Recurrent Memory for Online Interdomain Gaussian Processes. NeurIPS 2025

Interlude: Recurrent Memory for GPs!

Chen et al. Recurrent Memory for Online Interdomain Gaussian Processes. NeurIPS 2025

Bui et al. Streaming Sparse Gaussian Process Approximations. NIPS 2017

KV cache ⇔ growing number of data-space inducing points

Recurrent memory ⇔ interdomain inducing points (fixed state dim. = fixed #basis)

Interdomain Attention

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Interdomain Attention

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Yang et al. Gated Linear Attention Transformers with Hardware-Efficient Training. ICML 2024

Gu et al. On the Parameterization and Initialization of Diagonal State Space Models. NeurIPS 2022

Practical Implementation (with deep learning magic 🪄)

1. Kernel feature: $ξ (\cdot) = L 2 Norm (SiLU (conv 1 D (\cdot)))$ + RoPE

Tried RFF with (learnable) frequencies, didn’t work well

2. Apply RMSNorm to $ξ (k_{t})$ and $v_{t}$ before entering the recurrent memory

Mainly to stabilise the recurrent memory updates

$v_{t}$ in SVGPs are “free parameters” anyway

3. Complex S4D as the recurrent memory module

Several gating and selective update mechanism tried

4. Optional: construct $q$ from input $x$ with non-linear projection

Combined with (2): the kernel now becomes asymmetric

Extensive amount of engineering for GPU kernel implementations!

(our implementation is still sub-optimal even with chunking)

IA Empirical Performances: Text

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Perplexity = exp(avg. NLL)

IA Empirical Performances: Text

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Downstream tasks (reasoning) performance

IA Empirical Performances: Genomics

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Downstream tasks performance

IA Empirical Performances: Ablations

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission

Query-conditioned read-out

is the main improvement factor

Some deep learning magic 🪄

also helps

Reflections on Interdomain Attention

Recurrent memory $\approx$ interdomain inducing points

So can think about how GP methods inspire memory design

Think about sequence layers as online learners

Now can think about network layer design in a “statistical modelling” way!

Bayes can already say quite something for NN design:

“knowledge adaptation” $\approx$ online learning + memory management

Engineering & MLSys work may feel less “pleasant”:

We had to throw away some “elegant ideas”

But what if: hardware info should enter in your “prior beliefs” of models?

Kiyohara et al. Interdomain Attention: Beyond Token-Level Key-Value Memory. In submission