Birdie: Advancing State Space Language Modeling with Dynamic Mixtures of Training Objectives (2024)

Abstract: Efficient state space models (SSMs), including linear recurrent neural networks and linear attention variants, have emerged as potential alternative language models to Transformers. While efficient, SSMs struggle with tasks requiring in-context retrieval, such as text copying and associative recall, limiting their usefulness in practical settings. Prior work on how to meet this challenge has focused on the internal model architecture and not investigated the role of the training procedure. This paper proposes a new training procedure that improve the performance of SSMs on retrieval-intensive tasks. This novel pre-training procedure combines a bidirectional processing of the input with dynamic mixtures of pre-training objectives to improve the utilization of the SSM’s fixed-size state. Our experimental evaluations show that this procedure significantly improves performance on retrieval-intensive tasks that challenge current SSMs, such as phone book lookup, long paragraph question-answering, and infilling tasks. Our findings offer insights into a new direction to advance the training of SSMs to close the performance gap with Transformers.

## introduction due to their scaling properties @xcite and in-context learning @xcite , large transformer models using attention @xcite are now prominent in natural language processing (nlp) and achieve effective performance in natural language generation tasks (nlg), including language modeling, machine translation, and question and answering (q&a) @xcite @xcite . however, the softmax attention mechanism cost scales quadratically with sequence length during training, and its key-value (kv) cache grows linearly with sequence length during inference. this leads to increasing costs for training and deployment as model providers continue to increase the context length @xcite . this trend in increasing context length has sparked a strong interest in developing efficient alternative sequence models. the goal is to maintain high performance while scaling effectively with longer sequences. recent work has focused on recurrent models which offer two key advantages: subquadratic scaling for parallel processing and a fixed state size (in contrast to the growing kv cache in transformer models) that enables constantcost inference per step. these models come in different forms, ranging from state space model (ssm)-based methods, such as s4 @xcite , @xcite , or mamba @xcite ), to linear rnns, such as rwkv @xcite , hgru @xcite , and hawk @xcite , to linear attention variants, such as retnet @xcite and gla @xcite . these different methods vary in their exact parameterization and parallel computation, but all have an efficient, fixed-state size recurrence for inference. for brevity, we will generally refer to all of these methods as ssms regardless of their exact parameterization or parallel computation path. while some studies have shown the ability of ssms to match transformers in perplexity and some public benchmarks, an increasing line of work shows that current ssms struggle on tasks that require long-range in-context abilities @xcite , such as long-range retrieval @xcite , multi-query associative recall @xcite , and copying @xcite . these tasks are critical in nlp, where the ability to maintain and manipulate long-term dependencies is key to generating coherent text, following directions, copying sequences, and responding accurately to multiple queries. a typical approach to address these weaknesses has been to formulate hybrid models @xcite that interleave ssm layers with global attention layers @xcite @xcite , or sliding window attention @xcite @xcite . 2 however, models with global attention layers still scale quadratically with sequence length and have a growing kv cache. models that rely on sliding window attention also fail to perform in-context retrieval outside of the sliding window length @xcite . the predominant focus on architecture to improve performance on long-range in-context abilities misses an opportunity to investigate the role of the pre-training objective(s) and potential interaction with model architecture. we note that prior work on ssms predominantly utilizes a causal language modeling (clm) pre-training objective. in this paper we argue (and show) that in the presence of a fixed state size, a mixture of pre-training objectives can bias learning towards pertinent longrange interactions and that bidirectional processing of the context allows better utilization of the fixed state for such interactions. this paper makes the following key methodological contributions: the result is a new training procedure that significantly improves the performance of ssms on recall-intensive tasks, making them more competitive with transformers. we refer to this proce-dure as birdie. while we do still observe a performance gap with transformers on some tasks as the retrieval requirement becomes more difficult (e.g. increasing the number of retrievals required per example), our procedure makes the ssm performance degradation in these scenarios much less severe and expands the regime where these efficient methods can be useful. more broadly, our work points to considering the learning dynamics along with the inductive biases of ssm architectures in order to make better use of the fixed state size. ## background and related work this section relates background and prior work. ## methods we propose two key methodological components to reduce the gap between ssms and transformers on in-context retrieval tasks: bidirectional processing of the input prompt or prefix and new mixtures of pre-training objectives designed to improve the ability of ssms to perform retrieval. we then offer a new pre-training procedure that leverages rl for dynamic sampling of the pre-training objectives to reduce the burden of pre-selecting the optimal mixture ahead of time. we combine these components to define the birdie training procedure. in the final part of this section, we also describe a baseline gated ssm that allows for a simple implementation to test our methods. ## experiments and results here, we present our experimental setup and main findings. ## conclusion in this work, we investigated the significant impact of the training procedure on the downstream capabilities of state space models (ssms). while prior research highlighted major weaknesses of ssms on in-context retrieval tasks, we demonstrated that model training procedure accuracy instruct models gated ssm+ birdie 42.5% transformer birdie 42.0% transformer clm 42.0% hawk birdie 40.4% hawk clm 34.0% gated ssm+ clm 32.8% base models transformer clm 40.4% hawk birdie 39.7% transformer birdie 39.7% gated ssm+ birdie 36.6% hawk clm 29.6% gated ssm+ clm 29.5%