Research map / Architectures and efficiency / Mixture of experts
Ultra-Sparse Memory Network
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.
Introduces UltraMem, a large-scale ultra-sparse memory layer architecture (improved product-key memory with Tucker-decomposed retrieval and implicit value expansion) that matches or beats MoE at equal compute and parameters while being up to 6x faster at inference.
Problem
MoE models decouple parameters from compute but incur high memory-access costs at inference, while prior product-key memory (PKM) layers underperform MoE and do not scale well.
Approach
Builds UltraMem by improving PKM with a bag of training/structural tricks, Tucker-decomposed query-key retrieval (with rank-1 approximated top-m), implicit value expansion (virtual memory via on-demand reparameterization), multi-core scoring, distributed memory layers across transformer depth, and a new initialization, so that only a few values are accessed per token and memory access stays near d
Key results
- UltraMem-1.6B-x12 (21.4B params, 3.50 GFLOPs) reaches average score 48.26 vs MoE-1.6B-2in34 at 45.07 and C4 validation loss 2.24 vs 2.30 (10-benchmark table).
- UltraMem achieves up to 6x faster inference than MoE at batch size 64 with the same parameters and compute on A100 hardware.
- UltraMem-1.6B-x12 matches the performance of the 6.5B dense model, per the paper's claim.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use UltraMem instead of MoE at 680M activated-model scale | UltraMem-680M-x12 vs MoE-680M-2in33 · C4 validation | Avg: 38.43 -> 42.27 | Helped |
| use UltraMem instead of MoE at 1.6B activated-model scale | UltraMem-1.6B-x12 vs MoE-1.6B-2in34 · C4 validation | Avg: 45.07 -> 48.26 | Helped |
| remove score Softmax (ablation) | PKM-151M-x10 · C4 validation | loss: 2.828 -> 2.822 | Helped |
| halve value dimension, double value count, and add output projection (ablation) | PKM-151M-x10 · C4 validation | loss: 2.822 -> 2.800 | Helped |
| share query across two key sets (ablation) | PKM-151M-x10 · C4 validation | loss: 2.800 -> 2.803 | Hurt |
| split large memory into smaller skip-layer memory units (ablation) | PKM-151M-x10 · C4 validation | loss: 2.803 -> 2.788 | Helped |
| add query/key Layer Normalization (ablation) | PKM-151M-x10 · C4 validation | loss: 2.788 -> 2.789 | No clear effect |
| add Implicit Value Expansion (ablation) | PKM-151M-x10 · C4 validation | loss: 2.789 -> 2.772 | Helped |
| replace product-key retrieval with Tucker-decomposed retrieval (ablation) | PKM-151M-x10 · C4 validation | loss: 2.772 -> 2.764 | Helped |
| add Multi-Core Scoring (ablation) | PKM-151M-x10 · C4 validation | loss: 2.764 -> 2.761 | Helped |
| use improved value initialization (ablation) | PKM-151M-x10 · C4 validation | loss: 2.761 -> 2.758 | Helped |
| linearly decay value-parameter learning rate (ablation) | PKM-151M-x10 · C4 validation | loss: 2.758 -> 2.736 | Helped |
Surprising findings
- Inference time of UltraMem stays nearly constant as sparse parameters grow exponentially, while MoE's rises significantly.
Limitations
- Ablations and largest models are based on 151M and 1.6B dense backbones; the paper's own tables show UltraMem not always best (e.g., MoE-151M leads the 151M average in the appendix table).
- Approximated top-m via rank-1 Tucker core introduces approximation error that requires an auxiliary loss to control.
Experimental setup
Datasets: RedPajama C4 validation
Benchmarks: MMLU TriviaQA GPQA ARC BBH BoolQ HellaSwag WinoGrande DROP AGIEval
Metrics: accuracy C4 validation loss inference time memory access
Baselines: Dense transformer (151M, 680M, 1.6B, 6.5B) MoE (fine-grained, 2in32/2in33/2in34) PKM-x12
Builds on: Product Key Memory (PKM) Mixture of Experts Tucker decomposition PEER
Closest papers by meaning
- Memory Layers at Scale (ICML 2025 · similarity 0.86)
- IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs (ICLR 2024 · similarity 0.86)
- MemoryFormer : Minimize Transformer Computation by Removing Fully-Connected Layers (NeurIPS 2024 · similarity 0.86)
- MatFormer: Nested Transformer for Elastic Inference (NeurIPS 2024 · similarity 0.86)
- Inference-Time Hyper-Scaling with KV Cache Compression (NeurIPS 2025 · similarity 0.86)
- Memory-Augmenting Decoder-Only Language Models through Encoders (Student Abstract) (AAAI 2024 · similarity 0.86)
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.