English

MIDUS: Memory-Infused Depth Up-Scaling

Machine Learning 2026-05-12 v2 Artificial Intelligence

Abstract

Expanding pre-trained language models offers a practical way to increase capacity without training larger models from scratch. Depth Up-Scaling (DUS) does so by duplicating Transformer blocks and inserting them into a pre-trained backbone. This process also duplicates FFN-heavy blocks, increasing parameter and compute cost while adding capacity through a block-level dense residual branch. Yet prior work suggests that added capacity need not remain tied to dense FFN branches, while attention heads often play heterogeneous roles, motivating more efficient head-level residual corrections. We propose Memory-Infused Depth Up-Scaling (MIDUS), which replaces the duplicated FFN branches with memory layers and turns added depth into lightweight retrieval-based residual capacity. We introduce a Head-wise Memory Layer (HML), which combines multi-head product-key memory with Head-wise Implicit Value Expansion (HIVE). HML assigns each head a distinct key space, while HIVE realizes head-specific values from a shared latent bank through compact projections. Alongside empirical improvements in performance and efficiency, our head-importance and fixed-retrieval structural analyses characterize HML with HIVE as a structurally distinct, head-conditioned alternative to FFN-based residual expansion.

Cite

@article{arxiv.2512.13751,
  title  = {MIDUS: Memory-Infused Depth Up-Scaling},
  author = {Taero Kim and Hoyoon Byun and Youngjun Choi and Sungrae Park and Kyungwoo Song},
  journal= {arXiv preprint arXiv:2512.13751},
  year   = {2026}
}
R2 v1 2026-07-01T08:25:58.409Z