Blending Complementary Memory Systems in Hybrid Quadratic-Linear Transformers
This work addresses the design of neural memory systems for general-purpose sequence processing, offering incremental improvements by blending existing transformer architectures.
The paper tackled the problem of combining two complementary memory systems in transformers—KV-memory for precise retrieval and FW-memory for long sequences—by proposing three hybrid methods, demonstrating improved performance on language modeling, retrieval, and reinforcement learning tasks with models up to 1.3B parameters.
We develop hybrid memory architectures for general-purpose sequence processing neural networks, that combine key-value memory using softmax attention (KV-memory) with fast weight memory through dynamic synaptic modulation (FW-memory) -- the core principles of quadratic and linear transformers, respectively. These two memory systems have complementary but individually limited properties: KV-memory offers precise retrieval but is constrained by quadratic complexity in sequence length, while FW-memory supports arbitrarily long sequences and enables more expressive computation but sacrifices precise recall. We propose and compare three methods to blend these two systems into a single memory system, differing in how and when input information is delivered to each system, to leverage the strengths of both. We conduct experiments on general language modeling and retrieval tasks by training 340M- and 1.3B-parameter models from scratch, as well as on synthetic algorithmic tasks designed to precisely illustrate the benefits of certain hybrid methods over others. We also evaluate our hybrid memory systems on reinforcement learning in partially observable environments. Overall, we demonstrate how a well-designed hybrid can overcome the limitations of its individual components, offering new insights into the design principle of neural memory systems.