MIDUS: Memory-Infused Depth Up-Scaling

2025-12-15

AI Summary

To address the explosive parameter growth and prohibitive inference overhead arising from deep scaling of large language models (LLMs), this paper proposes the Head-level Memory Layer (HML)—a lightweight, architecture-preserving module that replaces repeatedly stacked feed-forward networks (FFNs) with head-specific, independent memory banks. HML integrates head-level information retrieval, attention-head role modeling, per-head value decomposition, and sparse memory access, enabling efficient depth upscaling while preserving the core Transformer functional structure. In continual pretraining (CPT) experiments, HML achieves substantial performance gains over strong depth-upscaling (DUS) baselines with minimal parameter increase, significantly reduces inference latency, and delivers robust accuracy improvements. Crucially, HML is the first approach to break the long-standing efficiency–performance trade-off inherent in conventional FFN replication paradigms.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Scaling large language models (LLMs) demands approaches that increase capacity without incurring excessive parameter growth or inference cost. Depth Up-Scaling (DUS) has emerged as a promising strategy by duplicating layers and applying Continual Pre-training (CPT), but its reliance on feed-forward networks (FFNs) limits efficiency and attainable gains. We introduce Memory-Infused Depth Up-Scaling (MIDUS), which replaces FFNs in duplicated blocks with a head-wise memory (HML) layer. Motivated by observations that attention heads have distinct roles both across and within layers, MIDUS assigns an independent memory bank to each head, enabling head-wise retrieval and injecting information into subsequent layers while preserving head-wise functional structure. This design combines sparse memory access with head-wise representations and incorporates an efficient per-head value factorization module, thereby relaxing the usual efficiency-performance trade-off. Across our CPT experiments, MIDUS exhibits robust performance improvements over strong DUS baselines while maintaining a highly efficient parameter footprint. Our findings establish MIDUS as a compelling and resource-efficient alternative to conventional FFN replication for depth up-scaling by leveraging its head-wise memory design.

Problem

• Scaling LLMs efficiently without excessive parameters
• Improving Depth Up-Scaling by replacing FFNs with memory
• Enhancing model performance while maintaining parameter efficiency

Innovation

• Head-wise memory layer replaces feed-forward networks
• Independent memory banks enable head-wise retrieval and injection
• Efficient per-head value factorization reduces parameter footprint