Ultra-Sparse Memory Network

2024-11-19

arXiv.org

AI Summary

Transformer inference suffers from high memory-access overhead and latency, while existing Mixture-of-Experts (MoE) approaches fail to alleviate memory-bandwidth bottlenecks. Method: This paper proposes UltraMem, a hyper-sparse memory network embedding millions to hundreds of millions of sparse memory slots. It decouples model parameter count from computational cost via dynamic routing, low-rank key-value compression, sparse activation, and gradient optimization. Contribution/Results: UltraMem is the first framework enabling efficient training and inference with up to 100M sparse memory slots, breaking MoE’s memory-bandwidth limitations. Empirical analysis reveals superior scaling laws for hyper-sparse architectures. Under identical compute budgets, a 20M-slot variant matches or exceeds state-of-the-art performance while achieving faster inference. Moreover, UltraMem provides a viable pathway toward deploying models with billion-scale memory slots.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

Abstract

It is widely acknowledged that the performance of Transformer models is logarithmically related to their number of parameters and computational complexity. While approaches like Mixture of Experts (MoE) decouple parameter count from computational complexity, they still face challenges in inference due to high memory access costs. This work introduces UltraMem, incorporating large-scale, ultra-sparse memory layer to address these limitations. Our approach significantly reduces inference latency while maintaining model performance. We also investigate the scaling laws of this new architecture, demonstrating that it not only exhibits favorable scaling properties but outperforms MoE. In experiments, the largest UltraMem we train has 20 million memory slots. The results show that our method achieves state-of-the-art inference speed and model performance within a given computational budget, paving the way for billions of slots or experts.

Problem

• Reduces inference latency significantly
• Maintains model performance effectively
• Outperforms Mixture of Experts scaling

Innovation

• Ultra-sparse memory layer
• Reduces inference latency
• Scales beyond MoE performance