使用分层记忆预训练分离长尾与常识知识

在现代语言模型的研发中，参数规模的扩大通常被视为提升模型性能的重要途径。更大的模型能够储存更多的世界知识并展现更强的推理能力。然而，这种单一依赖大规模参数的做法存在局限性：一方面，模型仅能在特定任务中利用所储存知识的一小部分；另一方面，在边缘设备上运行这类模型极为困难，因其受限于推理过程的内存与计算能力。因此，一种更加高效、灵活的模型预训练方法亟需被提出。

为了应对这一挑战，我们提出了基于记忆增强架构的预训练策略，同时该方法与现有硬件生态紧密结合。我们的研究引入了一种全新的层次记忆架构，设计了小型语言模型与大型层次参数记忆库协同工作的机制。在此架构中，小型模型主要学习并运用常识性知识和一般推理能力，而长尾型的专业知识则由参数化记忆单元管理。模型在预训练和推理过程中，根据上下文动态提取所需的记忆块，以进一步增强性能。

基于万亿级别的语料训练，我们展示了显著的性能提升。一款拥有160M参数的小型语言模型，通过访问一组18M参数的上下文相关记忆块，其性能与拥有超过两倍参数的常规语言模型媲美。而这些记忆块来自一个拥有4.6B参数的记忆库。

通过大量实验，我们进一步探究了不同类型、尺寸的参数化记忆模块对于Transformer架构的最优适配情况。实验表明，无论是在预训练阶段添加，还是在后续阶段（后期补充）添加，我们提出的层次记忆前馈架构都具有稳健的适应性。此外，我们的研究表明，参数规模可以扩展至超过21B参数，而模型性能在此基础上进一步提升。

这种创新性记忆架构不仅提高了模型在处理长尾信息时的效率，还在有限的硬件资源下展现了卓越性能，为边缘设备的推理与计算需求提供了新路径。

在资源受限的移动设备上对大语言模型（LLM）进行高效微调一直是一个难题。即便是对部分参数（如 LoRA）的回传微调，也需要消耗极高的内存，常被认为是不切实际的。在此背景下，一些替代性方法应运而生，例如零阶优化（Zeroth-Order Optimization，ZO）。尽管 ZO 方法显著减少了对内存的需求，但却牺牲了模型的收敛速度，较传统的反向传播慢了 10 至 100 倍。

针对这一问题，学术界正积极探索能够兼顾内存高效性与快速收敛的方法。这一领域的研究为实现大语言模型的多样化部署铺平了道路。