Blending Complementary Memory Systems in Hybrid Quadratic-Linear Transformers

TL;DR

This paper introduces hybrid memory architectures combining key-value and fast weight memories in transformers, enhancing sequence processing by leveraging their complementary strengths for language modeling, algorithmic, and reinforcement learning tasks.

Contribution

The paper proposes and compares three methods to effectively blend quadratic and linear transformer memory systems, demonstrating improved performance across multiple tasks.

Findings

Hybrid memory systems outperform individual components in language modeling.

Hybrid approaches enable processing of longer sequences with better recall.

Experimental results show improved performance in reinforcement learning environments.

Abstract

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.