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

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

Claim 1 (Gated-convolutions and attention). Gated-convolution models with two layers require model dimension to scale at least linearly in sequence length in order to solve associative recall, while attention models can solve it with near-constant dimensionality. We compare attention and BaseConv as well as three popular instantiations of gated-convolution architectures: RWKV, H3, and Hyena ( Peng et al., 2023 ; Fu et al., 2023c ; Poli et al., 2023a ) . In the top row of Fig. 2 , attention solves Mqar perfectly at all sequence lengths using a constant model dimension of 64. In contrast, Mqar does not achieve accuracy > 0.9 >0.9 unless d ≥ N d\geq N .

Claim 2 (Input-dependent filters). Using input-dependent filters in gated-convolution models can close some of the gap to attention. In Theorem 4.5 , we show that BaseConv with input-dependent filters could solve Mqar with improved scaling. In this solution, we construct a filter that spikes at position j j if matching keys are separated by j j tokens. We evaluate two approaches for constructing this filter: (1) programatically ( i.e. hard-coded comparisons between token ids) or (2) with autocorrelation, which could learn to perform fuzzy-matches (see Section H.8 ). In the bottom row of Fig. 2 , we see that BaseConv with programmatic input-dependent filters achieves near-constant scaling in model dimension and that BaseConv with autocorrelation input-dependent filters achieves improved scaling over BaseConv with input-independent filters.

These input-dependent filters cannot easily be made to satisfy causality and using an O ⁡ ( N ​ log ⁡ N ) O(N\log N) filter per gap could be expensive if each gap applies only to a small number of bigrams. A simpler and perhaps more efficient way to solve the problem would be to introduce a small amount of attention to an otherwise BaseConv model ( Fu et al., 2023a ) . As the first natural baseline, we evaluate an attention hybrid on synthetic Mqar and show that it achieves improved scaling in Figure 2 . Next, we put these empirical and theoretical insights into practice on Pile language modeling.

## 5 Closing the Associative Recall Gap

In this section, we evaluate hybrid BaseConv -Attention models that leverage different sparsity patterns. We show that hybrids with input-dependent sparsity patterns can close most of the gap to attention, while maintaining sub-quadratic scaling. We also show that BaseConv hybrids can outperform attention-only models by up to a full perplexity point, all while being dramatically simpler to implement and analyze than prior hybrids ( Fu et al., 2023a ) (see implementation in Appendix B . We describe the architectures in Section 5 and results on Pile language modeling in Section 5 .

##### Sparse BaseConv -Attention Hybrids

We evaluate hybrids composed primarily of BaseConv layers and three attention layers (6.3% of layers at 354M parameters and 10% at 168M parameters). We augment the attention layers with operators that selectively apply attention to some tokens based on a selection function f : ℝ N × d → { 0 , 1 } N f:\mathbb{R}^{N\times d}\rightarrow\{0,1\}^{N} . They take as input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbb{R}^{N\times d} and output 𝒚 ∈ ℝ N × d \bm{y}\in\mathbb{R}^{N\times d} : 𝒚 [ i , : ] = softmax ( 1 d 𝒒 [ i , : ] 𝒌 ⊤ ) 𝒗 ⋅ f ( 𝒖 ) [ i ] \bm{y}[i,:]=\mathrm{softmax}(\frac{1}{\sqrt{d}}\bm{q}[i,:]\bm{k}^{\top})\bm{v}\cdot f(\bm{u})[i] (2)

where 𝒒 , 𝒌 , 𝒗 \bm{q},\bm{k},\bm{v} are query, key, and value projections, as in attention. We evaluate four choices for f f :

(1) Full attention. First, we evaluate the performance of full attention hybrids. These hybrids correspond to fixing f ​ ( 𝒖 ) ​ [ i ] = 1 f(\bm{u})[i]=1 for all i i in Eq. 2 . Prior work has shown full attention hybrids to be effective, but they do not explain why supplementing gated convolutions with attention is necessary ( Fu et al., 2023a ) . Based on our findings in Section 3 , we hypothesize that sparse attention applied only to associative recall hits may suffice.

(2) Random selection. As a control, we evaluate sparse attention randomly applied to tokens. This corresponds to a stochastic selection function where f ​ ( 𝒖 ) ​ [ i ] f(\bm{u})[i] is drawn from a Bernoulli. Many prior works have employed sparse attention patterns like this one, which are independent of the input ( Zaheer et al., 2020 ; Child et al., 2019a ; Beltagy et al., 2020 , inter alia.) .

(3) Programmatic selection. Next, to evaluate our hypothesis that attention is needed for associative recall, we prototype a programmatic selection function that selects only those tokens that might be associative recall hits. Specifically, f ( 𝒙 [ i , : ] ) f(\bm{x}[i,:]) is 1 1 if the token x i x_{i} previously occurred in the sequence. In practice, we compare raw token ids, not token embeddings. f ​ ( 𝒙 ) ​ [ i ] = { 1 if there exists j < i such that x i = x j 0 otherwise f(\bm{x})[i]=\begin{cases}1&\text{if there exists}\quad j<i\quad\text{such that}\quad x_{i}=x_{j}\\ 0&\text{otherwise }\end{cases} (3)

(4) Learned selection. Finally, we prototype a learned selection function f ( 𝒖 ) [ i ] = σ ( 𝒖 [ i , : ] ⋅ 𝑾 ) f(\bm{u})[i]=\sigma(\bm{u}[i,:]\cdot\bm{W}) parameterized as a simple linear layer with sigmoid activation. We fix a hyperparameter k k and select the top- k k tokens with the highest score in each batch. This allows us to compute attention in 𝒪 ⁡ ( n ​ d ​ k ) \mathcal{O}(ndk) time. During training, we add small amount of Gaussian noise to the top- k k calculation to encourage exploration and use an auxiliary loss to encourage sparse selection: ℓ f ​ ( 𝐮 ) = 1 N ​ max ⁡ ( 0 , ∑ i = 1 N f ⁡ ( 𝒖 ) ​ [ i ] − k ) \ell_{f}(\mathbf{u})=\frac{1}{N}\max(0,\sum_{i=1}^{N}f(\bm{u})[i]-k) . This approach is most similar to the recently proposed SeqBoat architecture ( Ren et al., 2023 ) . We discuss the differences in Appendix A .

##### Downstream Evaluations

We evaluate the prototypes on Pile language modeling. We take BaseConv architectures at the 150M and 360M parameter scales and add input-dependent selection to three layers (details in Section C.2 ). Results are in Table 2 . We validate that the prototypes close the overall and AR quality gaps to attention when added to the BaseConv backbone.

At 360M parameters, BaseConv with just 3 3 attention layers can outperform the Transformer, while requiring fewer FLOPs. Our hybrid attention- BaseConv models outperform attention only models by 0.85 perplexity points while enabling an 18% reduction in total FLOPs vs attention. However, this uses full quadratic attention. We next show that sparse attention localized to potential AR tokens, is also sufficient to close the gap, validating our insights on the role of input-dependence for MQAR. At 360M parameters, programmatic selection closes 85% of the gap between pure BaseConv and attention on the AR slice, in contrast to the random selection control. Learned selection closes 72% a of the gap using just k = 256 k=256 (sub-quadratic) attention positions per example.

## 6 Discussion and Conclusion

We present an extensive analysis of gated convolution architectures in light of their recent popularity. We identify a persisting quality gap between efficient convolution and inefficient attention based architectures, largely due to a single failure mode associative recall. We design a new multi-query associative recall ( Mqar ) analysis tool, which correlates with downstream AR quality. We theoretically and empirically explain the gap is due to insufficient data-dependent mixing in gated convolutions and we show minimal architectures that close the gap on the Pile.

Our results in analyzing gated convolutions go beyond the conventional wisdom that attention is the “ right ” model. A significant amount of work focuses on improving the efficiency of attention ( Dao et al., 2022 ; Dao, 2023 ; Katharopoulos et al., 2020a ) and theoretically studying the exact power of attention ( Hahn, 2020 ; Merrill et al., 2022 ; Keles et al., 2023 ) . Attention is often used as the goalpost for what is needed downstream. We hope our contributions highlight the value of Mqar , and more broadly tasks tied to real language modeling, as a proxy to study.

## Acknowledgments

We thank Tri Dao, Daniel Fu, Neel Guha, Stefano Massaroli, Eric Nguyen, and Michael Zhang for helpful feedback and discussion during this work. We are grateful to Together Computer for making this work possible. We gratefully acknowledge the support of DARPA under Nos. FA86501827865 (SDH) and FA86501827882 (ASED); NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under No. N000141712266 (Unifying Weak Supervision); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Microsoft Azure, Swiss Re, Brown Institute for Media Innovation, Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fellowship Program, Texas Instruments Stanford Graduate Fellowship in Science and Engineering, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, and Infosys. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ONR, or the U.S. Government. AR’s work is supported by NSF grant# CCF-2247014. IJ’s work is supported by an NSF Graduate Fellowship.

## Appendix

The appendix includes the following content: 1. Appendix A provides an extended discussion of related work and concepts.

2. Appendix B provides a code implementation of the BaseConv architecture.

3. Appendix C gives details for the experiments, including model architectures and hyperparameters.

4. Appendix D provides additional analysis of how Mqar appears in real data across a variety of language distributions.

5. Appendix E provides a formal definition of the Mqar problem, synthetic construction procedure, and experimental details.

6. Appendix F provides additional synthetic experiments and analysis.

7. Appendix G provides experiments and analysis for how the Mqar gap changes as we scale the gated convolution and attention architectures.s

8. Appendix H gives proofs and additional discussion for the theoretical analysis in our work.

## Appendix A Extended Related Work

We are inspired by and build on prior work in mechanistic interpretability ( Section A.1 ), efficient architecture design ( Section A.2 ), and input-dependent architectures ( Section A.3 ).

### A.1 Mechanistic Interpretability, Synthetic Languages, and Associative Recall

Work in mechanistic interpretability aims to decompose the capabilities of a neural network into human-understandable algorithms that can be attributed to specific parameters in the model [ Olah, 2022 , Power et al., 2022 , Elhage et al., 2021 , Cammarata et al., 2020 ] . Some of these works, use synthetic data to validate mechanistic interpretations of neural networks [ Olsson et al., 2022 ] . This relates to a broader line of work using synthetic languages to study language model architectures [ Wang and Eisner, 2016 , White and Cotterell, 2021 , Allen-Zhu and Li, 2023 , Ravfogel et al., 2019 , Xie et al., 2021 ] . Mechanistic design puts mechanistic interpretations to use in designing new architectures and learning algorithms. Several works in architecture research have used synthetic tasks to validate architecture designs [ Kitaev et al., 2020 ] .

Our work is focused on one particular synthetic task: associative recall . Motivated by psychological models of how humans associate and retrieve information, work in the early days of neural network research focused on developing systems capable of associative recall [ Willshaw et al., 1969 , Feldman et al., 1981 , Hopfield, 1982 ] . For example, Hopfield networks, proposed in 1982, are a recurrent neural network explicitly designed to support associative (“content-addressable”) memory [ Hopfield, 1982 ] . More recently, several notable recurrent neural network mechanisms ( e.g. neural turing machines, LSTMs, and RNNs with fast weights) were evaluated on a synthetic version of associative recall [ Graves et al., 2014 , Ba et al., 2016 , Zhang and Zhou, 2017 , inter alia.] . These works use a formulation of associative recall very similar to the single-query associative recall in our work. Since the rise of large language models, several works have argued that LLMs ability to perform in-context learning is due, at least in part, to the associative recall capabilities of attention [ Elhage et al., 2021 , Olsson et al., 2022 ] .

### A.2 Efficient Language Modeling Architectures

We first briefly review the efficiency motivations for the recent excitement around gated convolution architectures. While attention requires compute that scales as 𝒪 ⁡ ( N 2 ) \mathcal{O}(N^{2}) in sequence length N N , convolutions scale as 𝒪 ⁡ ( N ​ log ⁡ N ) \mathcal{O}(N\log N) [ Cooley and Tukey, 1965 ] . Ideal, inference complexity is 𝒪 ⁡ ( 1 ) \mathcal{O}(1) in sequence length, as provided by recurrent neural networks. State-space models can be computed either as a convolution or recurrence, to achieve both sub-quadratic training and constant inference complexity in sequence length [ Gu et al., 2021 ] . Architectures that use implicit convolutional filters [ Poli et al., 2023a ] , can be converted to an SSM via a simple distillation step [ Massaroli et al., 2023 ] .

### A.3 Input-Dependence in Sequence Modeling Architectures

In an input-dependent sequence model, the way tokens are aggregated across the sequence is controlled by the data, not just the model parameters. We highlight several prior works related to our study that explore input-dependent sequence models. • Attention [ Vaswani et al., 2017 ] achieves input-dependent sequence mixng since the Q Q , K K , and V V terms are achieved via linear combinations of the input x x . For instance, Q = x ​ W Q Q=xW_{Q} for some learned weight matrix W Q W_{Q} .

• Input-Independent Convolution Architectures . Given the quadratic scaling of attention, Gu et al. [2021] propose S4, an architecture that use long convolutions to process the input. CKConv Romero et al. [2022] uses a convolution filter that is implicitly parametrized by an MLP of the relative positions of the observations in a time-series. A line of subsequent work improved upon S4 by changing the parametrization [ Gupta et al., 2022 , Gu et al., 2022 , Mehta et al., 2022 , Ma et al., 2022 , Fu et al., 2023c ] or changing the SISO convolutions to MIMO convolutions [ Smith et al., 2023 ] . The convolution filter in each of these architectures is input-independent. The linear convolution layer alone cannot perform MQAR [ Fu et al., 2023a ] .