Research map / Architectures and efficiency / Attention and transformers
MemoryFormer : Minimize Transformer Computation by Removing Fully-Connected Layers
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.
Proposes MemoryFormer, a transformer that replaces all fully-connected layers with LSH-based in-memory lookup tables (Memory Layers), trained from scratch on the Pile to cut FLOPs while keeping zero-shot accuracy.
Problem
Transformer FC layers dominate compute in typical sequence lengths, and existing efficiency work mostly targets attention; the paper seeks a cheaper replacement for linear projections.
Approach
Replace each linear projection with Memory Layers: sign-based LSH hashes split sub-vectors into indices into learnable lookup tables, whose retrieved entries are weighted by a differentiable sign-similarity probability so gradients reach both tables and inputs.
Key results
- MemoryFormer-base reaches 0.458 average zero-shot accuracy vs 0.435 for Pythia-410M, with 10.2 G vs 34.4 G total FLOPs per block at sequence length 2048.
- MemoryFormer-tiny averages 0.404 vs 0.375 for Pythia-70M at 4.7 G vs 10.7 G FLOPs.
- Against efficient attention variants on Pythia-410M base, MemoryFormer (0.458) beats Linformer, cosFormer and Performer (0.398-0.418).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace fully-connected layers with Memory Layers | Pythia-410M (MemoryFormer-base) · PIQA | accuracy: 0.675 -> 0.698 | Helped |
| Replace fully-connected layers with Memory Layers | Pythia-410M (MemoryFormer-base) · WinoGrande | accuracy: 0.534 -> 0.546 | Helped |
| Replace fully-connected layers with Memory Layers | Pythia-410M (MemoryFormer-base) · WSC | accuracy: 0.471 -> 0.385 | Hurt |
| Replace fully-connected layers with Memory Layers | Pythia-410M (MemoryFormer-base) · ARC-E, ARC-C, and LogiQA | accuracy | Helped |
| Use wider hash chunks (τ=10, K=51 rather than τ=4, K=128) (ablation) | MemoryFormer-tiny · Validation set | perplexity: 19.01 -> 18.67 | Helped |
| Increase learning rate from 1e-3 to 3e-3 (ablation) | MemoryFormer-tiny · Validation set at 8000 training steps | perplexity: 19.86 -> 18.82 | Helped |
| Increase learning rate from 3e-3 to 4e-3 (ablation) | MemoryFormer-tiny · Validation set at 8000 training steps | perplexity: 18.82 -> 18.84 | Hurt |
| Add expanding bits to the Memory Block (0 to 3) (ablation) | MemoryFormer-tiny · Validation set at 8000 training steps | perplexity: 19.89 -> 18.54 | Helped |
| Add GeLU between the two Memory Layers (ablation) | MemoryFormer-tiny · PIQA | accuracy: 0.602 -> 0.595 | Hurt |
| Add GeLU between the two Memory Layers (ablation) | MemoryFormer-tiny · WinoGrande and WSC | accuracy: 0.522 -> 0.522; 0.375 -> 0.375 | No clear effect |
| Add GeLU between the two Memory Layers (ablation) | MemoryFormer-tiny · ARC-E | accuracy: 0.437 -> 0.441 | Helped |
| Add GeLU between the two Memory Layers (ablation) | MemoryFormer-tiny · ARC-C | accuracy: 0.228 -> 0.211 | Hurt |
Surprising findings
- Adding a GeLU between the two Memory Layers made no meaningful difference (avg 0.402 vs 0.404), supporting the claim that the hashing step is already nonlinear.
- Scaling the learning rate to 3x baseline was needed for best perplexity, attributed to sparse hash-table gradients.
Limitations
- Evaluated only on small models (up to 410M) with zero-shot multiple-choice benchmarks; no scaling to LLM sizes.
- Reports FLOPs rather than measured latency or memory cost; lookup tables grow exponentially with tau.
- Benchmark results are noisy (e.g. WSC and ARC-C show non-monotonic or near-chance values).
Experimental setup
Datasets: The Pile
Benchmarks: PIQA WinoGrande WSC ARC-E ARC-C LogiQA
Metrics: zero-shot accuracy average accuracy inference FLOPs validation perplexity
Baselines: Pythia-70M Pythia-160M Pythia-410M Linformer cosFormer Performer
Builds on: Pythia Reformer LookupFFN Linformer Performer
Closest papers by meaning
- Memory-Augmenting Decoder-Only Language Models through Encoders (Student Abstract) (AAAI 2024 · similarity 0.89)
- Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads (NeurIPS 2025 · similarity 0.88)
- Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.88)
- PolySketchFormer: Fast Transformers via Sketching Polynomial Kernels (ICML 2024 · similarity 0.88)
- DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging (NeurIPS 2024 · similarity 0.87)
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.