Research map / Architectures and efficiency / RNNs and normalization

Memory Layers at Scale

ICML 2025 · Architectures and efficiency › RNNs and normalization · approach: Efficient network compression and training · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Scales trainable product-key memory layers (up to 128B memory parameters) that replace FFN layers in Llama-style transformers, showing large factual QA gains over dense, MoE, and PEER baselines at matched compute.

Problem

Memory layers (trainable key-value lookups that add parameters without FLOPs) have not been scaled enough to be useful in modern architectures; they are memory-bandwidth bound and hard to parallelize, and their benefits for factual knowledge at scale are unknown.

Approach

Make memory layers practical at scale via trainable product keys, a parallel sharded EmbeddingBag over GPUs, custom CUDA kernels, a single shared memory pool across multiple (3) layers replacing FFNs, and a gated silu-based "Memory+" variant.

Key results

• At 1.3B base size, Memory+ (3.377B total params) reaches 13.68 NQ and 42.89 TQA versus 7.76 and 32.64 for the dense model, and 12.33 / 42.46 for PEER.
• A 1.3B Memory+ model with 64M keys (138.7B total params) reaches 20.78 NQ and 62.14 TQA, approaching Llama2 7B (25.10 NQ, 64.00 TQA) trained on 2T tokens with about 10x more FLOPs.
• At 8B base scale with 16M memory values, Memory+ at 1T tokens scores 68.15 TQA and 63.04 MMLU versus 63.62 and 59.68 for the dense model at 1T tokens; the custom kernels make the embedding bag operation 6x faster than PyTorch's.

Measured effects reported in the paper

Surprising findings

• Sparse and dense layers are complementary: replacing more than 3 FFN layers with shared memory layers degrades performance.
• Memory+ at 1.3B with 1M values already matches or beats Llama2 7B-level factual QA gains are not reached; the gain over dense is larger on factual QA than on general reasoning tasks such as PIQA.

Limitations

• Memory layer implementation still requires substantial engineering to be efficient for production; not yet shown faster than dense FFNs.
• Evaluation focuses on short-form QA and likelihood; long-form generation is left for future work.
• Key sampling improvements (random keys, softmax sink) were excluded due to training speed cost and inconsistent gains at larger sizes.

Experimental setup

Datasets: Llama2-like pretraining data mix Llama3-like pretraining data mix

Benchmarks: NaturalQuestions TriviaQA HotpotQA MMLU HellaSwag OBQA PIQA HumanEval MBPP

Metrics: exact match F1 accuracy pass@1 negative log-likelihood

Baselines: dense Llama-style transformer mixture-of-experts PEER Llama2 7B Llama3.1 8B

Builds on: product-key memory networks PEER mixture-of-experts expert choice Llama

Closest papers by meaning

• Attention-Level Speculation (ICML 2025 · similarity 0.87)
• Memory Grafting: Scaling Language Model Pre-training via Offline Conditional Memory (NeurIPS 2026 · similarity 0.87)
• FFN Fusion: Rethinking Sequential Computation in Large Language Models (NeurIPS 2025 · similarity 0.87)
• Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers (NeurIPS 2025 · similarity 0.87)
• Scaling Embedding Layers in Language Models (NeurIPS 2025 · similarity 0.87)
• Scaling Inference-Efficient Language Models (ICML 2025 · similarity 0.87)