Language Modeling With Factorization Memory
This work addresses the computational efficiency and scalability challenges in language modeling for both short and long contexts, representing an incremental improvement over existing RNN and Transformer approaches.
The authors tackled the problem of efficient language modeling across both short and long contexts by proposing Factorization Memory, an RNN architecture that achieves performance comparable to Transformers on short-context tasks while demonstrating superior generalization in long-context scenarios.
We propose Factorization Memory, an efficient recurrent neural network (RNN) architecture that achieves performance comparable to Transformer models on short-context language modeling tasks while also demonstrating superior generalization in long-context scenarios. Our model builds upon Mamba-2, enabling Factorization Memory to exploit parallel computations during training while preserving constant computational and memory complexity during inference. To further optimize model efficiency and representational capacity, we develop a sparse formulation of Factorization Memory that updates only a subset of recurrent states at each step while preserving the strong performance of its dense counterpart. To our knowledge, this represents the first RNN architecture that successfully combines sparse memory activation with competitive performance across both short and long-context settings. This work provides a systematic empirical analysis of Factorization Memory in comparison to Transformer and Mamba-2 architectures.