Memory Augmented Language Models through Mixture of Word Experts (2024)

Abstract: Scaling up the number of parameters of language models has proven to be an effective approach to improve performance. For dense models, increasing their size proportionally increases their computational footprint. In this work, we seek to aggressively decouple learning capacity and FLOPs through Mixture-of-Experts (MoE) style models with large knowledge-rich vocabulary based routing functions. Our proposed approach, dubbed Mixture of Word Experts (MoWE), can be seen as a memory augmented model, where a large set of word-specific experts play the role of a sparse memory. We demonstrate that MoWE performs significantly better than the T5 family of models with similar number of FLOPs in a variety of NLP tasks. Moreover, MoWE outperforms traditional MoE models on knowledge intensive tasks and has similar performance to complex memory augmented approaches that often require to invoke custom mechanisms to search the sparse memory.

## introduction increasing the parameter count of language models has been a primary driver of increased model quality @xcite @xcite . this is particularly apparent on knowledge intensive tasks, such as triviaqa @xcite , where language models with more parameters and learning capacity benefit from soaking up world knowledge from their pretraining data @xcite . however, increasing the model size also increases the cost of training and serving the model. in this work, we build on the mixture-of-experts (moe) paradigm to design a neural net architecture that enjoys the quality benefits from scaling the parameter count but remains flops and latency efficient. our proposed approach, which we name mixture-of-word-experts (mowe), follows two design principles: (1) a very large numand t5 .1.1-xxl, respectively, while using a significantly smaller number of flops. t5.1.1 results are from @xcite . we empirically demonstrate that mowe significantly outperforms t5 models @xcite with a comparable number of flops across a variety of nlp tasks. focusing on knowledge intensive tasks such as triviaqa @xcite and @xcite , we show that a mowe "base" sized outperforms t5-xl and a mowe "large" outperforms t5-xxl models (see figure 1 ), while being at least 4.3x and 6.6x faster to train, respectively. mowe outperforms vanilla moe models @xcite figure 2 : mowe layer: we replace the ffn layer in a subset of transformer blocks by a mowe layer, which is a sparse layer that processes tokens using multiple experts (ffns). each input token is processed by a single expert that is selected based on the input token id (at the corresponding sequence position) in the routing vocabulary. 2020; @xcite on knowledge intensive tasks, while matching performance on nlp task suites such as superglue @xcite . additionally, mowe also matches or outperforms recently proposed knowledge augmented models @xcite , while avoiding invoking any custom mechanism to search the sparse memory. the main contributions of this work are: ## conclusions we presented mowe, a novel neural net architecture effectively bridging the efficiency of sparse moe models with the knowledge retrieval capabilities of memory-augmented models. by leveraging a large, knowledge-rich vocabulary for routing and employing tens of thousands of wordspecific experts, mowe demonstrates significant performance gains on knowledge-intensive tasks. our findings highlight the effectiveness of lexicaldriven routing and the potential of word experts as a form of sparse, integrated memory within language models, opening doors for future research. future investigations could analyze the properties of word experts, explore integration within decoder-only models, and extend mowe to other modalities and languages.