Zoology: Measuring and Improving Recall in Efficient Language Models — part 13/14

Cleaned-text code points [185795, 199956). Bibliography omitted here; source.md includes it.

There exists an ( 𝒪 ⁡ ( ( n + m ) ​ d ) , 𝒪 ⁡ ( n ​ d ⋅ ( log ⁡ m + log ⁡ n ) ) , 𝒪 ⁡ ( log ⁡ n ⁡ ( log ⁡ m + log ⁡ n ​ log ​ d ) ) , 𝒪 ⁡ ( n ​ d ) ) (\mathcal{O}((n+m)d),\mathcal{O}(nd\cdot(\log{m}+\log{n})),\mathcal{O}(\log{n}(\log{m}+\log{n}\log{d})),\mathcal{O}(nd)) -arithmetic circuit 14 14 14 Recall that a ( n , s , Δ , w ) (n,s,\Delta,w) -arithmetic circuit is an n n -variate circuit with size s s , depth at most Δ \Delta , and width w w . equivalent to pbs-key-values ( Algorithm 10 ) with inputs 𝐪 \bm{q} and 𝐤 \bm{k} of lengths n n and m m with d d -bit entries.

###### Proof.

The size of the circuit for pbs-key-values should equal the work-complexity of Algorithm 10 but we also need to account for the comparison gates in equation 14 . Further, the circuit for binary search also has a size of 𝒪 ⁡ ( d ⋅ n ) \mathcal{O}(d\cdot n) instead of 𝒪 ⁡ ( d ⋅ log ⁡ n ) \mathcal{O}(d\cdot\log{n}) work as in Algorithm 10 . Using Lemma H.32 , along with the fact that the comparison gates and the binary search are used at most 𝒪 ⁡ ( log ⁡ n ) \mathcal{O}(\log{n}) times, we deduce that we are adding 𝒪 ⁡ ( n ​ d ​ log ⁡ n ) \mathcal{O}(nd\log{n}) size to the circuit in addition to the work complexity of the parallel algorithm. Thus, the overall size of the arithmetic circuit for pbs-key-values is 𝒪 ⁡ ( d ​ n ​ log ⁡ m + n ​ d ​ log ⁡ n ) \mathcal{O}(dn\log{m}+nd\log{n}) . Further, the depth of the circuit here is determined by the runtime of Algorithm 10 along with the depth of the comparison gates and binary search 𝒪 ⁡ ( log ⁡ n ​ log ⁡ d ) \mathcal{O}(\log{n}\log{d}) , and finally, at most n n processors are used for the parallel algorithm in Algorithm 10 , which yields the width of the circuit. ∎

##### Circuit for Parallel-MQAR .

Now, we will call the circuit for Parallel-MQAR with the same name while the input gates contain the inputs of Algorithm 11 . Indeed, we can directly “translate” Algorithm 11 to an arithmetic circuit as the values for I k 𝒙 I_{k}^{\bm{x}} and I k 𝒙 I_{k}^{\bm{x}} for each 𝒙 \bm{x} and k k are predetermined from N N . Thus, we start by placing the corresponding sorting networks which feeds into the pbs-key-values ( 𝒒 [ s ⋯ t ] , 𝒌 [ x ⋯ y ] , n , m ) \texttt{pbs-key-values}\left(\bm{q}[s\cdots t],\ \bm{k}[x\cdots y],n,m\right) circuit for Algorithm 10 in Fig. 9 so that the output values from the calls to pbs-key-values result in the checks as in line 17 of Algorithm 11 . That is, we get outputs C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))] from each call to pbs-key-values . We can then use a comparison gate to check if this value equals 2 k 2^{k} , and if not, we have found a match C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))] for the query 𝒒 i \bm{q}_{i} which results in the output of the associated value 𝒗 C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] + 1 \bm{v}_{C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]+1} , exactly as in Algorithm 11 . That is, we first define the following variable as the output of the comparison gate: 𝒄 dec ⁡ ( i ) k := 𝟙 { C k [ J permuted k ​ 𝒙 ( dec ( i ) ) ] ≠ 2 k } . \bm{c}_{\mathrm{dec}(i)}^{k}:=\mathbbm{1}\{C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\neq 2^{k}\}. (18) Here, as C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] ≠ 2 k C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\neq 2^{k} implies that I permuted k ​ 𝒙 ​ ( C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] ) I_{\mathrm{permuted}}^{k\bm{x}}\left(C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\right) equals the index of the matching key 𝒌 j \bm{k}_{j} corresponding to the query 𝒒 i \bm{q}_{i} , the i i th output is then simply given by 𝒄 dec ⁡ ( i ) k ⋅ I permuted k ​ 𝒙 ​ ( C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] ) \bm{c}_{\mathrm{dec}(i)}^{k}\cdot I_{\mathrm{permuted}}^{k\bm{x}}\left(C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\right) , where the 0 0 output implies that there does not exist a matching key.

Here, we also briefly introduce the the sorting networks that we use to sort the keys and queries:

###### Definition H.34 (Informal) .

Sorting networks are circuits with gates and wires where the gates of the circuit are comparators ( Definition H.31 ) connecting two wires. Each such circuit can perform sorting on a fixed number of values.

We can then show the circuit schematically as in Fig. 10 .

We now dilineate the complexity of the circuit, starting with the complexity of the sorting networks.

###### Lemma H.35 ( Ajtai et al. [1983] ) .

Let A A be an array with d d -bit entries of size n n . Then, one can implement a sorting network to sort the array A A with size 𝒪 ⁡ ( d ⋅ n ​ log ⁡ n ) \mathcal{O}(d\cdot n\log{n}) and depth 𝒪 ⁡ ( log ⁡ d ​ log ⁡ n ) \mathcal{O}(\log{d}\log{n}) .

###### Proposition H.36 .

There exists an ( N ⋅ d , 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) , 𝒪 ⁡ ( log ⁡ d ​ log 2 ​ N ) , 𝒪 ⁡ ( N ​ d ​ log ​ N ) ) (N\cdot d,\mathcal{O}(Nd\cdot\log^{2}{N}),\mathcal{O}(\log{d}\log^{2}{N}),\mathcal{O}(Nd\log{N})) -arithmetic circuit that solves the multiple-query associative recall problem.

###### Proof.

We note here that for each k k , there are N 3 2 k + 1 \frac{\frac{N}{3}}{2^{k+1}} parallel problems of size 2 k 2^{k} for both the sorting networks and the 𝚙𝚋𝚜 − 𝚔𝚎𝚢 − 𝚟𝚊𝚕𝚞𝚎𝚜 {\tt pbs-key-values} circuit. Using Lemma H.35 , the cumulative size of these sorting networks is 𝒪 ⁡ ( d ⋅ N ​ log 2 ​ N ) \mathcal{O}(d\cdot N\log^{2}{N}) (see equation 13 ) with overall depth 𝒪 ⁡ ( log ⁡ d ​ log ⁡ N ) \mathcal{O}(\log{d}\log{N}) .

Similarly, the next layer again runs ∑ k = 0 log ⁡ N 3 − 1 N 3 2 k + 1 \sum_{k=0}^{\log{\frac{N}{3}}-1}\frac{\frac{N}{3}}{2^{k+1}} -many circuits for 𝚙𝚋𝚜 − 𝚔𝚎𝚢 − 𝚟𝚊𝚕𝚞𝚎𝚜 {\tt pbs-key-values} each of which has size 𝒪 ⁡ ( 2 k ​ d ​ ( log ⁡ 2 k + log ⁡ 2 k ) ) = O ⁡ ( d ⋅ 2 k ​ log ⁡ 2 k ) \mathcal{O}(2^{k}d(\log{2^{k}}+\log{2^{k}}))=O(d\cdot 2^{k}\log{2^{k}}) , depth 𝒪 ⁡ ( log 2 ⁡ 2 k ​ log ​ d ) \mathcal{O}(\log^{2}{2^{k}}\log{d}) and width 𝒪 ⁡ ( 2 k ) \mathcal{O}(2^{k}) ( Proposition H.33 ). Again, the cumulative size of this layer is given by 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) \mathcal{O}(Nd\cdot\log^{2}{N}) (see equation 13 ). Since we run each of these circuits in parallel, the depth of this layer is again 𝒪 ⁡ ( log ⁡ d ​ log 2 ⁡ ( N ) ) \mathcal{O}(\log{d}\log^{2}(N)) while the width is 𝒪 ⁡ ( N ⋅ log ⁡ N ) \mathcal{O}(N\cdot\log{N}) .

Finally, we perform N 3 ​ log ⁡ N 3 \frac{N}{3}\log{\frac{N}{3}} comparisons at the end of d d -bit strings in parallel which results in size 𝒪 ⁡ ( N ​ log ⁡ N ⋅ d ) \mathcal{O}(N\log{N}\cdot d) , depth 𝒪 ⁡ ( log ⁡ d ) \mathcal{O}(\log{d}) and width 𝒪 ⁡ ( N ​ log ⁡ N ⋅ d ) \mathcal{O}(N\log{N}\cdot d) ( Lemma H.32 ). Therefore, the resulting arithmetic circuit has size 𝒪 ⁡ ( d ⋅ N ​ log 2 ​ N + N ​ d ⋅ log 2 ⁡ N + N ​ log ⁡ N ⋅ d ) = O ⁡ ( N ​ d ​ log 2 ​ N ) \mathcal{O}(d\cdot N\log^{2}{N}+Nd\cdot\log^{2}{N}+N\log{N}\cdot d)=O(Nd\log^{2}{N}) , depth 𝒪 ⁡ ( log ⁡ d ​ log 2 ​ N ) \mathcal{O}(\log{d}\log^{2}{N}) and width 𝒪 ⁡ ( N ​ d ​ log ⁡ N ) \mathcal{O}(Nd\log{N}) . ∎

#### H.7.7 The Resulting BaseConv Model

As we have an arithemtic circuit for solving the multiple-query associative recall problem, we can now invoke Theorem H.21 to claim that there is a corresponding BaseConv model that solves the multiple-query associative recall problem with 𝒪 ~ ​ ( N ​ log ⁡ c ) \tilde{\mathcal{O}}(N\log{c}) parameters and 𝒪 ⁡ ( 1 ) ~ \tilde{\mathcal{O}(1)} layers.

###### Theorem H.37 .

There exists a ( N ​ d , 𝒪 ~ ​ ( 1 ) , 𝒪 ~ ​ ( 1 ) , 𝒪 ~ ​ ( N ​ d ) , 𝒪 ~ ​ ( 1 ) ) − BaseConv \left(Nd,\tilde{\mathcal{O}}(1),\tilde{\mathcal{O}}(1),\tilde{\mathcal{O}}(Nd),\tilde{\mathcal{O}}(1)\right)-\text{{BaseConv}} solves the multiple-query associative recall problem.

###### Proof.

Directly applying Theorem H.21 yields a BaseConv model with the number of layers 𝒪 ⁡ ( log ⁡ d ⋅ log 2 ⁡ N ⋅ log ⁡ N ​ d ​ log ​ N ) = 𝒪 ⁡ ( 1 ) \mathcal{O}(\log{d}\cdot\log^{2}{N}\cdot\log{Nd\log{N}})=\mathcal{O}(1) layers while the claim on the input and inner dimensions follow trivially. ∎

### H.8 Data-Dependent Convolutions

#### H.8.1 Introduction

In this section, we are again concerned with solving the multiple-query associative recall problem ( Section H.7.1 ). However, in contrast to Section H.7.4 , which yields a circuit that is unchanged and works for all inputs, we instead take the viewpoint of adapting the sequence mixing weights ( Section 2 ) with respect to the particular sequence that the model receives as input. More specifically, we take the distance between the tokens in the sequence as a measure for designing data-dependent convolutions.

##### Setup.

To formally setup the problem, as in our discussion of designing a parallel algorithm, we consider the following problem description of the multiple-query associative recall problem. Suppose we are given an input 𝒖 [ 0 ⋯ N − 1 ] ≜ { ( 𝒌 0 , 𝒗 0 , 𝒒 0 ) , … , ( 𝒌 N 3 − 1 , 𝒗 N 3 − 1 , 𝒒 N 3 − 1 ) } \bm{u}[0\cdots N-1]\triangleq\left\{\left(\bm{k}_{0},\bm{v}_{0},\bm{q}_{0}\right),\ldots,\left(\bm{k}_{\frac{N}{3}-1},\bm{v}_{\frac{N}{3}-1},\bm{q}_{\frac{N}{3}-1}\right)\right\} with each 𝒌 i , 𝒗 i , 𝒒 i ∈ C \bm{k}_{i},\bm{v}_{i},\bm{q}_{i}\in C . Here, each token is embedded using the standard one-hot encoding in { 0 , 1 } c \{0,1\}^{c} (i.e. we assume d = c d=c ). 15 15 15 Our arguments do need c = d c=d . However we do not need d = N d=N but we made this simplification for ease of presentation. Our goal is again to check, for each 1 ≤ i ≤ N 3 − 1 1\leq i\leq\frac{N}{3}-1 , whether there exists 0 ≤ j < i 0\leq j<i such that 𝒒 i ≡ 𝒌 j \bm{q}_{i}\equiv\bm{k}_{j} , and if so, output 𝒗 j \bm{v}_{j} . Here, we define the interaction distance between the i i th query 𝒒 i \bm{q}_{i} and the matching key 𝒌 j \bm{k}_{j} as i − j i-j . We then also assume that number of distinct interaction distances is bounded by t t . We can identify these distances using an autocorrelation [ Chatfield, 1995 ] , which has an elegant underlying formulation. We will briefly introduce the relevant mathematical machinery in the context of elucidating the data-dependent model that we seek to develop in the sequel.

##### Autocorrelations

We introduce autocorrelation convolutions. Let 𝒖 ~ ​ [ t ] := 𝒖 ​ [ − t ] \tilde{\bm{u}}[t]:=\bm{u}[-t] , then the cross correlation of two vectors 𝒖 \bm{u} and 𝒗 \bm{v} is given by 𝒖 ⋆ 𝒗 ≜ 𝒖 ~ ∗ 𝒗 . \bm{u}\star\bm{v}\triangleq\tilde{\bm{u}}\ast\bm{v}. The autocorrelation of a vector 𝒖 ∈ ℝ n \bm{u}\in\mathbbm{R}^{n} is the cross correlation of 𝒖 \bm{u} with itself. Moreover, in terms of polynomials, we have 𝒖 ~ ​ ( X ) = X n − 1 ⋅ 𝒖 ⁡ ( 1 / X ) \tilde{\bm{u}}(X)=X^{n-1}\cdot\bm{u}(1/X) . Thus, in analogy with our interpretation of convolution in terms of polynomial multiplication, we characterize the autocorrelation of a vector 𝒖 ∈ ℝ n \bm{u}\in\mathbbm{R}^{n} , given by 𝐰 ∈ ℝ n {\bf w}\in\mathbbm{R}^{n} as follows: 𝐰 = coeff ⁡ ( 𝒖 ⁡ ( X ) ⋅ 𝒖 ~ ​ ( X ) mod X n − 1 ) . {\bf w}=\mathrm{coeff}\left(\bm{u}(X)\cdot\tilde{\bm{u}}(X)\mod{X^{n}-1}\right).

#### H.8.2 BaseConv with kernels generated using Auto-Correlation

We are now ready to describe the model that solves the multiple-query associative recall problem using data-dependent kernels derived using auto-correlations.

##### The Input-Dependent Kernels.

There are several potential strategies to identify the best token-interaction distances for an input using a convolution. Here we will focus on an autocorrelation convolution for exposition. Autocorrelation allows us to identify the top t t distinct shifts of the sequence that result in highly overlapping values (e.g. matches between query and keys in our associative recall setting). We can then construct convolution filters that perform each of these t t shifts. That is, we define a function 𝚃𝚘𝚙 {\tt Top} such that 𝚃𝚘𝚙 ⁡ ( u ⋆ u , t ) {\tt Top}\left(u\star u,t\right) returns a list of the top t t shifts { s ℓ } ℓ ∈ [ t ] \{s_{\ell}\}_{\ell\in[t]} . We then use these top t t distances { s ℓ } ℓ ∈ [ t ] \{s_{\ell}\}_{\ell\in[t]} to define the following two kernels: 𝐡 k ​ ( X ) \displaystyle{\bf h}^{k}(X) ≡ ∑ ℓ ∈ [ t ] X s ℓ + ( ℓ − 1 ) ⋅ N 3 , \displaystyle\equiv\sum_{\ell\in[t]}X^{s_{\ell}+(\ell-1)\cdot\frac{N}{3}}, (19) 𝐡 v ​ ( X ) \displaystyle{\bf h}^{v}(X) ≡ ∑ ℓ ∈ [ t ] X s ℓ − 1 + ( ℓ − 1 ) ⋅ N 3 . \displaystyle\equiv\sum_{\ell\in[t]}X^{s_{\ell}-1+(\ell-1)\cdot\frac{N}{3}}. Here, we note that we only only have two kernels as we will assume N ′ = t ​ N 3 N^{\prime}=t\frac{N}{3} and the shift will be done in ”parallel.” Obviously, one can instead define t t distinct shift kernels but then there is a cost of 𝒪 ⁡ ( t ) \mathcal{O}(t) in the number of layers.

##### Projections.

We define the following projections 𝐊 , 𝐐 , 𝐕 ∈ { 0 , 1 } N × d {\bf K},{\bf Q},{\bf V}\in\{0,1\}^{N\times d} of the input that we shall use below using equation 9 . 𝐊 [ i , : ] \displaystyle{\bf K}[i,:] : = { 𝒖 [ i , : ] if ​ i ∈ 𝒦 , 𝟎 d otherwise , \displaystyle:=\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, (20) 𝐐 [ i , : ] \displaystyle{\bf Q}[i,:] : = { 𝒖 [ i , : ] if ​ i ∈ 𝒬 , 𝟎 d otherwise , \displaystyle:=\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{Q},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, 𝐕 [ i , : ] \displaystyle{\bf V}[i,:] : = { 𝒖 [ i , : ] if ​ i ∈ 𝒱 , 𝟎 d otherwise , \displaystyle:=\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{V},\\ {\bf 0}^{d}&\text{otherwise}\end{cases},

Finally, we present the BaseConv model that solves the multiple-query associative recall problem with input-dependent kernels using 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layer and 𝒪 ⁡ ( t ⋅ N ​ d ) \mathcal{O}(t\cdot Nd) -many parameters.

###### Theorem H.38 .

There exists a BaseConv model with gated data-dependent convolutions that solves the multiple-query associative recall problem on inputs from { 0 , 1 } 3 ​ N × c \{0,1\}^{3N\times c} with the total number of distinct interaction distances bounded by t t in 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layers and 𝒪 ⁡ ( t ⋅ N ​ c ) \mathcal{O}(t\cdot Nc) total parameters.

###### Proof.