arXiv:2412.09764cs.CLcs.AI2024-12被引 36

用可训练记忆层提升模型性能,不增加算力消耗。

Memory Layers at Scale

  • 引入可并行计算的记忆层,实现低成本信息存取。
  • 在事实类任务上表现优于双倍算力的密集模型。
  • 适合需要高效存储与检索知识的大规模语言模型应用。

记忆层通过可训练的键值查找机制,在不增加计算量(FLOPs)的前提下为模型增添额外参数。稀疏激活的记忆层可补足计算密集型的前馈层,以低成本提供专用的信息存储与检索能力。本文将记忆层推进至当前主流规模,实验证明:在下游任务中,增强后的语言模型性能超越计算量超过两倍的密集模型,且在相同算力与参数预算下优于混合专家模型。尤其在事实性任务上提升显著。本文提供了完全可并行化的记忆层实现,支持高达1280亿记忆参数,预训练至1万亿标记,对比基础模型最大达80亿参数。

原文摘要 · Abstract (English)

Memory layers use a trainable key-value lookup mechanism to add extra parameters to a model without increasing FLOPs. Conceptually, sparsely activated memory layers complement compute-heavy dense feed-forward layers, providing dedicated capacity to store and retrieve information cheaply. This work takes memory layers beyond proof-of-concept, proving their utility at contemporary scale. On downstream tasks, language models augmented with our improved memory layer outperform dense models with more than twice the computation budget, as well as mixture-of-expert models when matched for both compute and parameters. We find gains are especially pronounced for factual tasks. We provide a fully parallelizable memory layer implementation, demonstrating scaling laws with up to 128B memory parameters, pretrained to 1 trillion tokens, comparing to base models with up to 8B parameters.

记忆机制大模型高效架构

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。