SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ryan Singh
School of Engineering and Informatics,
rs773@sussex.ac.uk
Christopher L. Buckley
VERSES AI Research Lab,
Los Angeles, CA, USA.
Attention mechanisms play a crucial role in cognitive systems by allowing them to flexibly allocate cognitive resources. Transformers, in particular, have become a dominant architecture in machine learning, with attention as their central innovation. However, the underlying intuition and formalism of attention in Transformers is based on ideas of keys and queries in database management systems. In this work, we pursue a structural inference perspective, building upon, and bringing together, previous theoretical descriptions of attention such as; Gaussian Mixture Models, alignment mechanisms and Hopfield Networks. Specifically, we demonstrate that attention can be viewed as inference over an implicitly defined set of possible adjacency structures in a graphical model, revealing the generality of such a mechanism. This perspective unifies different attentional architectures in machine learning and suggests potential modifications and generalizations of attention. Here we investigate two and demonstrate their behaviour on explanatory toy problems: (a) extending the value function to incorporate more nodes of a graphical model yielding a mechanism with a bias toward attending multiple tokens; (b) introducing a geometric prior (with conjugate hyper-prior) over the adjacency structures producing a mechanism which dynamically scales the context window depending on input. Moreover, by describing a link between structural inference and precision-regulation in Predictive Coding Networks, we discuss how this framework can bridge the gap between attentional mechanisms in machine learning and Bayesian conceptions of attention in Neuroscience. We hope by providing a new lens on attention architectures our work can guide the development of new and improved attentional mechanisms.
Designing neural network architectures with favourable inductive biases lies behind many recent successes in Deep Learning. The Transformer, and in particular the attention mechanism has allowed language models to achieve human like generation abilities previously thought impossible [1, 2]. The success of the attention mechanism as a domain agnostic architecture has prompted adoption across a diverse range of tasks beyond language modelling, notably reaching state-of-the-art performance in visual reasoning and segmentation tasks [3, 4].
This depth and breadth of success indicates the attention mechanism expresses a useful computational primitive. Recent work has shown interesting theoretical links to kernel methods [5, 6, 7], Hopfield networks [8], and Gaussian mixture models [9, 10, 11, 12, 13], however a formal understanding that captures the generality of this computation remains outstanding. In this paper, we show the attention mechanism can naturally be described as inference on the structure of a graphical model, agreeing with observations that transformers are able to flexibly choose between models based on context [14, 15]. This Bayesian perspective complements previous theory [16, 8, 12], adding new
methods for reasoning about inductive biases and the functional role of attention variables. Further, understanding the core computation as inference permits a unified description of multiple attention mechanisms in the literature as well as narrowing the explanatory gap to ideas in neuroscience.
This paper proceeds in three parts: First in Sec.3, we show that 'soft' attention mechanisms (e.g. self-attention, cross-attention, graph attention, which we call transformer attention hereafter) can be understood as taking an expectation over possible connectivity structures, providing an interesting link between softmax-based attention and marginal likelihood. Second in Sec.4, we extend the inference over connectivity to a Bayesian setting which, in turn, provides a theoretical grounding for iterative attention mechanisms (slot-attention and block-slot attention) [17, 18, 19], Modern Continuous Hopfield Networks [8] and Predictive Coding Networks. Finally in Sec.5, we leverage the generality of this description in order to design new mechanisms with predictable properties.
Intuitively, the attention matrix can be seen as the posterior distribution over edges  $E$  in a graph,  $\mathcal{G} = (K \cup Q, E)$  consisting of a set of query and key nodes  $Q, K$  each of dimension  $d$ . Where the full mechanism computes an expectation of a function defined on the graph  $V: \mathcal{G} \to \mathbb{R}^{d \times |\mathcal{G}|}$  with respect to this posterior.
$$
\begin{array}{l} A t t e n t i o n (Q, K, V) = \overbrace {\operatorname {s o f t m a x} \left(\frac {Q W _ {Q} W _ {K} ^ {T} K ^ {T}}{\sqrt {d}}\right)} ^ {p (E \mid Q, K)} W _ {V} K \\ = \mathbb {E} _ {p (E | Q, K)} [ V ] \\ \end{array}
$$
Crucially, when  $\mathcal{G}$  is seen as a graphical model, the posterior over edges becomes an inference about dependency structure and the functional form becomes natural. This formalism provides an alternate Bayesian theoretical framing within which to understand attention models, shifting the explanation from one centred around retrieval to one that is fundamentally concerned with in-context inference of probabilistic relationships (including retrieval). Within this framework different attention architectures can be described by considering different implicit probabilistic models, by making these explicit we hope to support more effective analysis and the development of new architectures.
A key benefit of the perspective outlined here is to tie together different approaches taken in the literature. Specifically, structural variables can be seen as the alignment variables discussed in previous Bayesian descriptions [16, 20, 21], on the other hand Gaussian Mixture Models (GMMs) can be seen as a specific instance of the framework developed here. This description maintains the explanatory power of GMMs by constraining the alignment variables to be the edges of an implicit graphical model, while offering the increased flexibility of alignment approaches to describe multiple forms of attention.
Latent alignment and Bayesian Attention, several attempts have been made to combine the benefits of soft (differentiability) and stochastic attention, often viewing attention as a probabilistic alignment problem. Most approaches proceed by sampling, e.g., using the REINFORCE estimator [20] or a top  $K$  approximation [22]. Two notable exceptions are [16] which embeds an inference algorithm within the forward pass of a neural network, and [21] which employs the re-parameterisation trick for the alignment variables. In this work, rather than treating attention weights as an independent learning problem, we aim to provide a parsimonious implicit model that would give rise to the attention weights. Additionally showing that 'soft' attention weights arise naturally in variational inference from either collapsed variational inference or a mean-field approximation.
Relationship to Gaussian mixture model, previous works that have taken a probabilistic perspective on the attention mechanism note the connection to inference in a gaussian mixture model [11, 10, 12, 13]. Indeed [23] directly show the connection between the Hopfield energy and the variational free energy of a Gaussian mixture model. Although Gaussian mixture models, a special case of the framework we present here, are enough to explain cross attention they do not capture slot or self-attention, obscuring the generality underlying attention mechanisms. In contrast, the description presented here extends to structural inductive biases beyond what can be expressed in a Gaussian mixture model, additionally offering a route to describing the whole transformer block.
Attention as bi-level optimisation, mapping feed-forward architecture to a minimisation step on a related energy function has been called unfolded optimisation [24]. Taking this perspective can lead to insights about the inductive biases involved for each architecture. It has been shown that the cross-attention mechanism can be viewed as an optimisation step on the energy function of a form of Hopfield Network [8], providing a link between attention and associative memory. while [25] extend this view to account for self-attention. Our framework distinguishes Hopfield attention, which does not allow an arbitrary value matrix, from transformer attention. Although there remains a strong theoretical connection, we interpret the Hopfield Energy as an instance of variational free energy, aligning more closely with iterative attention mechanisms such as slot-attention.
3 Transformer Attention
3.1 Attention as Expectation
We begin by demonstrating transformer attention can be seen as calculating an expectation over graph structures. Specifically, let  $x = (x_{1},\dots,x_{n})$  be observed input variables,  $\phi$  be some set of discrete latent variables representing edges in a graphical model of  $x$  given by  $p(x\mid \phi)$ , and  $y$  a variable we need to predict. Our goal is to find  $\mathbb{E}_{y|x}[y]$ , however the graph structure  $\phi$  is unobserved so we calculate the marginal likelihood.
$$
\mathbb {E} _ {y \mid x} [ y ] = \sum_ {\phi} p (\phi \mid x) \mathbb {E} _ {y \mid x, \phi} [ y ]
$$
Importantly, the softmax function is a natural representation for the posterior,
$$
p (\phi \mid x) = \frac {p (x , \phi)}{\sum_ {\phi} p (x , \phi)} = s o f t m a x (\ln p (x, \phi))
$$
in order to expose the link to transformer attention, let the model of  $y$  given the graph  $(x,\phi)$  be parameterised by a function  $\mathbb{E}_{y|x,\phi}[y] = v(x,\phi)$ .
$$
\mathbb {E} _ {y \mid x} [ y ] = \sum_ {\phi} \operatorname {s o f t m a x} (\ln p (x, \phi)) v (x, \phi) = \mathbb {E} _ {\phi \mid x} [ v (x, \phi) ] \tag {1}
$$
In general, transformer attention can be seen as weighting  $v(x,\phi)$  by the posterior distribution  $p(\phi \mid x)$  over different graph structures. We show Eq.1 is exactly the equation underlying self and cross-attention by presenting the specific generative models corresponding to them. In this description the latent variables  $\phi$  are identified as edges between observed variables  $x$  (keys and queries) in a pairwise Markov Random Field, parameterised by matrices  $W_{K}$  and  $W_{Q}$ , while the function  $v$  is parameterised by  $W_{V}$ .
Pairwise Markov Random Fields are a natural tool for modelling the dependencies of random variables, with prominent examples including Ising models (Boltzmann Machine) and multivariate Gaussians. While typically defined given a known structure, the problem of inferring the latent graph is commonly called structural inference.
Formally, given a set of random variables  $X = (X_{v})_{v\in V}$  with probability distribution  $[p]$  and a graph  $G = (V,E)$ . The variables form a pairwise Markov Random Field (pMRF) [26] with respect to  $G$  if the joint density function  $P(X = x) = p(x)$  factorises as follows
$$
p (x) = \frac {1}{Z} \exp \left(\sum_ {v \in V} \psi_ {v} + \sum_ {e \in E} \psi_ {e}\right)
$$
where  $Z$  is the partition function  $\psi_v(x_v)$  and  $\psi_e = \psi_{u,v}(x_u,x_v)$  are known as the node and edge potentials respectively. Bayesian structural inference also requires a structural prior  $p(\phi)$  over the space of possible adjacency structures,  $\phi \in \Phi$ , of the underlying graph.
Factorisation, without constraints this space grows exponentially in the number of nodes  $(2^{|V|}$  possible graphs leading to intractable softmax calculations), all the models we explore here implicitly assume a factorised prior<sup>1</sup>. We briefly remark that Eq.1 respects factorisation of  $[p]$  in the following
1Additionally placing zero probability mass on much of the space, for example disconnected graphs.
sense; if the distribution admits a factorisation (a partition of the space of graphs  $\Phi = \prod_{i} \Phi_{i}$ ) with respect to the latent variables  $p(x, \phi) = \prod_{i} e^{f_{i}(x, \phi_{i})}$  where  $\phi_{i} \in \Phi_{i}$ , and the value function distributes over the same partition of edges  $v(x, \phi) = \sum_{i} v_{i}(x, \phi_{i})$  then each of the factors can be marginalised independently:
$$
\mathbb {E} _ {\phi \mid x} [ v (x, \phi) ] = \sum_ {i} \mathbb {E} _ {\phi_ {i} \mid x} [ v _ {i} ] \tag {2}
$$
To recover cross-attention and self-attention we need to specify the structural prior, potential functions and a value function. (In order to ease notation, when  $\Phi_i$  is a set of edges involving a common node  $x_i$ , such that  $\phi_i = (x_i, x_j)$  represents a single edge, we use the notation  $\phi_i = [j]$ , suppressing the shared index.)
3.2 Cross Attention and Self Attention
We first define the model that gives rise to cross-attention:
- Key nodes  $K = \left( {{x}_{1},..,{x}_{n}}\right)$  and query nodes  $Q = \left( {{x}_{1}^{\prime },\ldots ,{x}_{m}^{\prime }}\right)$
- Structural prior  $p(\phi) = \prod_{i=1}^{m} p(\phi_i)$ , where  $\Phi_i = \{(x_1, x_i'),.., (x_n, x_i')\}$  is the set of edges involving  $x_i'$  and  $\phi_i \sim \text{Uniform}(\Phi_i)$  such that each query node is uniformly likely to connect to each key node.
- Edge potentials  $\psi(x_j, x_i') = x_i'^T W_Q^T W_K x_j$ , in effect measuring the similarity of  $x_j$  and  $x_i'$  in a projected space.
- Value functions  $v_{i}(x, \phi_{i} = [j]) = W_{V} x_{j}$ , a linear transformation applied to the node at the start of the edge  $\phi_{i}$ .
Taking the expectation with respect to the posterior in each of the factors defined in Eq.2 gives the standard cross-attention mechanism,
$$
\mathbb {E} _ {p (\phi_ {i} | Q, K)} [ v _ {i} ] = \sum_ {j} \operatorname {s o f t m a x} _ {j} \left(x _ {i} ^ {\prime T} W _ {Q} ^ {T} W _ {K} x _ {j}\right) W _ {V} x _ {j}
$$
If the key nodes are in fact the same as the query nodes and the prior is instead over a directed graph we recover self-attention (A.8.1).
4 Iterative Attention
We continue by extending attention to a latent variable setting, where not all the nodes are observed. In essence applying the attention trick, i.e., a marginalisation of structural variables, to a variational free energy (Evidence Lower Bound). This allows us to recover models such as slot attention [17] and block-slot attention [18]. These mechanisms utilise an EM-like procedure using the current estimation of latent variables to infer the structure and then using the inferred structure to improve estimation of latent variables. Interestingly, Modern Continuous Hopfield Networks fit within this paradigm rather than the one discussed in Sec.3; collapsed variational inference produces an identical energy function to the one proposed by Ramsauer et al. [8].
4.1 Collapsed Inference
We present a version of collapsed variational inference [27], where the collapsed variables  $\phi$  are again structural, showing how this results in a Bayesian attention mechanism. In contrast to the previous section, we have a set of (non-structural) latent variables  $z$ . The goal is to infer  $z$  given the observed variables,  $x$ , and a latent variable model  $p(x,z,\phi)$ . Collapsed inference proceeds by marginalising out the extraneous latent variables  $\phi$  [27]:
$$
p (x, z) = \sum_ {\phi} p (x, z, \phi) \tag {3}
$$
We define a gaussian recognition density  $q(z) \sim N(z; \mu, \Sigma)$  and optimise the variational free energy  $\mathcal{F}(\lambda) = \mathbb{E}_q[\ln q_\lambda(z) - \ln p(x, z)]$  with respect to the parameters,  $\lambda = (\mu, \Sigma)$ , of this distribution. Application of Laplace's method yields approximate derivatives of the variational
Figure 1: Comparison of models involved in different attention mechanisms. In each case, the highlighted edges indicate  $\Phi_{i}$  the support of the uniform prior over  $\phi_{i}$ . Attention proceeds by calculating a posterior over these edges, given the current state of the nodes, before using this inference to calculate an expectation of the value function  $v$ . For iterative attention mechanisms the value function can be identified as the gradient of a variational free energy, in contrast, transformer attention uses a learnable function.
free energy  $\nabla_{\mu}\mathcal{F}\approx -\nabla_{\mu}\ln p(x,\mu)$  and  $\nabla_{\Sigma}\mathcal{F}\approx -\nabla_{\mu}^{2}\ln p(x,\mu)$ , here we focus on the first order terms<sup>2</sup>. Substituting in Eq.3:
$$
\begin{array}{l} \nabla_ {\mu} \mathcal {F} \approx - \nabla_ {\mu} \ln \sum_ {\phi} p (x, \mu , \phi) (4) \\ = - \frac {1}{\sum_ {\phi} p (x , \mu , \phi)} \sum_ {\phi} \nabla_ {\mu} p (x, \mu , \phi) (5) \\ \end{array}
$$
In order to make the link to attention, we employ the log-derivative trick, substituting  $p(\cdot) = e^{\ln p(\cdot)}$  and re-express Eq.5 in two ways:
$$
\begin{array}{l} = - \sum_ {\phi} \operatorname {s o f t m a x} _ {\phi} (\ln p (x, \mu , \phi)) \nabla_ {\mu} \ln p (x, \mu , \phi) (6) \\ = \mathbb {E} _ {\phi | x, \mu} \left[ - \nabla_ {\mu} \ln p (x, \mu , \phi) \right] (7) \\ \end{array}
$$
The first form reveals the softmax which is ubiquitous in all attention models. The second, suggests the variational update should be evaluated as the expectation of the typical variational gradient (the term within the square brackets) with respect to the posterior over the parameters represented by the random variable  $\phi$ . In other words, iterative attention is exactly transformer attention applied iteratively where the value function is the variational free energy gradient. We derive updates for a general pMRF before again recovering (iterative) attention models in the literature by specifying particular distributions.
Free Energy of a marginalised pMRF, recall the factorised pMRF,  $p(x,\phi) = \frac{1}{Z}\prod_{i}e^{f_{i}(x,\phi_{i})}$ . Again, independence properties simplify the calculation, the marginalisation can be expressed as a product of local marginals,  $\sum_{\phi}p(x,\phi) = \frac{1}{Z}\prod_{i}\sum_{\phi_{i}}e^{f_{i}(x,\phi_{i})}$ . Returning to the inference setting, the nodes are partitioned into observed nodes,  $x$ , and variational parameters  $\mu$ . Hence the (approximate) collapsed variational free energy Eq.5, can be expressed as,  $F(x,\mu) = -\sum_{i}\ln \sum_{\phi_{i}}e^{f_{i}(x,\mu ,\phi_{i})} + C$  and it's derivative:
$$
\frac {\partial F}{\partial \mu_ {j}} = - \sum_ {i} \sum_ {\phi_ {i}} \operatorname {s o f t m a x} \left(f _ {i}\right) \frac {\partial f _ {i}}{\partial \mu_ {j}}
$$
Finally, we follow [8] in using the Convex-Concave Procedure (CCCP) to derive a simple fixed point equation which necessarily reduces the free energy.
Quadratic Potentials and the Convex Concave Procedure, assuming the node potentials are quadratic  $\psi (x_{i}) = -\frac{1}{2} x_{i}^{2}$  and the edge potentials have the form  $\psi (x_i,x_j) = x_iWx_j$  , and define  $\tilde{f}_i = \sum_{e\in \Phi_i}\psi_e$  . Consider the following fixed point equation,
$$
\mu_ {j} ^ {*} = \sum_ {i} \sum_ {\phi_ {i}} \operatorname {s o f t m a x} \left(\tilde {f} _ {i}\right) \frac {\partial \tilde {f} _ {i}}{\partial \mu_ {j}} \tag {8}
$$
As the first order terms are independent of the second order ones, see A.7.1 for details.
since (under mild conditions) node potentials are convex and edge potentials are concave (A.7.1.1), we can invoke the CCCP [28] to show this fixed point equation descends on the energy  $F(x, \mu_j^*) \leq F(x, \mu_j)$  with equality if and only if  $\mu_j^*$  is a stationary point of  $F$ . We follow Sec.3 in specifying specific structural priors and potential functions that recover different iterative attention mechanisms.
4.2 Modern Continuous Hopfield Network
Let the observed, or memory, nodes  $x = (x_{1},\dots,x_{n})$  and latent nodes  $z = (z_{1},\dots,z_{m})$  have the following structural prior  $p(\phi) = \prod_{i=1}^{m}p(\phi_{i})$ , where  $\phi_{i} \sim \text{Uniform}\{(x_{1},z_{i}),\dots,(x_{n},z_{i})\}$ , meaning each latent node is uniformly likely to connect to a memory node. Define edge potentials  $\psi(x_{j},z_{i}) = z_{i}^{T}x_{j}$ . Application of Eq.8:
$$
\mu_ {i} ^ {*} = \sum_ {j} s o f t m a x _ {j} (\mu_ {i} ^ {T} x _ {j}) x _ {j}
$$
When  $\mu_{i}$  is initialised to some query  $\xi$  the system the fixed point update is given by  $\mu_i^* (\xi) = \mathbb{E}_{\phi_i|x,\xi}[x_{[j]}]$ . If the patterns  $x$  are well separated,  $\mu_i^* (\xi)\approx x_{j'}$ , where  $x_{j'}$  is the closest vector and hence can be used as an associative memory.
4.3 Slot Attention
Slot attention [17] is an object centric learning module centred around an iterative attention mechanism. Here we show this is a simple adjustment of the prior beliefs on our edge set. With edge potentials of the form  $\psi(x_j, z_i) = z_i^T W_Q^T W_K x_j$ , replace the prior over edges with  $p(\phi) = \prod_{j=1}^{n} p(\phi_j)$ ,  $\phi_j \sim \text{Uniform}\{(x_j, z_1),.., (x_j, z_m)\}$ . Notice, in comparison to MCHN, the prior over edges is swapped, each observed node is uniformly likely to connect to a latent node, in turn altering the index of the softmax.
$$
\mu_ {i} ^ {*} = \sum_ {j} s o f t m a x _ {i} \left(\mu_ {i} ^ {T} W _ {Q} ^ {T} W _ {K} x _ {j}\right) W _ {Q} ^ {T} W _ {K} x _ {j}
$$
while the original slot attention employed an RNN to aid the basic update shown here, the important feature is that the softmax is taken over the 'slots'. This forces competition between slots to account for the observed variables, creating object centric representations.
4.4 Predictive Coding Networks
Predictive Coding Networks (PCN) have emerged as an influential theory in Computational Neuroscience [29, 30, 31]. Building on theories of perception as inference and the Bayesian brain, PCNs perform approximate Bayesian inference by minimising a variational free energy of a graphical model, where incoming sensory data are used as observations. Typical implementations use a hierarchical model with Gaussian conditionals, resulting in a local prediction error minimising scheme. The minimisation happens on two distinct time-scales, which can be seen as E-step and M-steps on the variational free energy: a (fast) inference phase encoded by neural activity corresponding to perception and a (slow) learning phase associated with synaptic plasticity. Gradient descent on the free energy gives the inference dynamics for a particular neuron  $\mu_{i}$ , [32]
$$
\frac {\partial \mathcal {F}}{\partial \mu_ {i}} = - \sum_ {\phi^ {-}} k _ {\phi} \epsilon_ {\phi} + \sum_ {\phi^ {+}} k _ {\phi} \epsilon_ {\phi} w _ {\phi}
$$
Where  $\epsilon$  are prediction errors,  $w$  represent synaptic strength,  $k$  are node specific precisions representing uncertainty in the generative model and  $\phi^{-},\phi^{+}$  represent pre-synaptic and post-synaptic terminals resectively. Applying a uniform prior over the incoming synapses results in a slightly modified dynamics,
$$
\frac {\partial \mathcal {F}}{\partial \mu_ {i}} = - \sum_ {\phi^ {-}} \operatorname {s o f t m a x} (- \epsilon_ {\phi} ^ {2}) k _ {\phi} \epsilon_ {\phi} + \sum_ {\phi^ {+}} \operatorname {s o f t m a x} (- \epsilon_ {\phi} ^ {2}) k _ {\phi} \epsilon_ {\phi} w _ {\phi}
$$
where the softmax function induces a normalisation across prediction errors received by a neuron. This dovetails with theories of attention as normalisation in Psychology and Neuroscience [33, 34, 35]. In contrast previous predictive coding based theories of attention have focused on the precision terms,  $k$ , due to their ability to up and down regulate the impact of prediction errors [36, 37, 38]. Here we
Figure 2: Multihop Attention: (left) Graphical description of the toy problem,  $x_{2}$  is generated causally from  $x_{1}$  and  $x_{0}$ , which are used to generate  $y$ . (centre) Comparison of the attention employed by Multihop which takes two steps on the attention graph (top) contrasted with Self Attention (bottom). Multihop Attention has the correct bias to learn the task approaching the performance of two-layer Self Attention, while a single layer of Self Attention is unable (top right). Empirically examining the attention weights, Multihop Attention is able to balance attention across two positions, while self-attention favours a single position.