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

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

In total, we only need 𝒪 ⁡ ( c 2 ) \mathcal{O}(c^{2}) -many parameters to perform these multiplications and the linear bias in the first layer is a hyperparameter that is static and unlearned; the time complexity comes from the multiplication of 𝐐𝐊 ⊤ {\bf QK^{\top}} in 𝒪 ⁡ ( N 2 ​ c ) {\mathcal{O}}(N^{2}c) , and projections in 𝒪 ⁡ ( N ​ c 2 ) {\mathcal{O}}(Nc^{2}) . Finally, we only need 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layers for this solution. ∎

In the sequel, we develop a parallel algorithm to solve the multiple-query associative recall problem with 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) \mathcal{O}(Nd\cdot\log^{2}{N}) work complexity and 𝒪 ⁡ ( d ⋅ log 2 ⁡ N ) \mathcal{O}(d\cdot\log^{2}{N}) time. We then convert the algorithm into a BaseConv model via the route of arithmetic circuits, which then solves the multiple-query associative recall problem with 𝒪 ~ ​ ( 1 ) \tilde{\mathcal{O}}(1) layers and 𝒪 ~ ​ ( N ​ d ) \tilde{\mathcal{O}}(Nd) parameters.

#### H.7.3 Initial Attempt: A Sequential Algorithm

We will first discuss the algorithm that simply uses an associative array to solve the multiple-query associative recall problem. Specifically, we want to use a data structure that allows for logarithmic insertion and membership query. Here, we do not specify a choice but data structures including self-balancing binary search trees which allow for 𝒪 ⁡ ( log ⁡ N ⋅ d ) \mathcal{O}(\log{N}\cdot d) 𝚒𝚗𝚜𝚎𝚛𝚝 {\tt insert} and 𝚏𝚒𝚗𝚍 {\tt find} operations for d d -bit entries should be sufficient.

###### Proposition H.28 .

Algorithm 9 solves the multiple-query associative recall problem ( Mqar ) \left(\textsc{Mqar}\right) in 𝒪 ⁡ ( d ​ N ​ log ⁡ N ) \mathcal{O}(dN\log{N}) time for an input sequence 𝐮 ∈ { 0 , 1 } N × d \bm{u}\in\{0,1\}^{N\times d} .

###### Proof.

For any i ∈ { 0 , … , N 3 − 1 } i\in\{0,\ldots,\frac{N}{3}-1\} , we know that both insertion and lookup operations take 𝒪 ⁡ ( log ⁡ ( i ) ⋅ d ) \mathcal{O}(\log(i)\cdot d) time. Overall, the runtime of the algorithm is ∑ i = 0 N 3 − 1 𝒪 ⁡ ( log ⁡ ( i ) ⋅ d ) = 𝒪 ⁡ ( log ⁡ ( N ! ) ⋅ d ) = 𝒪 ⁡ ( N ​ log ​ N ⋅ d ) . \sum_{i=0}^{\frac{N}{3}-1}\mathcal{O}(\log(i)\cdot d)=\mathcal{O}(\log(N!)\cdot d)=\mathcal{O}(N\log{N}\cdot d). ∎

#### H.7.4 Algorithm via Parallel Binary Search

Our plan is to convert the algorithm for solving the multiple-query associative recall problem in the RAM model into an arithmetic circuit, which by Theorem H.21 will lead to a BaseConv model that solves the multiple-query associative recall problem. With respect to Algorithm 9 , it may be the case that the arithmetic circuit has a large number of layers Ω ⁡ ( N ) \Omega(N) . Unfortunately, this would imply that the resulting BaseConv model may have near quadratic complexity in N N . Instead, we now initiate our effort into designing a BaseConv model with both small enough number of parameters and number of layers. Here, we will first subdivide the problem using dyadic intervals into 𝒪 ⁡ ( N ) \mathcal{O}(N) subproblems and reduce each such subproblem into a multiple search problem [ Akl and Meijer, 1990 ] . To this end, we briefly introduce the multiple search problem below. Given two array of numbers A ≜ a 0 ≤ … ≤ a n − 1 A\triangleq a_{0}\leq\ldots\leq a_{n-1} and B ≜ ( b 0 ≤ … ≤ b m − 1 ) B\triangleq(b_{0}\leq\ldots\leq b_{m-1}) with n ≤ m n\leq m , for each a j ∈ A a_{j}\in A , the goal is to find the smallest element in B B that is larger than or equal to a j a_{j} . The multiple search problem is solved by a parallel binary search ( pbs ) algorithm in [ Akl and Meijer, 1990 ] with work complexity 𝒪 ⁡ ( n ⋅ log ⁡ m ) \mathcal{O}(n\cdot\log{m}) and time 𝒪 ⁡ ( log ⁡ n ​ log ⁡ m ) \mathcal{O}(\log{n}\log{m}) . Specifically, for sorted arrays A [ 0 ⋯ n − 1 ] A[0\cdots n-1] and B [ 0 ⋯ m − 1 ] B[0\cdots m-1] , 𝚙𝚋𝚜 {\tt pbs} constructs the array C [ 0 ⋯ n − 1 ] C[0\cdots n-1] defined as C ⁡ [ i ] ≜ { min 0 ≤ j < m { j | A [ i ] ≤ B [ j ] } if ​ A ​ [ i ] ≤ B ⁡ [ m − 1 ] m otherwise. C[i]\triangleq\begin{cases}\min_{0\leq j<m}\{j\lvert\ A[i]\leq B[j]\}&\text{if }A[i]\leq B[m-1]\\ m&\text{otherwise.}\end{cases} (12) The algorithm itself runs in exclusive-read exclusive-write (EREW) PRAM model—no two processors are allowed to read from or write into the same memory location at the same time.

We now augment the algorithm copied from [ Akl and Meijer, 1990 ] for our purposes below.

Let Σ \Sigma be the set { 0 , 1 } \{0,1\} and denote the set of binary strings of size n n as Σ n \Sigma^{n} . We define prefix ⁡ ( 𝒙 ) \mathrm{prefix}(\bm{x}) for n n -bit strings as the set of all initial substrings of 𝒙 ∈ Σ n \bm{x}\in\Sigma^{n} which includes the empty string and 𝒙 \bm{x} itself. Next, let dec : { 0 , 1 } n → ℕ \mathrm{dec}:\{0,1\}^{n}\to\mathbbm{N} be the decimal representation of an n n -bit string 𝒙 \bm{x} with 𝒙 ⁡ [ 0 ] \bm{x}[0] denoting the least significant bit. We also use 𝚜𝚘𝚛𝚝 ⁡ ( A ) {\tt sort}(A) as a procedure that sorts an array A A . Finally, wlog, we assume that N N is a power of 2 2 . We are now ready to present a parallel algorithm that solves the multiple-query associative recall problem below.

###### Remark H.29 .

In lines 8 and 11 , we keep track of the sorted permutation of the indices of keys and queries, respectively. This helps us in the retrieval of the index of the matching key as in line 17 .

#### H.7.5 Correctness and Complexity

###### Proposition H.30 .

Algorithm 11 solves the multiple-query associative recall problem with work complexity 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) \mathcal{O}(Nd\cdot\log^{2}{N}) and time 𝒪 ⁡ ( d ⋅ log 2 ⁡ N ) \mathcal{O}(d\cdot\log^{2}{N}) .

###### Proof.

The correctness of pbs implies the correctness of Algorithm 11 if we can show that, for each 1 ≤ i < N 3 1\leq i<\frac{N}{3} , we check for a match among the keys { 𝒌 j } j ∈ [ i − 1 ] \{\bm{k}_{j}\}_{j\in[i-1]} . To this end, for each 1 ≤ i < N 3 1\leq i<\frac{N}{3} , let the set of all iterator indices associated with an index i i be defined as K i ≜ { ( k , x ) | i ∈ J k 𝒙 } K_{i}\triangleq\{(k,x)\lvert i\in J_{k}^{\bm{x}}\} with J k 𝒙 J_{k}^{\bm{x}} as noted in line 10 . Then, we define the corresponding set for keys as ℐ i ≜ ⋃ ( k , 𝒙 ) ∈ K i I k 𝒙 ​ with I k 𝒙 s defined as in line 7 . \mathcal{I}_{i}\triangleq\bigcup_{(k,\bm{x})\in K_{i}}I_{k}^{\bm{x}}\text{ with $I_{k}^{\bm{x}}$s defined as in line \ref{step: i-def}}. That is, for all the calls to pbs-key-values that i i is part of (given by K i K_{i} ) where the algorithm checks for a match among the keys in ℐ i \mathcal{I}_{i} , it then suffices to show that ℐ i = [ i − 1 ] \mathcal{I}_{i}=[i-1] .

Here, first note that if some index j ∈ I k 𝒙 ⊆ ℐ i j\in I_{k}^{\bm{x}}\subseteq\mathcal{I}_{i} for some 𝒙 ∈ Σ log ⁡ ( N 3 ) − k \bm{x}\in\Sigma^{\log\left(\frac{N}{3}\right)-k} , then, by definition, 𝒙 ∈ prefix ⁡ ( bin ⁡ ( j ) ) \bm{x}\in\mathrm{prefix}(\mathrm{bin}(j)) . Here, let 𝒙 1 := 𝒙 | 𝒙 ⁡ [ log ⁡ ( N 3 ) − k ] = 1 \bm{x}^{1}:=\bm{x}\lvert_{\bm{x}[\log\left(\frac{N}{3}\right)-k]=1} where we set the ( log ⁡ ( N 3 ) − k ) (\log\left(\frac{N}{3}\right)-k) -th index of 𝒙 \bm{x} to be 1. Consequently, as we have i ∈ J k 𝒙 i\in J^{\bm{x}}_{k} for the same k k and 𝒙 \bm{x} as in I k 𝒙 I_{k}^{\bm{x}} ( cf. line 7 ), we must have 𝒙 1 ∈ prefix ⁡ ( bin ⁡ ( i ) ) \bm{x}^{1}\in\mathrm{prefix}(\mathrm{bin(i)}) . Thus we get j < i j<i , whence we can claim that ℐ i ⊆ [ i − 1 ] \mathcal{I}_{i}\subseteq[i-1] .

For the other direction, for any i i , let b b denote the position of the most significant bit in bin ⁡ ( i ) \mathrm{bin}(i) which differs from bin ⁡ ( j ) \mathrm{bin}(j) for any j ∈ [ i − 1 ] j\in[i-1] . Then, there must exist a binary string that is in the prefix set of both bin ⁡ ( i ) \mathrm{bin}(i) and bin ⁡ ( i ) \mathrm{bin}(i) . That is, there exists 𝒙 ∈ prefix ⁡ ( bin ⁡ ( i ) ) ∩ prefix ⁡ ( bin ⁡ ( j ) ) \bm{x}\in\mathrm{prefix}(\mathrm{bin}(i))\cap\mathrm{prefix}(\mathrm{bin}(j)) with 𝒙 ∈ Σ b \bm{x}\in\Sigma^{b} . Thus, we then must have bin ⁡ ( j ) ∈ I log ⁡ ( N 3 ) − b 𝒙 \mathrm{bin}(j)\in I^{\bm{x}}_{\log\left(\frac{N}{3}\right)-b} and bin ⁡ ( i ) ∈ J log ⁡ ( N 3 ) − b 𝒙 \mathrm{bin}(i)\in J^{\bm{x}}_{\log\left(\frac{N}{3}\right)-b} with 𝒙 \bm{x} as the corresponding witness. Hence, we have [ i − 1 ] ⊆ ℐ i [i-1]\subseteq\mathcal{I}_{i} .

Overall, we have shown that ℐ i = [ i − 1 ] \mathcal{I}_{i}=[i-1] . Since this holds for all 1 ≤ i < N 3 1\leq i<\frac{N}{3} , we can conclude that Algorithm 11 solves the multiple-query associative recall problem.

Next, it is easy to see that we execute lines 7 to 17 ∑ k = 0 log ⁡ ( N 3 ) − 1 N 3 2 k + 1 \sum_{k=0}^{\log\left(\frac{N}{3}\right)-1}\frac{\frac{N}{3}}{2^{k+1}} -many times. We note that sorting n n values each of size d d can be done with work complexity n ​ log ⁡ n ⋅ d n\log{n}\cdot d . We note that, at each instance, we are sorting sorting 2 k 2^{k} values. Meanwhile, remembering 2 k 2^{k} sorted permutation of indices can be done in linear time using arrays. Moreover, each call to pbs-key-values has n = m = 2 k n=m=2^{k} which has work complexity n ​ log ⁡ m n\log{m} . Finally, we know that the work complexity of lines 19 to 21 is 𝒪 ⁡ ( N ) \mathcal{O}(N) . Thus, the overall work complexity of Algorithm 11 is d ⋅ ∑ k = 0 log ⁡ ( N 3 ) − 1 N 3 2 k + 1 ​ 𝒪 ​ ( 2 k ⋅ log ⁡ 2 k ) = d ⋅ 𝒪 ⁡ ( N ) ⋅ ∑ k = 0 log ⁡ N / 3 − 1 𝒪 ⁡ ( k ) = 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) . d\cdot\sum_{k=0}^{\log\left(\frac{N}{3}\right)-1}\frac{\frac{N}{3}}{2^{k+1}}\mathcal{O}(2^{k}\cdot\log{2^{k}})=d\cdot\mathcal{O}(N)\cdot\sum_{k=0}^{\log{N/3}-1}\mathcal{O}(k)=\mathcal{O}(Nd\cdot\log^{2}{N}). (13)

We will now analyze the depth of Algorithm 11 . We know that the depth of computation for Algorithm 10 is 𝒪 ⁡ ( log ⁡ n ​ log ⁡ m ) \mathcal{O}(\log{n}\log{m}) for input sizes. Moreover, we have 𝒪 ⁡ ( 1 ) \mathcal{O}(1) depth for the computation in 19 to 21 as each entry in 𝚒𝚍𝚡 {\tt idx} can have at most one entry. Since the nested for loops iterating over k k s and the associated 𝒙 \bm{x} s runs in parallel, the depth of Algorithm 11 is dominated by the largest depth among all calls to pbs-key-values and to sort . The largest such call to pbs-key-values is of size n = m = 2 log ⁡ ( N 3 ) − 1 = N / 6 n=m=2^{\log\left(\frac{N}{3}\right)-1}=N/6 which yields a depth of d ⋅ log 2 ⁡ N 3 d\cdot\log^{2}{\frac{N}{3}} . Moreover, using sorting networks ( Definition H.34 ), we know that the largest depth is for sorting N 6 \frac{N}{6} values of size d d given by d ⋅ Θ ⁡ ( log ⁡ N ) d\cdot\Theta(\log{N}) ( Lemma H.35 ). Thus, we can conclude that Algorithm 11 takes 𝒪 ⁡ ( d ⋅ log 2 ⁡ N ) \mathcal{O}(d\cdot\log^{2}{N}) , time where N N is the length of the input. ∎

#### H.7.6 Conversion to Arithmetic Circuit

We will convert Algorithm 11 to an arithmetic circuits modularly. In particular, after writing out an explicit circuit for Algorithm 10 , we will uses this circuit as a black-box along with circuits for sorting networks.

##### Circuit for pbs-key-values .

We will denote the corresponding arithmetic circuit for Algorithm 10 as pbs-key-values as well with the input gates comprising of each entry from 𝒒 [ s ⋯ t ] \bm{q}[s\cdots t] and 𝒌 [ x ⋯ y ] \bm{k}[x\cdots y] and the i i -th output gate yielding the value C ⁡ [ i ] C[i] as in Algorithm 10 .

Here, we first convert the comparisons for the if statements in Algorithm 10 . To this end, we briefly introduce comparators.

###### Definition H.31 (Comparators) .

A comparator is a device with inputs 𝒙 \bm{x} and 𝒚 \bm{y} and outputs 𝒙 ′ \bm{x}^{\prime} and 𝒚 ′ \bm{y}^{\prime} , that performs the following function: 𝒙 ′ \displaystyle\bm{x}^{\prime} = min ⁡ ( 𝒙 , 𝒚 ) , \displaystyle=\min(\bm{x},\bm{y}), 𝒚 ′ \displaystyle\bm{y}^{\prime} = max ⁡ ( 𝒙 , 𝒚 ) . \displaystyle=\max(\bm{x},\bm{y}).

Using comparators, we can use bit-wise XOR and AND to define the result of the comparisons in lines 6 and 11 as the following fixed variables: ℓ x \displaystyle\ell_{x} : = 𝟙 { 𝒒 [ mid ] ≤ 𝒌 x } , \displaystyle:=\mathbbm{1}\{\bm{q}[\mathrm{mid}]\leq\bm{k}_{x}\}, (14) g y \displaystyle g_{y} : = 𝟙 { 𝒒 [ mid ] > 𝒌 y } , \displaystyle:=\mathbbm{1}\{\bm{q}[\mathrm{mid}]>\bm{k}_{y}\}, z \displaystyle z : = bin-search ​ ( { 𝒌 i } i = s t , 𝒒 ⁡ [ mid ] ) . \displaystyle:=\texttt{bin-search}({\{\bm{k}_{i}\}_{i=s}^{t},\ \bm{q}[\mathrm{mid}]}). This then allows us to infer the index of the key array for each recursive call to pbs-key-values in lines 9 , 14 , 19 , and 20 from Algorithm 10 . Specifically, let z s z_{s} and z t − 1 z_{t}-1 denote the starting and ending indices for the keys as inputs to the recursive calls in Algorithm 10 ) below: pbs-key-values ( 𝒒 [ mid + 1 ⋯ t ] , 𝒌 [ z t ⋯ y ] ) ; (lines 9 and 20 ) \displaystyle\texttt{pbs-key-values}(\displaystyle\bm{q}[\mathrm{mid}+1\cdots t],\ \bm{k}[z_{t}\cdots y]);\text{ (lines \ref{step: rec-1} and \ref{step: right})} (15) pbs-key-values ( 𝒒 [ s ⋯ mid − 1 ] , 𝒌 [ x ⋯ z s − 1 ] ) ; (lines 14 and 19 ) \displaystyle\texttt{pbs-key-values}(\bm{q}[s\cdots\mathrm{mid}-1],\ \bm{k}[x\ \cdots z_{s}-1]);\text{ (lines \ref{step: rec-2} and \ref{step: left})} (16) Here, z t z_{t} and z s z_{s} can assume values dependent on the results of the comparisons in lines 6 and 11 . Specifically, we have z t \displaystyle z_{t} = ℓ x ⋅ x + ( 1 − ℓ x ) ​ ( 1 − g y ) ⋅ z = { x if ​ 𝒒 ​ [ mid ] ≤ 𝒌 x ​ (line 9 ) z if ​ 𝒒 ​ [ mid ] ∈ ( 𝒌 x , 𝒌 y ] ​ (line 20 ) 0 otherwise , \displaystyle=\ell_{x}\cdot x+(1-\ell_{x})(1-g_{y})\cdot z=\begin{cases}x&\text{if }\bm{q}[\mathrm{mid}]\leq\bm{k}_{x}\text{ (line \ref{step: rec-1})}\\ z&\text{if }\bm{q}[\mathrm{mid}]\in(\bm{k}_{x},\bm{k}_{y}]\text{ (line \ref{step: right})}\\ 0&\text{otherwise}\end{cases}, z s \displaystyle z_{s} = g y ⋅ ( y + 1 ) + ( 1 − ℓ x ) ​ ( 1 − g y ) ⋅ z = { y + 1 if ​ 𝒒 ​ [ mid ] > 𝒌 y ​ (line 14 ) z if ​ 𝒒 ​ [ mid ] ∈ ( 𝒌 x , 𝒌 y ] ​ (line 19 ) 0 otherwise . \displaystyle=g_{y}\cdot(y+1)+(1-\ell_{x})(1-g_{y})\cdot z=\begin{cases}y+1&\text{if }\bm{q}[\mathrm{mid}]>\bm{k}_{y}\text{ (line \ref{step: rec-2})}\\ z&\text{if }\bm{q}[\mathrm{mid}]\in(\bm{k}_{x},\bm{k}_{y}]\text{ (line \ref{step: left})}\\ 0&\text{otherwise}\end{cases}. Here, z s z_{s} or z t z_{t} getting a value 0 0 signifies that the branch is dead, and we do not execute the recursive call.

Finally, let the arrays C t [ mid + 1 ⋯ t ] C_{t}[\mathrm{mid}+1\cdots t] and C s [ s ⋯ mid − 1 ] C_{s}[s\cdots\mathrm{mid}-1] denote the outputs to the recursive calls in equation 15 and equation 16 , respectively. We can then succinctly express the outputs for each index of the output array C C as C ⁡ [ i ] = { ℓ x ⋅ x + z s ⋅ ( C 1 ​ [ i ] ) i ∈ [ s ⋯ mid − 1 ] g y ⋅ ( y + 1 ) + z t ⋅ ( C 2 ​ [ i ] ) i ∈ [ mid + 1 ⋯ t ] ℓ x ⋅ x + ( 1 − ℓ x ) ​ ( 1 − g y ) ⋅ z + g y ⋅ ( y + 1 ) i = mid C[i]=\begin{cases}\ell_{x}\cdot x+z_{s}\cdot(C_{1}[i])&i\in[s\cdots\mathrm{mid}-1]\\ g_{y}\cdot(y+1)+z_{t}\cdot(C_{2}[i])&i\in[\mathrm{mid}+1\cdots t]\\ \ell_{x}\cdot x+(1-\ell_{x})(1-g_{y})\cdot z+g_{y}\cdot(y+1)&i=\mathrm{mid}\\ \end{cases} (17)

We can thus state the circuit schematically in Fig. 9 .

Now, before accounting for the complexity of the circuit for Algorithm 10 , we must first assert the complexity of the comparators that we use in Fig. 9 .

###### Lemma H.32 ( Cormen et al. [2022] ) .

For binary strings 𝐱 , 𝐲 ∈ Σ d \bm{x},\bm{y}\in\Sigma^{d} of length d d , there exists a comparison network of size 𝒪 ⁡ ( d ) \mathcal{O}(d) , width 𝒪 ⁡ ( d ) \mathcal{O}(d) , and depth 𝒪 ⁡ ( log ⁡ d ) \mathcal{O}(\log{d}) .

We use this lemma to compute the complexity of the arhmetic circuit from Fig. 9 .

###### Proposition H.33 .