Title: MIDUS: Memory-Infused Depth Up-Scaling
Primary area: foundation or frontier models, including LLMs
Keywords: Depth Up-Scaling, Memory Layer, Large Language Model

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 dense feed-forward networks (FFNs) limits efficiency and attainable gains. We introduce MIDUS (Memory-Infused Depth Up-Scaling), which replaces FFNs with a head-wise memory layer (HML). Motivated by observations that attention heads exhibit diverse 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-specific functional structure. This design breaks the usual efficiency–performance trade-off by combining sparse memory access with head-specific representations, and includes an efficient per-head value factorization module. In large-scale CPT, MIDUS consistently outperforms strong DUS baselines while using fewer trainable parameters and significantly less training-time memory, and quality improves further by increasing memory capacity even at fixed depth. These results establish head-wise memory as a highly effective and efficient mechanism for DUS of general-purpose LLMs.

TL;DR: MIDUS (Memory-Infused Depth Up-Scaling) advances DUS by replacing dense FFNs with head-wise memory layers, overcoming the conventional efficiency–performance trade-off and enabling more efficient and generalizable LLMs.