SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Expanding Derivation As in the main text, let  $p(\phi \mid q) \sim \text{Geo}(q)$  and  $p(q) \sim \text{Beta}(\alpha, \beta)$ , such that we have the full model  $p(x, \phi, q; \alpha, \beta) = p(x \mid \phi)p(\phi \mid q)p(q; \alpha, \beta)$ . In order to find  $p(\phi \mid x)$  we employ a truncated Mean Field Variational Bayes [39], assuming a factorisation  $p_t(\phi, q) = p_t(\phi)p_t(q)$ , and using the updates:
$$
\begin{array}{l} \ln p _ {t + 1} (\phi) = \mathbb {E} _ {p _ {t} (q)} [ \ln p (x \mid \phi) + \ln p (\phi \mid q) ] + C _ {1} \\ \ln p _ {t + 1} (q) = \mathbb {E} _ {p _ {t} (\phi)} [ \ln p (\phi \mid q) + \ln p (q; \alpha , \beta) ] + C _ {2} \\ \end{array}
$$
By conjugacy the second equation simplifies to a simple update of the beta distribution
$$
\begin{array}{l} \Rightarrow p _ {t + 1} (q) = B e t a \left(\alpha_ {t + 1}, \beta_ {t + 1}\right) \\ \alpha_ {t + 1} = \alpha_ {t} + 1 \\ \beta_ {t + 1} = \beta_ {t} + \mathbb {E} _ {p _ {t} (\phi)} [ \phi ] \\ \end{array}
$$
While the second update can be seen as calculating the posterior given  $q_{t} = \mathbb{E}_{p_{t}(q)}[q]$ ,
$$
\begin{array}{l} \ln p _ {t + 1} (\phi) = \ln p (x \mid \phi) + \mathbb {E} _ {p _ {t} (q)} [ \ln p (\phi \mid q) ] + C _ {2} \\ = \ln p (x \mid \phi) + \phi \mathbb {E} _ {p _ {t} (q)} [ \ln q ] + C _ {2} \\ = \ln p (\phi \mid x, q _ {t}) \\ \end{array}
$$
Finally, we use a truncation to approximate the infinite sum  $\mathbb{E}_{p_t(\phi)}[\phi] = \sum_k p_t(\phi = k)k \approx \sum_{<H} p_t(\phi = k)k$ . Where we set the horizon according to the current distribution of  $q$ . For example in our experiments we chose  $H(q_t) = \ln 0.05 / \ln (1 - q_t)$  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 in the paper.
8.1 Self Attention
- Nodes  $K = Q = \left( {{x}_{1},..,{x}_{n}}\right)$
- Structural prior, over a fully connected, directed graph  $p(\vec{\phi}) = \prod_{i=1}^{n} p(\vec{\phi}_i)$ , where  $\vec{\Phi}_i = \{(x_1, x_i),.., (x_n, x_i)\}$  is the set of edges involving  $x_i$  and  $\vec{\phi}_i \sim \text{Uniform}(\vec{\Phi}_i)$ , such that each node is uniformly likely to connect to every other node in a given direction.
- Edge potentials  $\psi(x_{j}, x_{i}) = x_{i}^{T} W_{Q}^{T} W_{K} x_{j}$ , in effect measuring the similarity of  $x_{j}$  and  $x_{i}'$  in a projected space.
- Value functions  $v_{i}(x, \phi_{i} = [j]) = W_{V} x_{j}$ , a linear transformation applied to the node at the start of the edge  $\phi_{i}$ .
Again, taking the posterior expectation in each of the factors defined in two Eq.2 gives the standard self-attention mechanism
$$
\mathbb {E} _ {p \left(\phi_ {i} \mid Q\right)} \left[ v _ {i} \right] = \sum_ {j} \operatorname {s o f t m a x} _ {j} \left(x _ {i} ^ {T} W _ {Q} ^ {T} W _ {K} x _ {j}\right) W _ {V} x _ {j}
$$
8.2 Positional Encodings and Graph Neural Networks
In Table.1 we show that positional encodings and graph attention are naturally incorporated in this framework. Absolute positional encoding as suggested by Vaswani et al. [1] can be seen as modifying the edge potentials with a vector that depends on position, while relative position encodings can be seen as a 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(\phi_i = [j]) = p_{i,j}$ , it can be fully specified by the matrix  $(P)_{i,j} = p_{i,j}$ . This leads to the modified attention update
$$
\sum_ {j} \operatorname {s o f t m a x} _ {j} \left(x _ {i} Q ^ {T} K x _ {j} + \ln p _ {i j}\right) x _ {j}
$$
However this requires local parameters for each node  $z_{i}$ . A more natural prior assign a different probability to the relative distance of  $i$  from  $j$ . This is achieved with  $P = \text{circ}(p_1, p_2, \dots, p_n)$ , where  $\text{circ}$  is the circulant matrix of  $(p)_{i \leq n}$ . Due to properties of circulant matrices  $\ln P_{ij}$  can be reparameterised with the hartley transform
$$
\ln p _ {i, j} = \sum_ {k} \beta_ {k} [ \cos (k \theta_ {i, j}) + \sin (k \theta_ {i, j}) ] = \beta \cdot b ^ {(i, j)}
$$
Where  $b_{k}^{(i,j)} = \cos (k\frac{i - j}{2\pi n}) + \sin (k\frac{i - j}{2\pi n})$  can be thought of as a relative position encoding, and  $\beta$  are parameters to be learnt.
8.2.1 Block-Slot Attention
Singh et al. [18] suggest combining an 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 us to 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  $X = \{x_{1},\dots,x_{n}\}$ ,  $Z = \{z_{1},\dots,z_{m}\}$  and  $M = \{m_{1},\dots,m_{l}\}$  with a prior over edges  $p(E) = \prod_{j=1}^{m}p(E_{j})\prod_{k=1}^{l}p(\tilde{E}_{k})$ ,  $E_{j} \sim \text{Uniform}[(x_{j},z_{1}),..,(x_{j},z_{m})]$ ,  $\tilde{E}_{k} \sim \text{Uniform}[(z_{1},m_{k}),..,(z_{m},m_{k})]$ , with edge potentials between  $X$  and  $Z$  given by  $\psi(x_{j},z_{i}) = z_{i}Q^{T}Kx_{j}$  and between  $Z$  and  $M$ ,  $\psi(z_{i},m_{k}) = z_{i} \cdot m_{k}$  applying (8) gives
$$
\begin{array}{l} \mu_ {i} ^ {*} = \sum_ {j} s o f t m a x _ {i} (\mu_ {i} Q ^ {T} K x _ {j}) Q ^ {T} K x _ {j} \\ + \sum_ {k} \operatorname {s o f t m a x} _ {k} \left(\mu_ {i} \cdot m _ {k}\right) m _ {k} \\ \end{array}
$$
In the original block-slot attention each slot  $z_{i}$  is broken into blocks, where each block can access block-specific memories i.e.  $z_{i}^{(b)}$  can has possible connections to memory nodes  $\{m_k^{(b)}\}_{k \leq l}$ . Allowing objects to be represented by slots which in turn disentangle features of each object in different blocks. We presented a single block version above, however it is easy to see that the update extends to the multiple block version applying (8) gives
$$
\begin{array}{l} \mu_ {i} ^ {*} = \sum_ {j} s o f t m a x _ {i} (\mu_ {i} Q ^ {T} K x _ {j}) Q ^ {T} K x _ {j} \\ + \sum_ {k, b} s o f t m a x _ {k} (\mu_ {i} ^ {(b)} \cdot m _ {k} ^ {(b)}) m _ {k} ^ {(b)} \\ \end{array}
$$
9 Experimental Details
Following the notation in the text; data generation parameters:
- Total number of tokens: 10
- Embedding dimension (dimension of each  $x$ ):10
- Output dimension (dimension of  $y$ ): 1
-  $\sigma^2$  (autoregressive noise): 1
- Random matrix initialisation was performed with torchRAND
Training parameters (across all models):
- batch size:200
number of batches: 10
- optimiser: ADAM
learning rate:  $1e - 3$
Number of different random seeds: 10
Model: To make analysis easier, all models were prevented from self-attending to the final token.
Total number of tokens: 50
- Embedding dimension(s) (dimension of each  $x$ ): [10, 50]
-  $p$  the parameter for generating a geometric shuffle:[0.5, .2, .1, .04]
- Output dimension (dimension of  $y$ ): 1
-  $\sigma^2$  (autoregressive noise): .1
- Random matrix initialisation was performed with torchRAND
- batch size:1*
number of batches: 10000
- optimiser: ADAM
learning rate:  $5e - 4$
Model: To make analysis easier, all models were prevented from self-attending to the final token. For expanding attention the hyperparameters were set as  $\alpha = .1$ ,  $\beta = .9$  these were chosen to have a mean value at roughly a quarter of 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 synchronisation locks, as the fastest samples wait for the longest ones to complete. In order to take full advantage of a dynamic window over a batch, intelligent asynchronous processing would be necessary.
Algorithm 1 Attention
Require:  $X, W_{q}, W_{k}, W_{v}$ $Q \gets W_{q}X$ $K \gets W_{k}X$ $V \gets W_{v}X$ $A \gets \text{softmax}(K^{T}Q)$ $Y \gets AV$
Algorithm 2 Multihop
Require:  $X, W_{q}, W_{k}, W_{v}, N$ $Q \gets W_{q}X$ $K \gets W_{k}X$ $V \gets W_{v}X$ $A \gets \text{softmax}(K^{T}Q)$ $P \gets \sum_{k < N} A^{k}$ $Y \gets PV$
Algorithm 3 Expanding
Require:  $X,x_{q},W_{q},W_{k},W_{v},\alpha ,\beta ,k = 0$ $q\gets W_qx_q$ $K\gets W_kX$ $V\leftarrow W_vX$    
while  $k$  not converged do  
 $p = \alpha /(\alpha +\beta)$ $k\gets \ln (.05) / \ln (1 - p)$ $g_{p}[i] = -i / k$ $A\gets \mathrm{softmax}(q^{T}K[:k,:] + g_{p})$ $\alpha \gets \alpha +1$ $\beta \gets \beta +\sum_{i <   k}A[i]i$   
end while  
 $y\gets AV$