SlowGuess/NeurIPS_2023_OCR huggingface dataset:
see the softmax terms play a functionally equivalent role to precision variables, inheriting their ability to account for bottom-up and top-down attention, while exhibiting the fast winner-takes-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 a principled manner to design new attention mechanisms. Recall transformer attention can be written as the marginal probability  $p(y \mid x) = \sum_{\phi} p(\phi \mid x) \mathbb{E}_{y \mid x, \phi}[y]$ , the specific mechanism is therefore informed by three pieces of data: (a) the value function  $p(y \mid x, \phi)$ , (b) the likelihood  $p(x \mid \phi)$  and (c) the prior  $p(\phi)$ . Here, we 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  $\mathbb{E}_{y|x,\phi}[y_i]$  in transformer attention, a linear transformation is applied to the most likely neighbour,  $x_{j}$ , of  $x_{i}$ . A natural extension is to include a two-hop neighbourhood, additionally using the most likely neighbour  $x_{k}$  of  $x_{j}$ . The attention mechanism then takes a different form  $\mathbb{E}_{p(\phi_j|\phi_i)p(\phi_i|x)}[V(x_{\phi_i} + x_{\phi_j})] = (P_{\phi} + P_{\phi}^{2})VX$ , where  $P_{\phi}$  is the typical attention matrix. While containing the same number of parameters as a single-layer of transformer attention, for some datasets two-hop attention should be able to approximate the behaviour of two-layers of transformer attention.
Task Setup We simulate a simple dataset that has this property using the following data generation process: Initialise a projection matrix  $W_{y} \in \mathbb{R}^{d \times 1}$  and a relationship matrix  $W_{r} \in \mathbb{R}^{d \times d}$ .  $X$  is then generated causally, using the relationship  $x_{i+1} = W_{r}x_{i} + N(0,\sigma)$  to generate  $x_{0}, x_{1}$  and  $x_{2}$ , while the remaining nodes are sampled from the noise distribution  $N(0,\sigma)$ . Finally, the target  $y$  is generated from the history of  $x_{2}$ ,  $y = W_{y}(x_{1} + x_{0})$  and the nodes of  $X$  are shuffled. Importantly  $W_{r}$  is designed to be low rank, such that performance on the task requires paying attention to both  $x_{1}$  and  $x_{0}$ , Figure 2.
5.2 Expanding Attention
One major limitation of transformer attention is the reliance on a fixed context window. From one direction, a small context window does not represent long range relationships, on the other hand a large window does an unnecessary amount of computation when modelling a short range relationship. By replacing the uniform prior with a geometric distribution  $p(\phi \mid q) \sim Geo(q)$ ,
Figure 3: Expanding Attention: (left) Graphical description of the toy problem,  $x_{2}$  and  $y$  are generated from  $x_{1}$  which is shuffled with a (exponentially decaying) recency bias. (centre) Comparison of the geometric prior, with different shades of red representing the iterative refinements during inference, used by Expanding and uniform 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.
and a conjugate hyper-prior  $p(q) \sim Beta(\alpha, \beta)$  we derive a mechanism that dynamically scales depending on input. We use a (truncated) mean-field variational inference procedure [39] to iteratively approximate  $p(\phi \mid x)$  using the updates: 1.  $q_{t} \gets \frac{\beta_{t}}{\alpha_{t} + \beta_{t}}$ , 2.  $p_{t} = p(\phi \mid x, q_{t})$ , 3.  $\alpha_{t+1} \gets \alpha_{t} + 1$ ,  $\beta_{t+1} \gets \beta_{t} + \sum_{<H(q_{t})} i(p_{t})_{i}$ . Where  $\alpha$  and  $\beta$  are hyperparameters determining the strength of the prior and  $H$  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^{2}d)$  where  $n$  is the size of the size of the context window and  $d$  is dimension over which the inner product is computed. In comparison, expanding attention  $O(n(md + k))$  where  $m$  is the size of the window at convergence, and  $k$  is the number of steps to converge. If, as is typical,  $d$  is large such that  $d >> k$  the time complexity of expanding attention should be favourable.
Task Setup Input and target sequence are generated similarly to above (without  $x_0$ ). Here  $x_1$  is moved away from  $x_2$  according to a draw from a geometric distribution, Figure 3.
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 to the framework presented here.
Feedforward layer, it has previously been noticed the feedforward component can also be understood as a key-value memory where the memories are stored as persistent weights [40, 41]. This is due to the observation  $ff(x) = W_2 \sigma(W_1 x)$  is equivalent to attention when the non-linearity  $\sigma$  is a softmax, although a ReLU is typically used. We speculate the framework presented here could be extended explain this discrepancy, intuitively the ReLU relates to an edge prior that fully factorises into binary variables.
Residual connections have been shown to encourage iterative inference [42]. This raises the possibility transformer attention, rather than having an arbitrary transformation  $v$  as presented in Sec.3, is in fact approximately implementing the iterative inference of Sec.4 through a form of iterative amortised inference [43]. The view that the transformer is performing iterative refinement is additionally supported by empirical studies of early-decoding [44].
Temperature and positional encodings, both positional encodings and the temperature scaling can be seen as adjustments to the prior edge probability. In the case of relative positional encodings, by breaking the permutation invariance of the prior (A.8.2). 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 a single transformer head, an interesting future direction would be to investigate whether multiple layers and multiple heads typically used in a transformer can also be interpreted within this framework. Additionally, the extension to iterative inference employed a 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.
6.3 Conclusion
In this paper, we presented a probabilistic description of the attention mechanism, formulating attention as structural inference within a probabilistic model. This approach builds upon previous research that connects cross attention to inference in a Gaussian Mixture Model. By considering the discrete inference step in a Gaussian Mixture Model as inference on marginalised structural variables, we bridge the gap with alignment-focused descriptions. This framework naturally extends to 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 a learning problem, changing on a 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.
This work was supported by The Leverhulme Trust through the be.AI Doctoral Scholarship Programme in biomimetic embodied AI. Additional thanks to Alec Tschantz, Tomasso Salvatori, Miguel Aguilera and Tomasz Korbak for their invaluable feedback and discussions.
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is 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, and 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 as a Foreign Language: BEiT Pretraining for All Vision and Vision-Language Tasks, August 2022.
Table 1: Different attention modules
Name	Graph (G)	Prior (p(φ))	Potentials (ψ)	Value v(x,φ)
Cross Attention	Key nodes K, query nodes Q	Uniform	xtWTWQWkxj	Vxj
Self Attention	K = Q, directed edges	Uniform	xtWTWQWkxj	Vxj
Graph Attention, Sparse Attention	K = Q, directed edges	Uniform (restricted)	xtWTWQWkxj	Vxj
Relative Positional Encodings	K = Q, directed edges	Categorical	xtWTWQWkxj	Vxj
Absolute Positional Encodings	K = Q	Uniform	xtWTWQWkxj	Vxj
xi = xi + ei
Classification Layer	NN output fθ(X), classes y	Uniform	fθ(X)iTytj	yj
MCHN	Observed nodes X, latent nodes Z	Uniform (observed)	ztWTWQWkxj	∂F/∂z
Slot Attention	Observed nodes X latent nodes Z	Uniform (latent)	ztWTWQWkxj	∂F/∂z
Block-Slot Attention	Observed nodes X, latent nodes Z, memory nodes M	Uniform (latent)	ztWTWQWkxj, mkTWTWQZi	∂F/∂z
PCN	Observed nodes X, multiple layers of latent nodes {Z(l)}l≤L	Uniform (latent)	ztWTWQWkxj	∂F/∂z
Multihop Attention	K = Q, directed edges	Uniform	xtWTWQWkxj	Vxj + Vxk
Expanding Attention	K = Q, directed edges	Geometric x Beta	xtWTWQWkxj	Vxj
Here we include some more detailed derivations of claims made in the paper, and list the hyperparameters 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 is useful:
Lemma 7.1. The function  $\ln p(x) = \ln \sum_{\phi}p(x,\phi) = \ln \sum_{\phi}\exp E_{\phi}(x)$  has derivatives (i)  $\frac{\partial}{\partial x}\ln p(x) = \mathbb{E}_{\phi |x}\left[\frac{\partial}{\partial x} E_{\phi}\right]$  and (ii)  $\frac{\partial^2}{\partial x^2}\ln p(x) = Var_{\phi |x}\left[\frac{\partial}{\partial x} E_{\phi}\right] + \mathbb{E}_{\phi |x}\left[\frac{\partial^2}{\partial x^2} E_{\phi}\right]$
Proof. Let  $E = (E_{\phi})$  the vector of possible energies, and  $p = (p_{\phi}) = (p(\phi \mid x))_{\phi}$  the vector of conditional probabilities. Consider  $\ln p(\phi \mid x)$  written in canonical form,
$$
\ln p (\phi \mid x) = \left\langle E _ {\phi} (x), \mathbb {1} _ {\phi} \right\rangle - A [ E _ {\phi} (x) ] + h (\phi)
$$
Where  $A[E(x)] = \ln Z(E)$  is the cumulant generating function. By well known properties of the cumulant:  $\frac{\partial A}{\partial E_i} = p(\phi = i \mid x) = p_i$ . Hence by the chain rule for partial derivatives,  $\frac{\partial A}{\partial x} = \sum_{\phi} p(\phi \mid x) \frac{\partial}{\partial x} E_{\phi}$ , which is (i).
To find the second derivative we apply again the chain-rule  $\frac{d}{dt} f(g(t)) = f''(g(t)) g'(t)^2 + f'(g(t)) g''(t)$ . Again by properties of the cumulant  $\frac{\partial^2 A}{\partial E_i \partial E_j} = Cov(\mathbb{1}_i, \mathbb{1}_j) = [diag(p) - p^T p]_{i,j} =$
$\mathbb{V}_{i,j}$ . Hence the second derivative is
$$
\frac {\partial^ {2} A}{\partial^ {2} x} = \frac {\partial E}{\partial x} ^ {T} \mathbb {V} \frac {\partial E}{\partial x} + \mathbb {E} \left[ \frac {\partial^ {2} E _ {\phi}}{\partial x ^ {2}} \right] \tag {9}
$$
Second order Laplace Approximation With these derivatives in hand we can calculate the second order laplace approximation of the free energy  $\mathcal{F} = \mathbb{E}_q[\ln q_\lambda (z) - \ln p(x,z)]$ .
$$
\begin{array}{l} \mathcal {F} \approx \mathbb {E} _ {q} [ \ln p (\mu , x) + \frac {\partial}{\partial z} \ln p (\mu , x) ^ {T} (z - \mu) + (z - \mu) ^ {T} \frac {\partial^ {2}}{\partial z ^ {2}} \ln p (\mu , x) (z - \mu) ] + H [ q ] \\ \approx \ln p (\mu , x) + t r (\Sigma_ {q} ^ {- 1} V a r _ {\phi | \mu , x} [ \frac {\partial}{\partial z} E _ {\phi} ]) + t r (\Sigma_ {q} ^ {- 1} \mathbb {E} _ {\phi | \mu , x} [ \frac {\partial^ {2}}{\partial z ^ {2}} E _ {\phi} ]) + \frac {1}{2} l o g | \Sigma_ {q} | + C \\ \end{array}
$$
We can see optimising the first order variational parameter in this approximation is independent of  $\Sigma_q$ , hence we can first find  $\mu$  and the fill in our uncertainty  $\Sigma_q = \frac{\partial^2}{\partial z^2} \ln p(\mu^*, x) = Var_{\phi|\mu,x}[\frac{\partial}{\partial z} E_\phi] + \mathbb{E}_{\phi|\mu,x}[\frac{\partial^2}{\partial z^2} E_\phi]$ . 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 a 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 = \{q(z;\lambda)q(\phi \mid z)\}$ , where  $q(z;\lambda)$  is parameterised, however  $q(\phi)$  is unconstrained.
$$
\begin{array}{l} \mathcal {F} = \min  _ {q \in Q} \mathbb {E} _ {q} [ \ln q (z, \phi) - \ln p (x, z, \phi) ] \\ = \min _ {q \in Q} \mathbb {E} _ {q (z)} [ \mathbb {E} _ {q (\phi)} [ \ln q (\phi) - \ln p (x, \phi \mid z) ] + \ln q (z) - \ln p (z) ] \\ \end{array}
$$
The inner expectation is maximised for  $q(\phi) = p(\phi \mid x, z)$  and the inner expectation evaluates to  $-\ln p(x \mid z)$  which recovers the marginalised objective
$$
\min _ {q \in Q} \mathbb {E} _ {q (z)} [ q (z) - \ln \sum_ {\phi} p (x, z, \phi) ]
$$
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  $Q = \{q(z;\lambda)q(\phi)\}$  coordinate wise:
$$
\begin{array}{l} \ln q _ {t + 1} (\phi) = \mathbb {E} _ {q _ {t} (z; \lambda_ {t})} [ \ln p (\phi \mid x, z) ] + C \\ \lambda_ {t + 1} = a r g \min _ {\lambda} \mathbb {E} _ {q _ {t} (\phi)} [ \mathbb {E} _ {q (z; \lambda)} [ \ln q (z) - \ln p (x, z \mid \phi) ] \\ = \arg \min  _ {\lambda} \mathbb {E} _ {q _ {t} (\phi)} [ \mathcal {F} _ {\phi} ] \\ \end{array}
$$
In the case of quadratic potentials,  $q_{t + 1}(\phi) = p(\phi \mid x,\lambda_t)$ , hence the combined update step can be written
$$
\arg \min  _ {\lambda} \mathbb {E} _ {p (\phi | x, \lambda_ {t})} [ \mathcal {F} _ {\phi} (\lambda) ]
$$
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 is used to minimise the typical variational free energy.
7.1.1 Convexity details for the CCCP
Given a pairwise pMRF with quadratic potentials  $\psi(x_i) = -\frac{1}{2} x_i^2$  and the edge potentials have the form  $\psi(x_i, x_j) = x_i W x_j$  and W p.s.d., s.t.  $\ln p(x, \phi) = -\frac{1}{2} \sum_{v \in \mathcal{G}} x_v^2 + \ln \sum_{\phi} \exp g_\phi(x)$ , where  $g_\phi(x) = \sum_{e \in \phi} \psi_e$ . We need the following lemma to apply the CCCP:
Lemma 7.2.  $\ln \sum_{\phi}\exp g_{\phi}(x)$  is convex in  $x$
Proof. We reapply Lemma.7.1, with  $E_{\phi} = g_{\phi}(x)$ , hence  $\frac{\partial^2}{\partial x^2}\ln \sum_{\phi}\exp g_{\phi}(x) = Var_{\phi |x}[\frac{\partial}{\partial x} g_{\phi}] + \mathbb{E}_{\phi |x}[\frac{\partial^2}{\partial x^2} g_{\phi}]$ . The first matrix is a variance, so p.s.d. The second term  $\mathbb{E}_{\phi |x}[\sum_{e\in \phi}\frac{\partial^2}{\partial x^2}\psi_e]$  is a convex sum of p.s.d matrices. Hence both terms are p.s.d, implying  $\ln \sum_{\phi}\exp g_{\phi}(x)$  is indeed convex.
7.2 PCN Detailed Derivation
Here we go through the derivations for the equations presented in section 4.4. PCNs typically assume a hierarchical model with gaussian residuals:
$$
z _ {0} \sim N (\hat {\mu} _ {0}, \Sigma_ {0})
$$
$$
z _ {i + 1} \mid z _ {i} \sim N \left(f _ {i} \left(z _ {i}; \theta_ {i}\right), \Sigma_ {i}\right)
$$
$$
y \mid z _ {N} \sim N \left(f _ {N} \left(z _ {N}; \theta_ {N}\right), \Sigma_ {N}\right)
$$
Under these conditions, a delta approximation of the variational free energy is given by:
$$
\mathcal {F} [ p, q ] = \mathbb {E} _ {q (z; \mu)} [ - \ln p (y, z) ] + H [ q ]
$$
$$
\mathcal {F} (\mu , \theta) \approx \sum_ {l = 0} ^ {N} \Sigma_ {l} ^ {- 1} \epsilon_ {l} ^ {2}
$$
Where  $\epsilon_{l} = (\mu_{l + 1} - f_{l}(\mu_{l};\theta_{l}))^{2}$ . The inference phase involves adjusting the parameters,  $\mu$  in the direction of the gradient of  $\mathcal{F}$ , which for a given layer is:
$$
\frac {\partial \mathcal {F}}{\partial \mu_ {l}} = \Sigma_ {l - 1} ^ {- 1} \epsilon_ {l - 1} - \Sigma_ {l} ^ {- 1} \epsilon_ {l} f ^ {\prime} (\mu_ {l}) \tag {10}
$$
Here, for ease of comparison, we consider the case where the link functions are linear,  $f_{i}(\cdot) = W_{i}(\cdot)$  and further the precision matrices are diagonal  $\Sigma_{i}^{-1} = diag(k_{i})$ . Under these conditions we can write the derivative component-wise as sums of errors over incoming and outgoing edges:
$$
\left(\frac {\partial \mathcal {F}}{\partial \mu_ {l}}\right) _ {i} = - \sum_ {\phi^ {-}} k _ {\phi} \epsilon_ {\phi} + \sum_ {\phi^ {+}} k _ {\phi} \epsilon_ {\phi} w _ {\phi}
$$
Where  $\phi^{-},\phi^{+}$  represent the set of incoming and outgoing edges respectively, and we redefine  $\epsilon_{\phi} = (\mu_i - \mu_j w_{ij})$  for an edge  $\phi = (z_i,z_j)$  and  $k_{\phi} = K(z_j)$  the precision associated with the node at the terminus of  $\phi$ .
Now if we instead assume a uniform prior over incoming edges, or concretely;
$$
z _ {0} \sim N (\hat {\mu_ {0}}, \Sigma_ {0})
$$
$$
\phi_ {l} ^ {i} \sim U n i f o r m (\left\{\left(z _ {l + 1} ^ {i}, z _ {l} ^ {0}\right), \left(z _ {l + 1} ^ {i}, z _ {l} ^ {1}\right), \dots \right\}
$$
$$
z _ {l + 1} ^ {i} \mid z _ {l}, \phi_ {l} ^ {i} \sim N (w _ {l} ^ {i j} z _ {l} ^ {\phi_ {l} ^ {i}}, 1 / k _ {l} ^ {i})
$$
$$
y \mid z _ {N} \sim N (f _ {N} (z _ {N}; \theta_ {N}), \Sigma_ {N})
$$
The system becomes a pMRF with edge potentials given by the prediction errors, recall applying Eq.4:
$$
\frac {\partial F}{\partial \mu_ {j}} = - \sum_ {i} \sum_ {\phi_ {i}} \operatorname {s o f t m a x} \left(f _ {i} (x, \mu , \phi_ {i})\right) \frac {\partial f _ {i}}{\partial \mu_ {j}}
$$
Here for a node in a given layer, it participates in one  $\Phi_{l-1}^j$  and all the  $\Phi_{l+1}^k$  from the layer above, where every  $f_i(x,\mu,\phi_i)$  here is a squared prediction error corresponding to the given edge  $e_l^{ij} =$
$k_{l}^{ij}(z_{l}^{i} - w_{l}^{ij}z_{l - 1}^{j})^{2}$  , hence:
$$
\begin{array}{l} \frac {\partial F}{\partial \mu_ {j}} = - \sum_ {i \in \Phi_ {l - 1} ^ {j}} s o f t m a x _ {i} (- (\epsilon_ {l - 1} ^ {i j}) ^ {2}) \epsilon_ {l - 1} ^ {i j} k _ {j} \\ + \sum_ {k \in [ l ]} \sum_ {i ^ {\prime} \in \Phi_ {l} ^ {k}} s o f t m a x _ {i ^ {\prime}} (- (\epsilon_ {l} ^ {i ^ {\prime} k}) ^ {2}) \epsilon_ {l} ^ {i ^ {\prime} k} w _ {l} ^ {i ^ {\prime} k} \mathbb {1} (i ^ {\prime} = j) \\ \end{array}
$$
$$
\begin{array}{l} \frac {\partial F}{\partial \mu_ {j}} = - \sum_ {i \in \Phi_ {l - 1} ^ {j}} s o f t m a x _ {i} (- (\epsilon_ {l - 1} ^ {i j}) ^ {2}) \epsilon_ {l - 1} ^ {i j} \\ + \sum_ {k \in [ l ]} \operatorname {s o f t m a x} _ {i ^ {\prime}} \left(- \left(\epsilon_ {l} ^ {i ^ {\prime} k}\right) ^ {2}\right) \epsilon_ {l} ^ {i ^ {\prime} k} w _ {l} ^ {i ^ {\prime} k} \\ \end{array}
$$
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  $\phi$ )
While we derived these equations for individual units to draw an easy comparison to standard Predictive Coding, we note it is 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-Clifford theorem indicates a deeper duality between attention as mediated by precision matrices and as structural inference.
7.3 New Designs
Multihop Derivation  $\mathbb{E}_{y|x,\phi}[y_i]$  in transformer attention, a linear transformation is applied to the most likely neighbour,  $x_{j}$ , of  $x_{i}$ . A natural extension is to include a two-hop neighbourhood, additionally using the most likely neighbour  $x_{k}$  of  $x_{j}$ . Formally, the value function  $v$  no longer neatly distributes over the partition  $\Phi_{i}$ , however the attention mechanism then takes a different form:  $\mathbb{E}_{p(\phi_j|\phi_i)p(\phi_i|x)}[V(x_{\phi_i} + x_{\phi_j})] = (P_\phi +P_\phi^2)VX$ . Where we use  $\phi_{j(i)} = \phi_j$  to denote the edge set of the node at the end of  $\phi_{i}$ . To see this note:
$$
\begin{array}{l} \mathbb {E} _ {p (\phi \mid x)} [ V (x _ {\phi_ {i}} + x _ {\phi_ {j}}) ] = \sum_ {\phi} \prod_ {k} p \left(\phi_ {k} \mid x\right) V \left(x _ {\phi_ {i}} + x _ {\phi_ {j}}\right) \\ = \sum_ {\phi} \prod_ {k} p \left(\phi_ {k} \mid x\right) V \left(x _ {\phi_ {i}} + x _ {\phi_ {j}}\right) \\ = \sum_ {\phi} \prod_ {k} p \left(\phi_ {k} \mid x\right) V x _ {\phi_ {i}} + \sum_ {\phi} \prod_ {k} p \left(\phi_ {k} \mid x\right) V x _ {\phi_ {j}} \\ \end{array}
$$
by independence properties
$$
= \sum_ {\phi_ {i}} p \left(\phi_ {i} \mid x\right) V x _ {\phi_ {i}} + \sum_ {\phi_ {i}, \phi_ {j}} p \left(\phi_ {i} \mid x\right) p \left(\phi_ {j} \mid x\right) V x _ {\phi_ {j}}
$$
Denoting the typical attention matrix,  $P$ , where  $p_{ij} = p(\phi_i = [j] \mid x)$
$$
\begin{array}{l} = \sum_ {k} \sum_ {j} p _ {j k} p _ {i j} V x _ {k} + \sum_ {j} p _ {i j} V x _ {j} \\ = \left(P _ {\phi} + P _ {\phi} ^ {2}\right) V X \\ \end{array}
$$