Zoology: Measuring and Improving Recall in Efficient Language Models β part 14/14
Cleaned-text code points [199956, 215348). Bibliography omitted here; source.md includes it.
We note that we have the input dimension d = c d=c . Now, for any input sequence π β { 0 , 1 } N Γ d \bm{u}\in\{0,1\}^{N\times d} , we get the input-dependent kernels as in equation 19 using autocorrelation of the input. We will now outline the following computations for the BaseConv layers: π \displaystyle\bm{y} = π»πππππ π β ( π ) β ( π‘ K β π»πππππ π β ( π ) ) \displaystyle={\tt Linear}_{{\bf Q}}(\bm{u})\odot\left({\bf h}^{K}*{\tt Linear}_{{\bf K}}(\bm{u})\right) = π β ( π‘ K β π ) \displaystyle={\bf Q}\odot\left({\bf h}^{K}*{\bf K}\right) (21) π \displaystyle\bm{z} = π»πππππ π β ( π ) β ( π‘ V β π»πππππ π β ( π ) ) \displaystyle={\tt Linear}_{{\bf E}}(\bm{y})\odot\left({\bf h}^{V}*{\tt Linear}_{{\bf V}}(\bm{u})\right) = π β ( π‘ V β π ) , \displaystyle={\bf E}\odot\left({\bf h}^{V}*{\bf V}\right), (22) where we have the linear projections π»πππππ π β ( π ) = π {\tt Linear}_{{\bf Q}}(\bm{u})={\bf Q} , π»πππππ π β ( π ) = π {\tt Linear}_{{\bf K}}(\bm{u})={\bf K} , π»πππππ π β ( π ) = π {\tt Linear}_{{\bf V}}(\bm{u})={\bf V} and π»πππππ π β ( π ) = π {\tt Linear}_{{\bf E}}(\bm{y})={\bf E} defined as π [ i , : ] := π»πππππ π ( π ) [ i , : ] = { π d if β β j β [ d ] β such that β π β [ i , j ] = 1 π d otherwise . {\bf E}[i,:]:={\tt Linear}_{{\bf E}}(\bm{y})[i,:]=\begin{cases}{\bf 1}^{d}&\text{if }\exists\ j\in[d]\text{ such that }\bm{y}[i,j]=1\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. (23) Here, we will first present the argument for the special case when we have t = 1 t=1 as that will help us elucidate the general case. To this end, as the kernels from equation 19 for t = 1 t=1 are given by π‘ k β ( X ) \displaystyle{\bf h}^{k}(X) β‘ X s 1 ; \displaystyle\equiv X^{s_{1}}; (24) π‘ v β ( X ) \displaystyle{\bf h}^{v}(X) β‘ X s 1 β 1 , \displaystyle\equiv X^{s_{1}-1}, we observe that convolving with these kernels π‘ β π {\bf h}\ast\bm{y} is equivalent to operating with the following primitives ( Section H.3 ): πππππ β _ β πππ π β ( π , s 1 ) ; \displaystyle\tt{\displaystyle shift\_down}(\bm{y},s_{1}); (25) πππππ β _ β πππ π β ( π , s 1 β 1 ) . \displaystyle\tt{\displaystyle shift\_down}(\bm{y},s_{1}-1). We note that we shift down instead of shifting up as the index of the top-left entry is ( 0 , 0 ) (0,0) . We can then write down the computations performed in equation 21 and equation 22 as follows: π \displaystyle\bm{y} = π β πππππ β _ β πππ π β ( π , s 1 ) \displaystyle={\bf Q}\odot{\tt shift\_down}({\bf K},s_{1}) (26) π \displaystyle\bm{z} = π β πππππ β _ β πππ π β ( π , s 1 β 1 ) , \displaystyle={\bf E}\odot{\tt shift\_down}({\bf V},s_{1}-1), (27) We begin by examining π \bm{y} below: π [ i , : ] \displaystyle\bm{y}[i,:] = ( π β πππππ _ πππ π ( π , s 1 ) ) [ i , : ] \displaystyle=\left({\bf Q}\odot{\tt shift\_down}({\bf K},s_{1})\right)[i,:] = π [ i , : ] β πππππ _ πππ π ( π , s 1 ) [ i , : ] \displaystyle={\bf Q}[i,:]\odot{\tt shift\_down}({\bf K},s_{1})[i,:] (28) = ( { π [ i , : ] if β i β π¬ π d otherwise ) β ( { π [ i β s 1 , : ] if β i β s 1 β π¦ , π d otherwise ) \displaystyle=\left(\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{Q}\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right)\odot\left(\begin{cases}\bm{u}[i-s_{1},:]&\text{if }i-s_{1}\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right) (29) = { π [ i , : ] β π [ i β s 1 , : ] if β i β π¬ β and β i β s 1 β π¦ , π d otherwise \displaystyle=\begin{cases}\bm{u}[i,:]\odot\bm{u}[i-s_{1},:]&\text{if }i\in\mathcal{Q}\text{ and }i-s_{1}\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases} Here, we use the fact that the Hadamard product is row-independent in equation 28 , and the definitions of the projections from equation 20 in equation 29 . Examining the j j th entry, we get π β‘ [ i , j ] β π β‘ [ i β s 1 , j ] \displaystyle\bm{u}[i,j]\odot\bm{u}[i-s_{1},j] = { 1 if β i β π¬ , i β s 1 β π¦ β and β π i β‘ π i β s 1 β‘ π j 0 otherwise. \displaystyle=\begin{cases}1&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\equiv{\bf e}_{j}\\ 0&\text{otherwise.}\end{cases} That is, we can express π [ i , : ] \displaystyle\bm{y}[i,:] = { π j if β i β π¬ , i β s 1 β π¦ β and β π i β‘ π i β s 1 β‘ π j π d otherwise . \displaystyle=\begin{cases}{\bf e}_{j}&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\equiv{\bf e}_{j}\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. (30) Consequently, as per the definition in equation 23 , we get π [ i , : ] = { π d if β i β π¬ , i β s 1 β π¦ β and β π i β‘ π i β s 1 π d otherwise \displaystyle{\bf E}[i,:]=\begin{cases}{\bf 1}^{d}&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} (31) We can now finally specify the output π \bm{z} from equation 27 as follows: π [ i , : ] \displaystyle\bm{z}[i,:] = ( π β πππππ _ πππ π ( π , s 1 β 1 ) ) [ i , : ] \displaystyle=\left({\bf E}\odot{\tt shift\_down}({\bf V},s_{1}-1)\right)[i,:] = π [ i , : ] β πππππ _ πππ π ( π , s 1 β 1 ) [ i , : ] \displaystyle={\bf E}[i,:]\odot{\tt shift\_down}({\bf V},s_{1}-1)[i,:] (32) = ( { π d if β i β π¬ , i β s 1 β π¦ β and β π i β‘ π i β s 1 π d otherwise ) β ( { π [ i β s 1 + 1 , : ] if β i β s 1 + 1 β π± , π d otherwise ) \displaystyle=\left(\begin{cases}{\bf 1}^{d}&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right)\odot\left(\begin{cases}\bm{u}[i-s_{1}+1,:]&\text{if }i-s_{1}+1\in\mathcal{V},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right) (33) = { π [ i β s 1 + 1 , : ] if β i β π¬ , i β s 1 β π¦ , i β s 1 + 1 β π± β and β π i β‘ π i β s 1 π d otherwise \displaystyle=\begin{cases}\bm{u}[i-s_{1}+1,:]&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K},i-s_{1}+1\in\mathcal{V}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} = { π i β s 1 if β π i β‘ π i β s 1 π d otherwise \displaystyle=\begin{cases}\bm{v}_{i-s_{1}}&\text{if }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} (34) Again, we use the fact that the Hadamard product is row-independent in equation 32 , and the definitions of the projections from equation 20 in equation 33 . Overall, we have solved associative recall for all queries that have interaction distance exactly equal to s 1 s_{1} .
In order to generalize this to arbitrary t β€ N 3 t\leq\frac{N}{3} , we first increase the internal dimension so that the input to the kernels π β² β β ( N β
t ) Γ d \bm{u}^{\prime}\in\mathbbm{R}^{(N\cdot t)\times d} in equation 19 and the projections π β² , π β² , π β² β β ( N β
t ) Γ d {\bf K^{\prime},Q^{\prime},V^{\prime}}\in\mathbbm{R}^{(N\cdot t)\times d} are given by π β² β‘ ( π d π d π ) , π² β² β‘ ( π d π d π ) , π β² β‘ ( π π π ) , π β² β‘ ( π d π d π ) , \bm{u}^{\prime}\equiv\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ \bm{u}\end{pmatrix},\bm{K}^{\prime}\equiv\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ {\bf K}\end{pmatrix},{\bf Q}^{\prime}\equiv\begin{pmatrix}{\bf Q}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf Q}\\ \hline\cr\\ {\bf Q}\end{pmatrix},{\bf V}^{\prime}\equiv\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ {\bf V}\end{pmatrix}, We then observe that for π‘ β k β ( X ) := X s β {\bf h}^{k}_{\ell}(X):=X^{s_{\ell}} and π‘ β v β ( X ) := X s β β 1 {\bf h}^{v}_{\ell}(X):=X^{s_{\ell}-1} , we have π‘ k β ( X ) \displaystyle{\bf h}^{k}(X) β‘ β β β [ t ] π‘ β k β ( X ) β
X ( β β 1 ) β
N 3 , \displaystyle\equiv\sum_{\ell\in[t]}{\bf h}^{k}_{\ell}(X)\cdot X^{(\ell-1)\cdot\frac{N}{3}}, π‘ v β ( X ) \displaystyle{\bf h}^{v}(X) β‘ β β β [ t ] π‘ β k β ( X ) β
X ( β β 1 ) β
N 3 . \displaystyle\equiv\sum_{\ell\in[t]}{\bf h}^{k}_{\ell}(X)\cdot X^{(\ell-1)\cdot\frac{N}{3}}. In analogy with equation 25 , we can then equivalently write ( π‘ K β π ) \displaystyle\left({\bf h}^{K}\ast{\bf K}\right) β‘ ( π‘ t K π‘ 2 K π‘ 1 K ) β ( π d π d π ) \displaystyle\equiv\begin{pmatrix}{\bf h}^{K}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf h}^{K}_{2}\\ \hline\cr\\ {\bf h}^{K}_{1}\end{pmatrix}\ast\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ {\bf K}\end{pmatrix} β‘ ( π‘ t K β π π‘ 2 K β π π‘ 1 K β π ) \displaystyle\equiv\begin{pmatrix}{\bf h}^{K}_{t}\ast{\bf K}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf h}^{K}_{2}\ast{\bf K}\\ \hline\cr\\ {\bf h}^{K}_{1}\ast{\bf K}\end{pmatrix} β‘ ( πππππ β _ β πππ π β ( π , s t ) πππππ β _ β πππ π β ( π , s 2 ) πππππ β _ β πππ π β ( π , s 1 ) ) . \displaystyle\equiv\begin{pmatrix}{\tt shift\_down}({\bf K},s_{t})\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf K},s_{2})\\ \hline\cr\\ {\tt shift\_down}({\bf K},s_{1})\end{pmatrix}. Similarly, we also have ( π‘ V β π β² ) \displaystyle\left({\bf h}^{V}\ast{\bf V}^{\prime}\right) β‘ ( πππππ β _ β πππ π β ( π , s t β 1 ) πππππ β _ β πππ π β ( π , s 2 β 1 ) πππππ β _ β πππ π β ( π , s 1 β 1 ) ) . \displaystyle\equiv\begin{pmatrix}{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix}. That is, the argument for t = 1 t=1 now applies to each of the t t shifts as we now have ( cf. equation 21 ) π β² \displaystyle\bm{y}^{\prime} β‘ π β² β ( π‘ V β π ) \displaystyle\equiv{\bf Q}^{\prime}\odot\left({\bf h}^{V}\ast{\bf V}\right) β‘ ( π π π ) β ( πππππ β _ β πππ π β ( π , s t β 1 ) πππππ β _ β πππ π β ( π , s 2 β 1 ) πππππ β _ β πππ π β ( π , s 1 β 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf Q}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf Q}\\ \hline\cr\\ {\bf Q}\end{pmatrix}\odot\begin{pmatrix}{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} β‘ ( π β πππππ β _ β πππ π β ( π , s t β 1 ) π β πππππ β _ β πππ π β ( π , s 2 β 1 ) π β πππππ β _ β πππ π β ( π , s 1 β 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf Q}\odot{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf Q}\odot{\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\bf Q}\odot{\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} β‘ ( π t π 2 π 1 ) , \displaystyle\equiv\begin{pmatrix}\bm{y}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ \bm{y}_{2}\\ \hline\cr\\ \bm{y}_{1}\\ \end{pmatrix}, where, for each β β [ t ] \ell\in[t] , we have ( cf. equation 30 ) π β [ i , : ] β‘ { π j if i β π¬ , i β s β β π¦ and π i β‘ π i β s β β‘ π j , π d otherwise . \bm{y}_{\ell}[i,:]\equiv\begin{cases}{\bf e}_{j}&\text{if }i\in\mathcal{Q},i-s_{\ell}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\equiv{\bf e}_{j},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. We then analogously get π β² {\bf E}^{\prime} as follows: π β² β‘ π»πππππ π β ( π β² ) β‘ ( π»πππππ π β ( π t ) π»πππππ π β ( π 2 ) π»πππππ π β ( π 1 ) ) β‘ ( π t π 2 π 1 ) , {\bf E}^{\prime}\equiv{\tt Linear}_{{\bf E}}(\bm{y}^{\prime})\equiv\begin{pmatrix}{\tt Linear}_{{\bf E}}(\bm{y}_{t})\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt Linear}_{{\bf E}}(\bm{y}_{2})\\ \hline\cr\\ {\tt Linear}_{{\bf E}}(\bm{y}_{1})\end{pmatrix}\equiv\begin{pmatrix}{\bf E}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf E}_{2}\\ \hline\cr\\ {\bf E}_{1}\end{pmatrix}, where, for each β β [ t ] \ell\in[t] , we have ( cf. equation 31 ) π β [ i , : ] = { π d if β i β π¬ , i β s β β π¦ β and β π i β‘ π i β s β π d otherwise {\bf E}_{\ell}[i,:]=\begin{cases}{\bf 1}^{d}&\text{if }i\in\mathcal{Q},i-s_{\ell}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} The output in the general case is then given by π β² \displaystyle\bm{z}^{\prime} β‘ π β² β ( π‘ V β π β² ) \displaystyle\equiv{\bf E}^{\prime}\odot\left({\bf h}^{V}*{\bf V}^{\prime}\right) β‘ ( π t π 2 π 1 ) β ( πππππ β _ β πππ π β ( π , s t β 1 ) πππππ β _ β πππ π β ( π , s 2 β 1 ) πππππ β _ β πππ π β ( π , s 1 β 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf E}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf E}_{2}\\ \hline\cr\\ {\bf E}_{1}\end{pmatrix}\odot\begin{pmatrix}{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} β‘ ( π t β πππππ β _ β πππ π β ( π , s t β 1 ) π 2 β πππππ β _ β πππ π β ( π , s 2 β 1 ) π 1 β πππππ β _ β πππ π β ( π , s 1 β 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf E}_{t}\odot{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf E}_{2}\odot{\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\bf E}_{1}\odot{\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} β‘ ( π t π 2 π 1 ) , \displaystyle\equiv\begin{pmatrix}\bm{z}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ \bm{z}_{2}\\ \hline\cr\\ \bm{z}_{1}\\ \end{pmatrix}, where, for each β β [ t ] \ell\in[t] , we have ( cf. equation 34 ) π β [ i , : ] β‘ { π i β s β if β π i β‘ π i β s β π d otherwise \bm{z}_{\ell}[i,:]\equiv\begin{cases}\bm{v}_{i-s_{\ell}}&\text{if }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} Finally, we define the last output layer to compute π o β u β t β‘ π»πππππ sum β ( π β² ) β‘ β β β [ t ] π β \bm{z}_{out}\equiv{\tt Linear}_{\mathrm{sum}}(\bm{z}^{\prime})\equiv\sum_{\ell\in[t]}\bm{z}_{\ell} so that we have π o β u β t [ i , : ] β‘ { π i β s β if β π i β‘ π i β s β β for some β β β [ t ] π d otherwise \bm{z}_{out}[i,:]\equiv\begin{cases}\bm{v}_{i-s_{\ell}}&\text{if }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\text{ for some }\ell\in[t]\\ {\bf 0}^{d}&\text{otherwise}\end{cases} To recall, we retrieved the top t t interaction distances of the input π \bm{u} using auto-correlation and defined the corresponding convolution kernels (equation 19 ). We then shifted djown the keys π {\bf K} using the first kernel and gated with the corresponding queries π {\bf Q} so that we got a match exactly when there exists a key that is at s β s_{\ell} interaction distance from the corresponding query. After βsmearingβ this match to get π {\bf E} , we used it as a mask to retrieve the value in the next layer. Overall, since we have t t atomic kernels that perform t t shifts with each of these kernels using πͺ β‘ ( N β d ) \mathcal{O}(Nd) parameters, we can conclude that the output solves the associative recall problem for all queries with exactly β \ell interaction distance from the corresponding keys for all β β [ t ] \ell\in[t] using πͺ β‘ ( 1 ) \mathcal{O}(1) layers and πͺ β‘ ( t β
N β c ) \mathcal{O}(t\cdot Nc) parameters as we have d = c d=c . β
## Instructions for reporting errors
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:
Click the "Report Issue" ( ) button, located in the page header.
Tip: You can select the relevant text first, to include it in your report.
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .