Attention as Implicit Structural Inference
Authors: Ryan Singh
Christopher Buckley
Venue: NeurIPS 2023 poster (2023)
Keywords: Attention
Structural Inference
Variational Inference
Predictive Coding
Graphical Models

Abstract: 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.

Attention as Implicit Structural Inference
Ryan Singh School of Engineering and Informatics, University of Sussex.
Christopher L. Buckley
School Engineering and Informatics, University of Sussex. VERSES Al Research Lab, Los Angeles, CA, USA.
rs7730sussex ac.
Abstract
Attention mechanisms play crucial role in cognitive systems by allowing them flexibly allocate cognitive resources. Transformers, in particular, have become a dominant architecture in machine learning, with attention their central innova- tion. However, the underlying intuition and formalism of attention in Transformers is based ideas of keys and queries in database management systems. In this work, we pursue 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 demon- strate that attention be viewed inference over an implicitly defined of possible adjacency structures in graphical model, revealing the generality of such 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 mechanism with bias toward attending multiple tokens; (b) introducing geometric prior (with conjugate hyper-prior) over the adjacency structures pro- ducing mechanism which dynamically scales the context window depending on input. Moreover, by describing link between structural inference and precision- regulation 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. hope by providing new lens on attention architectures work can guide the development of new and improved attentional mechanisms.
Introduction
Designing neural network architectures with favourable inductive biases lies behind many recent successes Deep Learning. The Transformer, and in particular the attention mechanism has allowed language models achieve human like generation abilities previously thought impossible 2 The success of the attention mechanism as a domain agnostic architecture has prompted adoption across diverse range of tasks beyond language modelling. notably reaching state-of-the-art performance in visual reasoning and segmentation tasks 3
This depth and breadth of success indicates the attention mechanism expresses useful computational primitive. Recent work has shown interesting thcoretical links to kernel methods (5) 66] 71. Hopfield networks [8]. and Gaussian mixture models 19, 10 [T
12 [3] however formal understanding that captures the generality of this computation remains outstanding. In this paper, we show the attention mechanism
naturally be described inference the structure of a graphical model, agreeing with observations that transformers are able to flexibly choose between models based on
context [1 15). This Bayesian perspective complements previous theory [16 g [2]. adding
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

methods for reasoning about inductive biases and the functional role of attention variables. Further, understanding the core computation as inference permits 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 See 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 we extend the inference over connectivity Bayesian setting which, turn, provides theoretical grounding for iterative attention mechanisms (slot-attention and block-slor slot attention) [17]
18 19]. Modern
Continuous Hopfield Networks [8] and Predictive Coding Networks. Finally in Sec we leverage the generality of this description in order to design new mechanisms with predictable properties.
Intuitively, the attention matrix be seen as the posterior distribution over edges F. in a graph, g = (K U 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 G + dd with respect this posterior.
p(EIQ.K)
QWOWTKT vd
(Q, softmar
WyK
=E
Crucially, when 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.
2 Related Work
key benefit of the perspective outlined here is 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. [1]] on the other hand Gaussian Mixture Models (GMMs) can be seen 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 topK 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 alignment variables. this work, rather than treating attention weights as an independent learning problem, we aim to provide parsimonious implicit model that would give rise to the attention weights. Additionally showing that 'soft' attention weights arise naturally in variational inference mnaa-ficeld approximation.
from either collapsed variational inference
Relationship 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 [1]] 10 12 B. Indeed (33 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 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 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 form of Hopfield Network [8]. providing 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 arbritary value matrix, from transformer attention. Although there remains strong theoretical connection, we interpret the Hopfield Energy as an instance of variational free energy, aligning more closely with iterative attention mechanisms such slot-attention.
3 Transformer Attention
3.1 Attention as Expectation
We begin by demonstrating transformer attention can be seen calculating expectation over graph structures. Specifically, let z = (I1,.,, Fn) be observed input variables, 6 be some set of discrete latent variables representing edges in graphical model of x given by p(2; Φ), and y variable we need to predict. Our goal is to find Eyja [y] however the graph structure ø is unobserved so we calculate the marginal likelihood.
By 111l]=
Σ p( E
Importantly, the softmax function is a natural representation for the posterior,
plz, Φ) softmar(Inp(z.o)
p( ( = Es p(x, の)
in order to expose the link to transformer attention, let the model of y given the graph (z, Φ) be parameterised by function Ey n. ly] = v(x, Ø).
Eylziv] Σ sos so p(a, ()) Φ) = (e)elv((, +)]
(1)
In general, transformer attention can be seen as weighting v(x, d) by the posterior distribution p(e z) over different graph structures. We show [ is exactly the equation underlying self and cross-attention by presenting the specific generative models corresponding to them. In this description the latent variables & are identified as edges between observed variables æ (keys and queries) in pairwise Markov Random Field, parameterised by matrices WK and Wo: while the function parameterised by Wv.
Pairwise Markov Random Fields are a natural tool for modelling the dependeneies 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 set of random variables X = (X,)vev with probability distribution [p] and graph G = (V, E). The variables form a pairwise Markov Random Field (pMRF) 26 with respect to G the joint density function P(X x) p(x) factorises as follows
p(x) Z exp
1
Σ So
Σ the cEE
BEV
where Z is the partition function (xx)) and dhe &uu (Ta, Rv) are known as the node and edge potentials respectively. Bayesian structural inference also requires a structural prior p()) over the space of possible adjacency structures, o E 4, of the underlying graph.
Factorisation, without constraints this space grows exponentially in the number of nodes V1 possible graphs leading to intractable softmax caleulations). all the models we explore here implicitly assume factorised p We briefly remark that Eq respects factorisation of [p] in the following
Additionally placing zero probability mass much of the space, for example disconnected graphs.
3

---

sense; if the distribution admits a factorisation (a partition of the space of graphs 4 = II, 4,) with respect to the latent variables p(æ, Φ) II. ef((2)00)) where 2, and the value function distributes over the same partition of edges v(z, 6) = vil Φi) then each of the factors can be marginalised independently:
E (v +)] 1o
To recover cross-attention and self-attention we need specify the structural prior, potential functions and value function. (In order to ease notation, when O; is set of edges involving common node Je, such that Φ; = (3) represents a single edge, we use the notation de = [j]. suppressing the shared index.)
3.2
Cross Attention and Self Attention
We first define the model that gives rise to cross- tross-attention:
Key nodes K = (æ1, Fn) and query nodes Q = "'n
Structural prior p()) - [I p(o)), where D; - { 2/),, (In, x(}} is the set of edges involving z and o, Unif form such that each query node is uniformly likely to connect to each key node.
Edge potentials x()),x(() = /T W Wxz in effect measuring the similarity of z, and af in a projected space.
Q
Value functions v; (2 oi = [i]) = Wvz,, linear transformation applied to the node at the start of the edge dy.
Taking the expectation with respect to the posterior in each of the factors defined in Eq gives the standard cross-attention mechanism,
E pp q, = Σ softmar, W W2Ww,,WWW
the key nodes are fact the same the query nodes and the prior is instead over directed graph we recover eeff-attention (A..I).
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..., marginalisation of structural variables, 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 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 Sec(] collapsed variational inference produces an identical energy function to the one proposed by Ramsauer al. g
4.1 Collapsed Inference
We present a version collapsed variational inference 277, where the collapsed variables & are again structural, showing how this results in Bayesian attention mechanism. In contrast to the previous section, we have a set of (non-structural) latent variables z. The goal is to infer 2 given the observed variables, æ, and latent variable model p(x, 2, Φ). Collapsed inference proceeds by marginalising out the extraneous latent variables & 27]
p 2) Σ 2,
We define a gaussian recognition density q(z) N(2; Σ) and optimise the variational free energy F(A) = (x - [x)( with respect to the parameters, (H,,.), E)), of this distribution. Application of Laplace's method yields approximate derivatives of the variational

---

Kiq
Observed: X. Lateat Z
Observed: X. Lateat: Slut Alneatiun
Observed: X, Latent:
9.K
MEHN
Empiriaa Prione M
Croos Allension
Block-Skt
Figure I: Comparison of models involved in different attention mechanisms. cach case, the highlighted edges indicate &; support these edges, given the current state of the nodes, before using this inference to calculate an expectation of the value function For iterative attention mechanisms the value function can be identified as the gradient of variational free energy, in contrast, transformer attention
uniform prior over d;. Attention proceeds by calculating a posterior over
learnable function.
free energy V F V Inp(z, p)) and VyF x V2 Inp( pJ), here we focus on the first order terms Substituting in Eq V Vu Σ
p( m)
(4)
Σ VPP(E,
(5)
E(P
In order to make the link to attention, we employ the log derivaavve trick, substituting p(.) = cln
and re-express Eq in two ways:
Σ so softmazz(I V
1n (np(z,
(6)
B -- V ((,,p,,,,100))
The first form reveals the softmax which 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 ø. 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 general pMRF before again recovering (iterative) attention models in the literature by specifying particular distributions.
Free Energy of marginalised pMRF, recall the factorised pMRF, p(x, Φ) II, (fi(0))))) Again, independence properties simplify the calculation, the marginalisation can be expressed as a product of local marginals, Σg p(x, Φ) = \ II. Es, efi(s,o4) Returning to the inference setting, the nodes are partitioned into observed nodes, x, and variational parameters μ.. Hence the (approximate) collapsed variational free energy 5 can be expressed as, F(x, ()) - Es, chiam.ood C and it's derivative:
OF - = 0pj
softmaz(f) 0f
0pj
Finally, we follow [8] in using the Convex-Concan 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 (b(x,) - 2 and the edge potentials have the form (1I,, 2z) W. zy, and define f. Eeco, the Consider the following fixed point equation,
0 0p,
ΣΣ softmar(f) 6:
(8)
As the first order terms independent of the second order ones, A 7.1 for details.
5

---

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 Fz. (z.) F(r, (b)) with equality if and only μi stationary point of F. We follow Sec 3 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 r = (Z1, F.,)) and latent nodes 2 (E1,.., 2m) have the follow- ing structural prior p(o) = Ii==1 pl (())). where Φ; ~ latent node is uniformly likely to connect to a memory node. Define edge potentials ( zi) = 2 xj. Application of Eq
( (x1,=1) .., (Tn, })). meaning each
1 = Σ softmar, (1n x ),
When pi initialised some query & the system the fixed point update is given by (E) 1e [.]] If the patterns a are well separated, p (E) se Zy,, where Iy is the closest vector and hence can used as an associative memory.
4.3 Slot Attention
Slot attention 17 is an object centric learning module centred around an iterative attention mecha- nism. Here we show this is simple adjustment of the prior beliefs
edge set. With edge poten- tials of the form ((z,, x)) z! W WK,,, replace the prior over edges with p(ò) = 1[ p(O,), Φj Uniform 21 ).... (zzj }. Notice, in comparison to MCHN, the prior over edges is swapped, each observed node uniformly likely to connect to latent node, in turn altering the index of the softmax.
14 Σ softmars W:
while the original slot attention employed an RNN to aid the basic update shown here, the important feature is that the softmax taken over the 'slots'. This forces competition between slots account for the observed variables, creating object centric representations.
4.4 Predictive Coding Networks
Predictive Coding Networks (PCN) have emerged an influential theory in Computational Neuro- science (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 graphical model, where incoming sensory data are used as observations. Typical implementations use 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: (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 14: [32
OF 0p4
Σ KEEP + Σ kgeaw
Where are prediction errors, w represent synaptic strength, & are node specific precisions repre- senting uncertainty in the generative model and & o represent pre-synaptic and post-synaptic terminals resectively. Applying a uniform prior over the incoming synapses results in a slightly modified dynamics,
0F 0p:
Σ so softmar(-6e ( 3k + Σ softmar(- -c co
where the softmax function induces a normalisation across prediction errors received by neuron. This dovetails with theories of attention as normalisation in Psychology and Neuroscience K3 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

---

Stalkic Unr/omm03
Figure 2: Multihop Attention: (left) Graphical description of toy problem, 22 generated causally from Ii and 2u, 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 single layer of Self Attention unable (top right). Empirically examining the attention weights, Multihop Attention is able balance attention across two positions, while self-attention favours a single position.
see the softmax terms play functionally equivalent role to precision variables, inheriting their ability account for bottom-up and top-down attention, while exhibiting the fast winner-taker taker-tta--attt -all dynamics that are associated with cognitive attention.
5 New Designs
By identifying the attention mechanism in terms of an implicit probabilistic model, we can review and modify the underlying modelling assumptions in principled manner to design new attention mechanisms. Recall transformer attention can be written as the marginal probability p(y x) = Es p( z))yylr.oy]. the specific mechanism is therefore informed by three pieces of data: (a) the value function p(y Φ), (b) the likelihood p(:: ()) and (c) the prior p()). Here, explore modifying (a) and (c) and show they can exhibit favourable biases on toy problems.
5.1 Multi-hop Attention
Our description makes it clear that the value function employed by transformer attention can be extended to any function over the graph. For example, consider the calculation of Eg [8%] in transformer attention, linear transformation is applied to the most likely neighbour, Zj. of Tj. A natural extension is to include a two-hop neighbourhood, additionally using the most likely neighbour I& of z.j. The attention mechanism then takes different form Ep(e, 16:))(e [V (æz. Fz, )1 (Po + V where P& is the typical attention matrix. While containing the same number of parameters single-layer of transformer attention, for some datasets two-hop attention should be able to approximate the behaviour two-layers of transformer attention.
Task Setup We simulate simple dataset that has this property using the following data generation process: Initialise a projection matrix W,, Rdx1 and relationship matrix W, E Raxd X is then generated causally, using the relationship xi+ 1 = W, N a)) to generate 2o, and 22. while the remaining nodes are sampled from the noise distribution (0, a). Finally, the target is generated from the history of 22, (æ + xo) and the nodes of X are shuffled. Importantly W, is designed to be low rank, such that performance on the task requires paying attention to both æ and øg. Figure
5.2 Expanding Attention
One major limitation of transformer attention is the reliance on a fixed context window. From one direction, small context window does not represent long range relationships, on the other hand large window does unnecessary amount of computation when modelling a short range relationship. By replacing the uniform prior with geometric distribution p(ø q) Geolg).
7

---

yostal Glviss
Figure 3: Expanding Attention: (left) Graphical description of the toy problem, J2 and y generated from which is shuffled with (exponentially decaying) recency bias. (centre) Comparison of the geometric prior, with different shades prior used by Self Attention. (right) The relative number of operations used by Expanding Attention is beneficial when either the recency bias (1/p) or the number of feature dimensions (d) is large, training curves (overlaid) across each of these settings remained roughly equivalent.
representing the iterative refinements during inference, used by Expanding and uniform
and a conjugate hyper-prion p(g) Beta(a, (a, 3) we derive mechanism that dynamically scales depending on input. We use (truncated) mean field variational inference procedure (9) to iteratively approximate p(ø xx) using the updates: 1. 9t a
9: 2. pt = p( x. qc), 3. 04++1 +- Oz + 1, 3t-1 + 8, + E<)()))) (p:);, Where a and /3 are hyperparameters determining the strength of the prior and И is the truncation horizon. Since attention dot products can be cached and reused for each calculation of step 2. the iterative procedure is computationally cheap.
The attention mechanism has asymptotic time complexity O(n²d) d) where n is the size of the size of the context window and is dimension expanding attention (nn(md + k)) where m is the size the window at convergence, and A: the number of steps to converge. If, as is typical, d is large such that d >> the time complexity of expanding attention should be favourable.
which the inner product is computed. In comparison,
Task Setup Input and target sequence are generated similarly to above (without zo). Here is moved away from T2 according to a draw from geometric distribution, Figure
6
Discussion
6.1
The Full Transformer Block
Transformer attention is typically combined with residual connections and a feedforward network. both of which have been shown important in preventing 'token collapse'. Here we briefly touch upon how these features might relate the framework presented here.
Feedforward layer, it has previously been noticed the feedforward component can also be understood key-value memory where the memories stored persistent weights 440 41) This is due to the observation ff( = W20 (War) is equivalent to attention when the non-linearity a is softmax, although ReLUU is typically used. We speculate the framework presented here could be extended explain this discrepancy, intuitively the ReL.U relates to an edge prior that fully factorises into binary variables.
Residual connections have been shown to encourage iterative inference possibility transformer attention, rather than having arbitrary transformation as presented in 3 is in fact approximately implementing the iterative inference See through a form of iterative amortised inference [3). The view that the transformer is performing iterative refinement is additionally supported by empirical studies of carly-decoding [4]]
42
This raises the
Temperaturr and positional encodings, both positional encodings and the temperature scaling can be seen as adjustments to the prior edge probability, In the case relative positional encodings, by breaking the permutation invariance of the prior (A 8 While the temperature may be understood

---

in terms of tempered (or generalised) Bayesian inference 45 adjusting the strength of the prior relative to the likelihood.
6.2 Limitations
The connection to structural inference presented here is limited to the attention computation of single transformer head, an interesting future direction would be to investigate whether multiple layers and multiple heads typically used in a transformer Additionally, the extension to iterative inference employed crude approximation to the variational free energy, arguably destroying the favourable properties of Bayesian methods. Suggesting the possibility of creating iterative attention mechanisms with alternative inference schemes, possibly producing more robust mechanisms.
also be interpreted within this framework.
6.3 Conclusion
In this paper, we presented a probabilistic description of the attention mechanism, formulating attention as structural inference within probabilistic model. This approach builds upon previous research that connects cross attention to inference in Gaussian Mixture Model. By considering the discrete inference step in Gaussian Mixture Model as inference on marginalised structural variables, we bridge the gap with alignment focused descriptions. This framework naturally extends self-attention, graph attention, and iterative mechanisms, such as Hopfield Networks. We hope this work will contribute to a more unified understanding of the functional advantages and disadvantages brought by Transformers.
Furthermore, we argue that viewing Transformers from a structural inference perspective provides different insights into their central mechanism. Typically, optimising structure is considered learning problem, changing on relatively slow timescale compared to inference. However, understanding Transformers as fast structural inference suggests that their remarkable success stems from their ability to change effective connectivity on the same timescale as inference. This general idea can potentially be applied to various architectures and systems. For instance, Transformers employ relatively simple switches in connectivity compared to the complex dynamics observed in the brain 46). Exploring inference over more intricate structural distributions, such as connectivity motifs or modules in network architecture, could offer artificial systems even more flexible control of resources.
Acknowledgements
This work was supported by The Leverhulme Trust through the be.Al Doctoral Scholarship Pro- gramme in biomimetic embodied AI. Additional thanks Alec Tschantz, Tomasso Salvatori, Miguel Aguilera and Tomasz Korbak for their invaluable feedback and discussions.
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention All You Need, December 2017.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, Dario Amodei. Language Models are Few-Shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[3] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, June 2021.
[4] Wenhui Wang. Hangbo Bao, Li Dong, Johan Bjorck. Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, and Furu Wei. Image Foreign Lan- guage: BEiT Pretraining for All Vision and Vision- Lisnuu-Le Tasks, August 2022.

---

[5] Yifan Chen, Qi Zeng. Heng Ji, and Yun Yang. Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystrl"om Method, October 2021.
[6] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernll, 2019.
November 2019.
[7] Xing Han, Tongzheng Ren, Tan Minh Nguyen, Khai Nguyen, Joydeep Ghosh, and Nhat Ho. Robustify Transformers with Robust Kemee Density Estimation, October 2022.
[8] Hubert Ramsauer, Bernhard Schäft, Johannes Lehner, Philipp Scidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield Networks is All You Need, April 2021.
[9] Javier R. Movellan and Prasad Gabbur. Probabilistic Transformers, November 2020.
[10] Prasad Gabbur, Manjot Bilkhu, and Javier Movellan, Prohabilistic Attention for Interactive Segmentation, July 2021.
[11] Xia Li, Zhisheng Zhong, Jianlong Wu, Yibo Yang, Zhouchen Lin, and Hong Liu. Seppetation-Maxxiiaaal Attention Networks Semantic Segmentation. In 2019 IEEF/CVF International Conference on Computer Vision (ICCV), pages 91 91667,,,, Seoul, Korea (South), October 2019. IEFF. ISBN 978-1-72814-803-8. doi: 10.1109/CCV22201900000
[12] Alexander Shim. Probabilistic Interpretation Transformers, April 2022.
[13] Tam Minh Nguyen, Tan Minh Nguyen. Dung D. D. Le. Duy Khuong Nguyen. Viet-Anh Tran, Richard Baraniuk, Nhat Ho, and Stanley Osher. Improving Transformers with Probabilistic Attention Keys. In Proceedings of the 39th International Conference on Machine Learning, pages (6595-16621 PMLR, June 2022.
[14] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev. Andrey Zhmoginov, and Max Vladymyrov. Transformers Ieam in-context by gradient descent, December 2022.
[15] Shivam Garg. Dimitris Tsipras, Percy Liang, Gregory Valiant. What Can Transformers Learn In- Context? A Case Study of Simple Function Classes.
[16] Yoon Kim, Carl Denton, Luong Hoang, and Alexander M. Rush. Structured Attention Networks, February 2017.
[17] Francesco Locatello, Dirk Weissenbom, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-Centre Learning with Slot Attention, October 2020.
[18] Gautam Singh, Yeongbin Kim, and Sungjin Ahn. Neural Block-Slot Representations, November 2022.
[19] Helen C. Barron, Ryszard Auksztulewicz, and Karl Friston. Prediction and memory: A predictive coding account. Progress in Neurobiology, 192:101821, September 2020. ISSN 03010082. doi: 10.101 neeuroio.2020.101822
[20] Yuntian Deng, Yoon Kim, Justin Chiu, Demi Guo, and Alexander Rush. Latent Alignment and Variational Attention. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[21] Xinjie Fan, Shujian Zhang, Bo Chen, and Mingyuan Zhou. Bayesian Attention Modules. In Advances in Neural Information Processing Systems, volume 33, pages 16362-16376. Curran Associates, Inc., 2020.
[22] Shiv Shankar, Siddhant Garg, and Sunita Sarawagi. Surprisingly Easy Bard-AAttention for Sequence to Sequence Learning. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 640-645, Brussels, Belgium, October 2018. Association for Computational Linguisties. doi: 1/D18655/11DDI8-10666
[23] Louis Annahi, Alexandre Pitri, and Mathias Quoy. On the Relationship Between Variational Inference and Auto-Associative Memory, October 2022.
[24] Jordan Frecon, Gilles Gasso, Massimiliano Pontil, and Saverio Salzo. Bregman Neural Networks. In Proceedings of the 39th International Conference on Machine Learning, pages 6779-67779799 67722, PMLR, June 2022.
10

---

[25] Yongyi Yang, Zengfeng Huang, and David Wipf. Transformers from an Optimization Perspective, May 2022.
[26] Martin J. Wainwright and Michael I. Jordan. Graphical Models, Exponential Families, and Variational Inference. Foundations and T in Machine Learning. (1-2):1-305, November 2008. ISSN 1935- 8237, 1935-8245. doi: 10.1561/2200000001.
[27] Yee Teh, David Newman, and Max Welling. A Collapsed Variational Bayesian Inference Algorithm for Latent Dirichlet Allocation. In Advances in Neural Information Processing Systems, volume 19. MIT Press, 2006.
[28] Alan L Yuille and Anand Rangarajan. The Concave-Comee Procedure (CCCP). In Advances in Neural Information Processing Systems, volume 14, MIT Press, 2001.
[29] Rajesh Rao and Dana H. Ballard. Predictive coding in the visual cortex: functional interpretation of some extra-classical receptive- effects. Nature Neuroscience, 2(1):79 87, January 1999. ISSN 1546-1726. doi: 10. 1038//8/880.
[30] Karl Friston and Stefan Kiebel. Predictive coding under the free-energy principle. Philosophical Transac- tions of the Royal Sociery B: Biological Sciences, 364(1521):1211-4 May 2009, ISSN 0962-8436, doi: 10. 1a.1099//stb./e www.stt...0000000000000
[31] Christopher I.. Buckley, Chang Sub Kim, Simon McGregor, and Anil K. Seth. The free energy principle for action and pereeption: A mathematical review. Journal of Mathematical Psychology, 81:55-79, December 2017. ISSN 0022-2496. doi: 10. 2o.1016/j.mm.2017..00
[32] Beren Millidge, Yuhang Song, Tommaso Salvatori, Thomas Lukasiewiez, and Rafal Bogacz. A Theoretical Framework for Inference and Learning in Predictive Coding Networks, August 2022.
[33] John H. Reynolds and David J. Heeger. The Normalization Model of Attention. Neuron, 6182)1668-885 Jamuary 2009. ISSN 08966273. doi: 10.10166j.neuron..22 2009.011.000000.....000 ..02.
[34] Matteo Carandini and David J. Heeger. Normalization as a canonical neural computation. Nature Reviews Neuroscience, 13(1):51-62, January 2012. ISSN 1471-0048. doi: 10.1038/mmm3136
[35] Grace W. Lindsay. Attention in Psychology. Neuroscience, and Machine Learning. Frontiers Computa- tional Neuroxcience, 14, 2020. ISSN 1662-5188.
[36] Andy Clark. The many faces of precision (Replies to commentaries on "Whatever next? Neural prediction. situated agents, and the future cognitive science"). Frontiers in Psychology, 4, 2013. ISSN 1664-1078.
[37] Harriet Feldman and Karl Friston. Attention, Uncertainty, and Free-Energy. Frontiers in Human Neuro- science, 4. 2010. ISSN 1662-5161.
[38] M. Berk Mirza, Rick A. Adams, Karl Friston, and Thomas Parr. Introducing a Bayesian model of selective attention based on active inference. Scientific Reports, 9(1):13915, September 2019, ISSN 2045 -322. doi: 10. 10.1038/541598001 000--00995511100 50138-8.
[39] O. Zobay. Mean field inference for the Dirichlet process mixture model. Electronic Journal of Statistics, 3 (nonc):507-555 -55,7 January 2009, ISSN 1935-7524, 1935-7524. doi: 10.1214/08 EJS339.
[40] Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Augmenting Self-attention with Persistent Memory, July 2019.
[41] Mor Geva, Roci Schuster, Jonathan Berant, and Omer Levy. Transformer Feed-Forward Layers Are Key-Value Memories, September 2021.
[42] Stanistaw Jastrzębski, Devansh Arpit. Nicolas Ballas, Vikas Verma, Tong Che, and Yoshua Bengio. Residual Connections Encourage Iterative Inference, March 2018.
[43] Joc Marino, Yisong Yue, and Stephan Mandt. Iterative Amortized Inference. In Proceedings of the 35th PMLR, July 2018
International Conference on Machine Learning, pages
[44] Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting Latent Predictions from Transformers with the Tuned Lens, March 2023.
[45] Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. Generalized Variational Inference: Three arguments for deriving new Posteriors, December 2019.
[46] Emmanuelle Tognoli and J. A. Scott Kelso. The Metastable Brain. Newron, 81(1):35 48, January 2014. ISSN 0896-6273. doi: new.on..nerron.211..222..

---

Table 1: Different attention modules
Name
Graph (G)
Prior (p(ø))
Potentials (0)
Value v(x,
Cross Attention
Key nodes K, query nodes Q
Uniform
x/ WQ K
Væy
Self Attention
K = Q. directed edges
Uniform
x w
Væj
Graph Attention,
K = Q directed edges
Uniform (restricted)
2. WWWWKRY
Vxj
Sparse Attention
Relative tional Encodings
Posi-
K = Q, directed edges
Categorical
xf < Væj
Absolute Posi-
==Q Q
Uniform
8. W."WWwKR Væj = Ii +
tional Encodings
Classification
NN output fo(X), classes y
Uniform
fo(X) yo y5
Layer
MCHN
Observed nodes X, latent nodes Z
Uniform (observed)
2/ zTWTWxa3
oF
8z
Slot Attention
Observed nodes la- tent nodes Z
Uniform (latent)
x.W2Wka,
OF 82
Block-Slot Atten-
Observed nodes X, latent nodes Z, mem- ory nodes M
Uniform (latent)
of SWWWka) m m 02
aF
tion
PCN
Observed nodes X, multiple layers of la- tent nodes
Uniform (latent)
#WWWKIY
OF 0=
(2(0) Jigz
Multihop Atten-
K = Q, directed edges
Uniform
x <W'''
Væj + Vor
tion
Expanding Atten-
= Q. directed edges
Geometric x Beta
x. =."TTKaz Vx,
tion
7 Appendix
Here include some more detailed derivations of claims made in the paper, and list the hyperparam- eters used for the experiments.
7.1 Iterative Attention
In this section we provide a more detailed treatment of the Laplace approximation, and provide proper justification for invoking the CCCP. For both, the following lemma useful:
Lemma 7.1. The function pn)( = In E p(x, Φ) = In Σ exp Fe((z) has derivatives (i)
de In p(x) = x 0. Feo) and (ii) 8² In p(x) = Var o dz dz
p Fegl + Eay dr F&)
Proof. Let E = (Es) the vector of possible energies, and p = (p²) = (p(Φ x))s the vector of
conditional probabilities. Consider in p(ó x²) written in canonical form,
In x) = (z). 1s) - (xE h(o)
Where A[E(x)] = In (EE) is the cumulant generating function. By well known properties of the cumulant: 04 = p( = 1 - = pt. Hence by the chain rule for partial derivatives, 24
dz Σ p(
DE,
x)) Es, which is (i).
To find the second derivative apply again the chain-rule d ( (") A f'((())) (t). Again by properties of the cumulant aE,8E, = Cov(1, 1,) (diag(p) - -p-
12

---

Vi.- Hence the second derivative is
A 0x
0ET
DE
8 E
+ E
(9)
dr
dz
dx
Second order Laplace Approximation With these derivatives in hand calculate the second order laplace approximation of the free energy F = (x - In p(r, =)]-
0
F E,,E,lpuuu + x)) (= (4) (a- - pu)²
In p(u x) - 14] + [g]
Oz
222
8
1 Llog En +C 2
2 ~lup(u, + +r V Varola
Bol) + tr B Be 0z
Bxll + 022
We can see optimising the first order variational parameter this approximation is independent of Eg. a²
hence first find lr and the fill in our uncertainty Σg 8=2 In p(u² x²) Var ofar 8z E&] Eølp.= 1822 Es]. Finding this uncertainty can be costly in the general case where the hessian of E is not analytically available.
As alluded to in the paper, iterative attention mechanisms can also be viewed as an alternating maximisation procedure, which may provide route to more general inference schemes:
As Alternating Minimisation Collapsed Inference can also be seen as co-ordinate wise variational inference 27 Consider the family of distributions Q {u(2; A)a( =)), where ((2; X) is parameterised, however (()) is unconstrained.
F min E,g In g(2, Φ) In Imp(x, 2,
QEQ
min E E E ( In (Φ)) - In p(a, Φ 2)] + In ((2) - In p(=)]
.EQ
The inner expectation is maximised for (σ) p( xx, x2 2) and the inner expectation evaluates to In p( 2) which recovers the marginalised objective
min E (x)((((2) In 2(1 vEQ
This motivates an alternate derivation of iterative attention as structural inference which is less reliant on the Laplace approximation; Consider optimising over the variational family {9(2; A)g(o)) coordinate wise:
In 1t41 (0) E (= p( x, 2, C
A+11 arg min E (2())) In ((2) Inp(2, 1
= arg min E (0) [Fo
In the case of quadratic potentials, (+)) p( Ae), hence the combined update step can be written
arg arg min En (o [F (X)]
Each step necessarily reduces the free energy of the mean-field approximation, so this process converges. This derivation is independent of which approximation or estimation used to minimise the typical variational free energy.
7.1.1 Convexity details for the CCCP
Given a pairwise pMRF with quadratic potentials ((x;)) = - 1.? and the edge potentials have the form (z),,r,) = x; z,Wa, zj and W p.s.d., s.t. (nz = = Eveg x, In Es exp exp (x)), where
(x(x) EcEo vie. We need the following lemma to apply the CCCP:
13

---

Lemma 7.2. In Es exp) is convex in x.
Proof. We reapply Lemmal with Eg = go(2)) hence d7z En exp exp = Var gol [ o² Eøja dg gol. The first matrix a variance, so p.s.d. The second term øf |5 azd the is a convex sum of p.s.d matrices. Hence both terms are p.s.d, implying In Ee exp Io (x) is indeed convex. L
7.2 PCN Detailed Derivation
Here we through the derivations for the equations presented in section 4.4 PCNs typically assume hierarchical model with gaussian residuals:
N(jo, Eo)
2++1 2 N(fi(z;; Ex)
2N N(fN(ZN;BN) ΣN)
Under these conditions, delta approximation of the variational free energy is given by:
Fl [B(zm) 1- Inp 2)] + [[g]
F(u, 0)
Σ 7
Where 6 = (M4+1 - f((p; 0))) The inference phase involves adjusting the parameters, μ in the direction of the gradient of F, which for given layer is:
OF 0μp = - EJ (f)
(10)
Here, for of comparison, consider the case where the link functions linear, f((-) W.(-) and further the precision matrices are diagonal M, diag(k)). Under these conditions we can write the derivative component-win wiss
wise as sums of errors over incoming and outgoing edges
OF 0pm ) =
Σ KEEp + Σ KEES
@+
Where o o represent the set of incoming and outgoing edges respectively, and redefine Eo = 44 for an edge Φ = (2z, 24) and kg (2₂)) the precision associated with the node at the terminus of 6.
Now we instead assume uniform prior over incoming edges, or concretely;
2o N Σo)
6 Uniform(((2()))) (())
241 2, o N(wj 21 of
1/k))
y 22 N(fN(zz: 0x), En)
The system becomes pMRF with edge potentials given by the prediction errors, recall applying
Eq 4
ar ΣΣ softmar([[( n, (:)) 0pd
05:
0μz
Here for a node in a given layer, participates one 4/_, and all the t from the layer above, where every f((xz μt db)) here is squared prediction error corresponding to the given edge .j
14

---

(x)
wof 21-1)² hence:
OF
Σ
so voftmar,((- ((6
0pj
iE4
+
Σ
Σ
softmar;(-(eff)))) 1 =1)
=)
kE|l] PE
0F
Σ
softmar,(-(6f_1)99
0p-
+ Σ
woftmary(--ee KEW
Here incoming signals (nodes i) compete through the softmax, whilst the outgoing signal competes with other outgoing signals from nodes (nodes i') in the same layer for representation in the next layer (nodes k), see block-slot attention diagram for intuition. By abuse of notation (reindexing edges as ()
OF
Σ softmar( () + Σ Softmarr-- co " Mkpcoue ot
0p:
While we derived these equations for individual units to draw an easy comparison to standard Predictive Coding, we note likely more useful to consider blocks of units competing with each other for representation, similar to multidimensional token representations in typical attention mechanisms. We also briefly note here, the Hammersley-Cliffonn theorem indicates a deeper duality between attention as mediated by precision matrices and as structural inference.
7.3 New Designs
Multihop Derivation oll in transformer attention, linear transformation is applied to the most likely neighbour, Tj. of z;. A natural extension is to include a two-hop neighbourhood, additionally using the most likely neighbour Iı of zz. Formally, the value function u longer neatly distributes over the partition t; however the attention mechanism then takes a different form:
Eptc, Epto,le.)(ee (p(.) (x) (Io, + =o;)] (Pa + VX.. Where we use dz(0) = oy denote the edge set of the node at the end of de. To this note:
Ep(0lz) IV (Iro + xo,)]
ΣI + b
plón V + Io₂)
ΣI
p( IVV + xx0,)
&
EI p(óx VIz + ElIp(e V\) à by independence properties
@ k
p( V)V + Σ p(di x)p(0,
p..d,
Denoting the typical attention matrix, P. where Pij p(p: = bi] x))
ΣΣ
PiKPit VZx + Σ Pij PiyVa Vzj
K
= -P& PP)+P
Expanding Derivation As in the main text, let p( 9) ~ Geolg) and p(4) ~ Beta(a, B). such that have the full model p(æ, o, q; a. B) p(x d)p(ó q)p(g: a, B). In order to find p( x) we employ a truncated Mean Field Variational Bayes [39,, assuming a factorisation p)(0.g) = (0)) and using the updates:
In Pt+ (Ø) = Epttg) p( ()) - Imp(6 ( 4) +
(p++(()) = Em p( g) inp(g: ( +C2
15

---

By conjugacy the second equation simplifies to simple update of the beta distribution
Pt+1(g) (g) Beta(a Be+1)
at+1 = at +
Bu+1 3e + (p
While the second update can be seen as calculating the posterior given q = (o)0ll.
In P (6)
pn () + (e) (n p(ő ()] + C2
Inp( Φ) + pBp [in + C2
p( x., qc)
Finally, use truncation to approximate the infinite sum B En pilo k-)k E< pi k)k. Where we set the horizon according to the current distribution of q. For example in our experiments we chose H(q)) = In 0.05/ In(1 (:) the truncation that would capture 95% of the probability mass of the prior.
8 Attention Variants
Here we briefly discuss some variants of attention that there wasn't space for the paper.
8.1 Self Attention
Nodes = Q = (I1, H.))
Structural prior, over fully connected, directed graph p( - where क,
{(01, (zn, is the set of edges involving I; and o; ~ Uniform( i)), such that each node is uniformly likely to connect to every other node given direction.
Edge potentials ( 24) = x! WT in a projected space.
W WKTJ. in effect measuring the similarity of z, and a;
Value functions (z Φ: = (i) = Wvzz, a linear transformation applied to the node at the start of the edge di.
Again, taking the posterior expectation each of the factors defined in two Eq 2 gives the standard self-attention mechanism
Epto.iQ1 [vx] Σ so softmar, (1 WoWka,)
8.2 Positional Encodings and Graph Neural Networks
In Table[] we show that positional encodings and graph attention are naturally incorporated in this framework. Absolute positional encoding as suggested by Vaswani al. [ can be seen as modifying the edge potentials with a vector that depends on position, while relative position encodings can be seen as categorical prior, where the prior depends on the relative distance between nodes. Graph and Sparse attention operate similarly to graph attention, except the uniform prior is restricted to edges in the provided graph, or according to predefined sparsity pattern.
Relative Position Encodings If the prior over edges is categorical i.e. P(e: bi = Pi.j. it can be
fully specified by the matrix (P)t.) = PP.j- This leads the modified attention update
Σ softmaz,(x.0" KI, + ImpT))
However this requires local parameters for each node z; A more natural prior assign different probability to the relative distance of from j. This is achieved with P circlpt. P2.... Pn). where cire the circulant matrix of (p)is.. Due to properties of circulant matrices In P.y can be reparameterised with the hartley transform
Inpic
Σ Be(cos(kOL.,) + sin(k0.,)) = ((65)
Where ((())) cos(k ا-6 ) + sin(k ا-2 2nn
can be thought of a relative position encoding, and 3 are
2nm parameters to be learnt.
16

---

8.2.1
Block-Slot Attention
Singh et al. [18] suggest combining associative memory ability with an object-centric slot-like ability and provide an iterative scheme for doing so, alternating between slot-attention and hopfield updates. Our framework permits flexibly combine different attention mechanisms through different latent graph structures, allowing us to derive a version of block-slot attention.
In this setting we have three sets of variables X, the observations, Z the latent variables to be inferred and M which are parameters. Define the pairwise MRF fzi...., 2m and M = mm]] with a prior over edges p(E) = IT p(EE) p (Ex). E, Uni Uniformi f 21), (I,, zm)}, Eg Uni (21, my), ... (Zm: mz)}, with edge poten- tials between X and Z given by (x),, 2² Kx, and between Z and M, (2,, my) 2s mi applying (8) gives
(I1, In.). Z
Σ so fttmax((u,0) Kxz)Q" Kr,
Σ softmarx (p. mamm
In the original block-slot attention each slot 2; is broken into blocks, where each block can access block-specific memories i.e. z!") can has possible connections to memory nodes {m (b)
J-sz- Allowing objects to be represented by slots which in turn disentangle features of each object in different blocks. We presented single block version above, however is easy to that the update extends the multiple block version applying (8) gives
14 Σ so Lonax Kx²)Q" Kx,
+ Σ
(0) softmarkle (p: my )
k,b
9
Experimental Details
Multihop
Following the notation in the text; data generation parameters:
Total number of tokens: 10
Embedding dimension (dimension each w):10
Output dimension (dimension of y):
(autoregressive noise): 1
Random matrix initialisation was performed with torch.rand
Training parameters (across all models):
batch size:200
number of batches:
optimiser: ADAM
learning rate: le
Number of different random seeds: 10
Model: To make analysis easier, all models were prevented from self-attending to the final token.
Expanding
Following the notation in the text; data generation parameters:
Total number of tokens: 50
17

---

Embedding dimension(s) (dimension of each r):[10, 50]
the parameter for generating geometric shuffle:[0.5, 22 .4]
Output dimension (dimension of y/:
(autoregressive noise):
Random matrix initialisation performed with torch.rand
Training parameters (across all models):
batch size:1*
number of batches: 10000
optimiser: ADAM
learning rate: 5e -
Model: To make analysis easier, all models were prevented from self-attending to the final token. For expanding attention the hyperparameters were set as = .1, 3 = .9 these were chosen to have mean value at roughly quarter the (size 50) window.
*Training was performed with single samples, despite the iterative process being completely parallel (no shared state). Naive parallel implementation of expanding attention would encounter synchronisa- tion locks, as the fastest samples wait for the longest ones to complete. In order to take full advantage of dynamic window over batch, intelligent asynchronous processing would be necessary.
Algorithm 1 Attention
Require: X, Wq WR. W.
Q WgX
K t- WEX
V < W.X
A + softmar( (TQ)
Y < AV
Algorithm 2 Multihop
Require: X, X,W, q. WR. W,, N
Q < WgX
K + WEX
V <- W,X
A softmar(K" Q)
P + EKEN A
Y PV
Algorithm 3 EEppanding
Require: X, Wg. Ws, W., a, 8 = 0
q <- Wgng
K < WEX
V W.X
while & not converged do
p = a/(a+ + 3)
A + In(.05)/ In p)
9pl 1 = -i/k
so softmar(q K[: k. + gp)
- 11
3- +i Alili
end while 3
AV AV
18


Reviews:
[{"review_id": "hVagVqo9qJ", "reviewer": "Reviewer_tiWP", "summary": "This submission describes how many different transformer architecture variants can be see as implicit structural inference.  The inference is understood as taking an expectation over possible connectivity structures constrained by a prior over structures.  Several variants of attention are shown to be describable in this framework (with more in the appendix).  Two new 'designs' of attentional systems are described and illustrated with two briefly-described experiments.", "questions": "A more selective presentation of the core ideas (e.g. focusing on 2-3 examples rather than 5 with the others relegated to the appendix, coupled with a more extensive set of investigations addressing the above points would have resulted in a more useful and impactful contribution.\n\nI thought the paper would have benefited by showing how a graph was actually latently inferred by a learning experiment -- the title suggested that such an idea might be forthcoming, but it didn't seem to be.  If the attention allocation illustrated in the lower right of Fig 2 was intended to demonstrate what the authors meant by this, they should have made this more explicit.  But if this is what they meant, perhaps the finding is disappointing.  Indeed, the failure of one layer of one-hop attention to learn the problem in 5.1 suggests that any interesting graph that is learned by a transformer must depend on multiple layers and/or heads.\n\nThe authors have a whole page in the main text in which they could have developed their presentation more fully.  A revision might use this space to address some of these questions or expand the presentation and analysis of the experiments.\n\n", "limitations": "The limitations section of the paper does hint at some of the weaknesses mentioned and points to possible future directions.  There are no ethical concerns with this research.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The idea of unifying related architectures under a common Bayesian inference framework seems useful and the authors have thoughtfully engaged with a number of related architectures to characterize how they can be integrated under a common framework.\n\nThe idea of allowing context to select a graph over which to perform inference is important and interesting and I would like to see it developed further.\n\nThe experiments introduce ways of thinking about how transformer style architectures could be extended in novel ways.\n\n", "weaknesses": "The paper is largely limited to describing various existing models within the attention as expectation over structures framework.  The new experiments are described extremely briefly and the neural network architecture used was very hard to discern, even after a close reading of the appendix.\n\nThe idea of allowing context to select a graph over which to perform inference was not developed.  I thought at first that the expanding attention experiment was going to address this, but, if I understand correctly, there are several different experiments each with a fixed p, so that there is no run-time context sensitivity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. \n\n- *The paper is largely limited to describing various existing models within the attention as expectation over structures framework.*\n\nWe see our key contribution as a unifying theoretical framework helping to understand the fundamental computation underlying attention, which is why we thought it appropriate to dedicate so much space to recovering models. \n\n-  *The new experiments are described extremely briefly and the neural network architecture used was very hard to discern, even after a close reading of the appendix.*\n\nWe apologise for this, given a revision we could make the architectures clear through inclusion of pseudocode. \n\n- *The idea of allowing context to select a graph over which to perform inference was not developed. I thought at first that the expanding attention experiment was going to address this, but, if I understand correctly, there are several different experiments each with a fixed p, so that there is no run-time context sensitivity.*\n\nIndeed we see the key property unifying the different attention mechanisms is their use of context to determine a graph over which to perform inference at runtime.  We believe the theory we developed here describes exactly this.\n\nWhile the experiments here were not designed specifically to demonstrate this property (rather to use our understanding to design extensions to the attention mechanism) they also exhibit inference time context sensitivity.\n\nSpecifically, in the expanding task with different p; the p value is a parameter for the data generating (task) distribution*, not the model —- the model determined the context window based on an individual instance of the task, growing the window as needed (on a per data-point basis).\n\n*The parameter is of a geometric distribution, i.i.d. draws from which determine the distance of the signal token from final token. This distance is therefore still different per datapoint. \n\n- *Indeed, the failure of one layer of one-hop attention to learn the problem in 5.1 suggests that any interesting graph that is learned by a transformer must depend on multiple layers and/or heads*\n\nWe agree multiple heads or layers are crucial for learning complex functions, however we believe interesting graphs can be inferred by a single head.\n\nThe reason this is not evident in 5.1, the one-hop attention used a restricted internal dimension (keys, queries)  (supplementary L.126)  where typically the embedding space is large enough to incorporate multiple types of relationship within a single head leading to more interesting graphs."}, {"review_id": "4SRuSBVbjr", "reviewer": "Reviewer_psYk", "summary": "This paper proposes a probabilistic interpretation of attention mechanism, where the computation of attention can be expressed as the expectation of a value function defined on the nodes of a graph consisting of the query nodes and key nodes, and the expectation is taken with respect to the posterior distribution over the edges of the graph. Under such interpretation, different \"soft\" attention mechanisms fall under the same probabilistic framework, with different node configurations and prior distribution over the edges. Building upon existing works of linking attention mechanisms with variational learning with Gaussian mixture model, the authors additionally established the link between variational inference over the edges (graph structure) and slot attention and hopfield networks. Based on the probabilistic insights, the authors propose new modelling assumptions for designing new attention mechanism.", "questions": "- d is not defined on line 51;\n- Figure 2 and 3 can be made larger for readability;", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed probabilistic interpretation for attention mechanisms is sounding and widely applicable to various attention types;\n- The new perspective on the link between attention and predictive coding network based on the probabilistic framework is interesting;\n- The evaluation of the proposed new attention mechanisms on toy problems exhibit promising results;", "weaknesses": "- Despite the mathematical derivations are sounding, the paper can be hard to read since the connection based on the graph formulation is implicit, arguably this can be omitted due to page limits, but the authors should consider improving on this respect;", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review,\n\nsince both your questions were raised by more than one reviewer we have included the answer in the global response."}, {"review_id": "TYDFfJn9Sm", "reviewer": "Reviewer_khyL", "summary": "The paper proposes a framework for interpreting standard formulations of attention mechanisms through the lens of graphical models. The authors illustrate that their formulation unifies architectures and offers a way to easily generalize and improve the existing formulations.", "questions": "Refer to Weaknesses", "limitations": "\nThe authors do include a section on limitations, but I would say it isn't complete. I encourage the authors to refer to the questions & address them in the list of limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper offers an interesting direction of looking at existing formulations from the lens of probabilistic models, which should open up several exciting works in the future.\n- Interpreting the attention mechanism in terms of an implicit probabilistic model opens up the possibility of understanding & changing the underlying modeling assumptions in a principled manner to design new attention mechanisms.\n- The paper is well-written and easy to follow.\n- In addition, the authors also include quantitative analysis to demonstrate that their framework can assimilate the table structures to generate robust representations.\n", "weaknesses": "One of the main issues in the paper lies in the experimental evaluation & baselines. \n\nI encourage the authors to include a more rigorous evaluation of their approach, comparing the proposed framework with the formulations of stochastic attention in the literature. The authors chose to demonstrate the results on toy problems, and it is hard to evaluate the effectiveness of the approach based on those results.\n\nWithout appropriate experiments, it is hard to judge the utility of the proposed approach; though intuitively, it could do well, I anticipate issues in practice.\n\nHow scalable are the proposed changes to the attention mechanisms? How do these changes interplay with popular architectures?\n\nDoes the formulation impose any constraints on decoding algorithms? For instance, does it affect the scalability & run time?\n\nIn addition, a few relevant references need to be included. Please refer to the list below:\nShujian Zhang, Xinjie Fan, Bo Chen, Mingyuan Zhou: Bayesian Attention Belief Networks. ICML 2021\nHareesh Bahuleyan, Lili Mou, Olga Vechtomova, Pascal Poupart: Variational Attention for Sequence-to-Sequence Models. CoRR abs/1712.08207 (2017)\nShiv Shankar, Sunita Sarawagi Posterior Attention Models for Sequence to Sequence Learning  In ICLR, 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- *How scalable are the proposed changes to the attention mechanisms? How do these changes interplay with popular architectures?*\n\nWhile we see no reason, in principle, for scaling issues, since both modifications were designed with computational complexity in mind. (Multihop requires a single extra matrix multiplication which is cheap compared two two-layers, while expanding is specifically designed to minimise the amount of computation due to context window length.) It remains to be seen if the distributional properties exploited in the toy examples are present in natural data.\n\nMore broadly,  we see our key contribution as a theoretical framework which can be used to develop or reason about architectures, while our small experiments serve as proof of principle that this approach works. We would be happy to make this more explicit in the limitations section of the paper.\n\n- *Does the formulation impose any constraints on decoding algorithms? For instance, does it affect the scalability & run time?*\n\nWe are not quite sure what is meant here by decoding algorithms, but would be happy to discuss further. \n\n- *..comparing the proposed framework with the formulations of stochastic attention in the literature*\n- *In addition, a few relevant references need to be included*\n\nThank you, we will certainly include these references in the paper. Although we appear to have missed these, we would like to point out we did highlight the connection to stochastic attention (alignment) formulations (L.59-72) that we feel does capture the main approaches of prior work. Indeed for *Bayesian Attention Belief Networks* and *Posterior Attention Models for Sequence to Sequence Learning* we cited earlier contributions from (a subset of) the authors containing the key ideas."}, {"review_id": "FiRl7TqtQJ", "reviewer": "Reviewer_JZQo", "summary": "This work presents a theoretical framework of how the attention mechanism often used in transformers can be recast as inference over possible adjacency structures in graphical models. In particular, there is an implicit inference on the distribution of edges within a graphical model defined over nodes in the query and key nodes. They can explain many kinds of recent models in this framework such as slot attention, modern hopfield nets, etc. Then, they do two small experiments illustrating the use in using the framing. I am not familiar enough with the literature to know how novel this framing is. ", "questions": "* The analysis to Predictive Coding Networks is interesting. It would be nice to know how neuroscience models of hippocampal structural inference can be explained under this framework (Clone Structured Cognitive Graphs, Tolman Eichenbaum Machine, etc). ", "limitations": "Limitations were addressed in section 6.1. ", "rating": 5, "confidence": 1, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* The methods seem very theoretically rigorous. \n\n* Lots of connections to other models in the literature. ", "weaknesses": "* For the first toy problem, there could be more description as to why having a two-hop neighborhood would be advantageous, or what kind of data would have this property. The second toy problem's motivation was much more clear to me. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review,\n\n- *For the first toy problem, there could be more description as to why having a two-hop neighborhood would be advantageous, or what kind of data would have this property. The second toy problem's motivation was much more clear to me.*\n\nSince this concern was brought up by more than one reviewer we have addressed it in the global comment.\n\n- *The analysis to Predictive Coding Networks is interesting. It would be nice to know how neuroscience models of hippocampal structural inference can be explained under this framework (Clone Structured Cognitive Graphs, Tolman Eichenbaum Machine, etc).*\n\nThank you, we agree that the relationship to hippocampal function and indeed the transformer-TEM is an interesting direction that we are excited about pursuing. For now, it remains unclear exactly what the relationship is, other than it may be important for the hippocampus to perform discrete (relational) inferences in-context (cf. structural inference) while more abstract general properties are consolidated in cortex (i.e. the complementary learning systems hypothesis)."}, {"review_id": "qjm38E4rT8", "reviewer": "Reviewer_14Rb", "summary": "The paper shows how attention mechanisms can be interpreted as expectation values over learnable graph connectivity structures given a structural prior; that is from a perspective of (structural) variational inference. The authors first demonstrate this link for cross- and self-attention heads, then proceed with interpreting iterative attention mechanisms as performing gradient descent on an approximate variational free energy. Finally, based on this perspective of variational structural inference, the authors propose two new attention mechanisms which allow for more complicated connectivity structures. \n\nPlease note that due to receiving 6 papers to review from NeurIPS alone, I have allocated a time budget of 4h per paper, and my review is based on that. In particular I did not read the supplementary material in greater detail. I regret this situation and apologize for possible inaccuracies.", "questions": "- The equation showing the attention definition in the introduction is not directly referenced or explained. Also $d_k$ is undefined.\n- Since cross attention and self attention only differ by setting $x'=x$ for self attention, sections 3.2 and 3.3 could easily be merged, resulting in additional space being available for larger Figures 2 and 3 (which are essentially unreadable on A4), or for additional explanations in the beginning to not loose non-expert readers: For example clarifying why $p(\\phi|x)=softmax(\\ln p(x,\\phi))$ in line 97/98, or making the paragraph on pMRFs in lines 109-115 more accessible.\n- Would introducing temperature scaling in the softmax lead to a useful generalization?\n- In the expression for $F(x,\\mu)$ in line 177, the normalization $Z$ is missing (although of course it drops out in $\\partial_{\\mu}F$). \n- How eq. (7) follows from $\\partial_{\\mu}F = 0$ is not immediately clear.\n- In sect. 4.2, it was not clear to me why $z_i$ is can be replaced by $\\mu_i$ in the eq. following line 192. Also, is there an asterisk missing at the $\\mu_i$ on the rhs here and in eq. (7)? \n- In sect. 4.4, how does the uniform prior over incoming synapses relate to the standard version of predictive coding?\n- Is the task described in sect. 5.1 lines 243-249 motivated as a toy version of a real world task?\n- In sect. 5.2, the description currently does not reference the supplementary material.\n\n\nSome typos: \\\n l109: comma missing after partition function \\\n l159: Laplace \\\n l193: the system the fixed point \\\n l261: the size of the size of \\\n l262: \"the\" dimension \\\n l292: . instead of ,", "limitations": "A short but reasonable discussion of limitations is provided.\n", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "- A framework to interpret attention mechanisms in Transformers.\n- Principled design of new types of attention is possible by reasoning in the proposed framework.\n- Original contribution to literature linking iterative attention to variational inference. \n \n", "weaknesses": "- The proposed multi-hop and expanding attention designs are only evaluated on toy tasks with structure matching the design.  \n- Significant discussion and additional results are hidden in the supplementary material, and the presentation in the main text is dense at times. It might seem that the paper would be more coherent and better to understand in a journal/venue with larger page limit.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review, \n\n- *Since cross attention and self attention only differ by setting $x=x'$ for self attention, sections 3.2 and 3.3 could easily be merged*\n\nIn hindsight we agree with this, and would reduce self-attention and cross-attention to a single section.\n\n- *Would introducing temperature scaling in the softmax lead to a useful generalization?*\n\nWithin the framework there are two ways of accommodating temperature, either as a parameter of the graphical model (in the edge potentials) or in terms of  *tempered/safe* Bayes. Indeed including the temperature as a latent variable represents an interesting direction, however for space reasons we felt it would confuse the presentation here.\n\n- *How eq. (7) follows from  is not immediately clear.*\n\nThis follows from application of the Convex-Concave Procedure (instead of gradient descent), giving a fixed point equation which necessarily reduces the objective function. We would be happy to make this clearer in a revision. \n\n- *In sect. 4.2, it was not clear to me why $z_i$ is can be replaced by $\\mu_i$ in the eq. following line 192.*\n\nThis is due to the Laplace approximation where we are updating the variational parameters rather than the unknown variable z, however we could discuss this more explicitly for increased clarity.\n\n- *Also, is there an asterisk missing at the $\\mu_i$ on the rhs here and in eq. (7)?*\n\nWe use the asterisk on the L.H.S to indicate an approximate stationary point (solution), whereas the right hand side is determined by initialization and not necessarily a fixed point. We now realise this is not mentioned, so thanks for pointing it out, again we could make this clearer in the final presentation.\n\n- *how does the uniform prior over incoming synapses relate to the standard version of predictive coding?*\n\nIn standard formulations of predictive coding the equivalent term is governed by a precision matrix (of the generative model), learnt slowly across the data, however in practice this term is fixed as a scaled identity matrix and so acts as a uniform weighting over incoming signals (cf. uniform prior). \n\nAdditionally, even when they are learnt these matrices are treated as parameters, rather than latent variables, and so remain fixed during inference (no in-context updates) arguably making the structural inference perspective a more powerful model of attention.\n\n- *Is the task described in sect. 5.1 lines 243-249 motivated as a toy version of a real world task?*\n\nSince this concern was brought up more than once we have included the response in the global comment. "}]

Author rebuttal:
Thank you to the reviewers for their insightful comments. A couple of concerns were repeated across reviewers which we will address here.

We appreciate reviewers concern that the experiments were on toy data, however, we would like to stress that we view the main contribution of the paper as theoretical, providing a perspective on why attention mechanisms are so useful, and unifying different uses of the attention mechanism. While the experiments serve as a proof of principle that we can use this understanding to design new mechanisms, hopefully enabling future researchers to develop new architectures which scale to natural data.

- *What is a natural motivation for the task set-up in multihop attention?*

Requiring two previous states in order to generate the next one serves as a natural example. Consider a character-level sequence model for English language; attending to the letters “ee”, since they are almost always followed by a consonant,  greatly reduces the uncertainty of the next character compared to simply attending to a single “e”. Of course, such cases can be handled with higher capacity heads, multiple heads or multiple layers. However it is possible that using some multi-hop heads could aid learning efficiency if patterns such as this are common in the dataset.

More generally, if the sequence has useful information spread across two tokens, that is not exclusively available from either of them individually, multihop could serve as an alternative to increasing the number of parameters (e.g. capacity, layers or heads).

- *Figure size too small*

We will make sure we increase the figure sizes in future versions.

-  *$d$ not defined in first equation.*

Thanks to reviewers for pointing this out, we will make sure this is defined.



Meta-review: I recommend acceptance of the paper primarily due to its novel and theoretically-grounded approach to interpreting attention mechanisms through a probabilistic framework. The probabilistic viewpoint introduced in this work can have wider applications and motivate promising avenues for future research. The reviewers also found presentation to be clear and results to be novel, even though they are based on toy problems. However, there is room for improvement in experimental validation and in providing clearer motivations for the choice of toy problems. Despite these limitations, the reviewers and I believe the work's contributions are significant enough to warrant acceptance. For the final version, I recommend that the authors revise their work based on reviewer feedback. I personally would also like to see a discussion on MLPs/nonlinearities following attention but this is certainly at the authors' discretion.
Decision: Accept (poster)