Language Modeling With Factorization Memory

TL;DR

This work introduces Factorization Memory, a sparse-update RNN with a 2D memory that matches Transformer performance on short-context language modeling and outperforms it on long-context extrapolation. By combining dense and top-k sparse memory updates, it achieves scalable computation with maintained capacity, and it provides optimized CUDA/Triton kernels to facilitate reproducible research. Across empirical studies, Factorization Memory demonstrates competitive results on English/Japanese benchmarks and superior inference speed compared to Transformer and Mamba-2, particularly on long contexts. The approach advances efficient, memory-augmented sequence modeling, offering a compelling alternative to attention-based architectures for long-context tasks.

Abstract

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.

Figures (7)

• Figure 1: Factorization Memory - Layer Schematics. Left: In the dense formulation all $m$ memory states are updated at each timestamp. The updates are weighted with memory-input affinity scores, and the thickness of the arrows represents the strength of the update. Right: In the sparse formulation, only selected top-$k$ states are updated at each timestamp. Grey shading indicates that the state is used neither in update nor in merge operations.
• Figure 2: Loss Frontier: All models are trained with the context length of 1024 tokens, while varying the number of model parameters, learning rate, and training budget.
• Figure 3: Loss Frontier Scaling
• Figure 4: Loss-So-Far. Left: English. Right: Japanese. Factorization Memory consistently achieves performance comparable with Transformer and Mamba-2 at the training context length of $2^10 = 1024$ tokens. At the same time Factorization Memory shows better extrapolation at the long context. This pattern holds for all tested model sizes (see Appendix \ref{['sec:app:longcontext']}).
• Figure 5: Comparison of dense and sparse memory updates: test loss generally decreases with an increasing number of memory states, even for sparse updates. Notably, updating only 25% of memory states achieves the same loss as dense formulation when the number of memory states is sufficiently large while reducing computational cost by 75%.