LoKiFormer: Locality-aware Attention with Decoupled Knowledge Memory for Efficient Large Language Model Pretraining

LoKiFormer：用于高效大语言模型预训练的局部感知注意力与解耦知识记忆架构

Abstract: Large language models (LLMs) have achieved remarkable breakthroughs across various applications. However, their architectures remain inefficient in pretraining due to two main limitations: (i) self-attention lacks an explicit inductive bias for locality, leading to redundant modeling of sequence-internal local information; (ii) mixture-of-experts (MoE) implicitly couples knowledge storage with computational pathways, hindering flexible access to sequence-external global knowledge.

摘要： 大语言模型（LLM）在各种应用中取得了显著突破。然而，由于两个主要限制，其架构在预训练阶段仍然效率低下：(i) 自注意力机制缺乏对局部性的显式归纳偏置，导致对序列内部局部信息的冗余建模；(ii) 专家混合模型（MoE）将知识存储与计算路径隐式耦合，阻碍了对序列外部全局知识的灵活访问。

To overcome these limitations, we propose LoKiFormer, a novel LLM architecture that augments the standard decoder with two dedicated modules: 1) Local Fusion Attention (LFA), which incorporates a convolutional fusion to attention, explicitly capturing local patterns and allowing the attention to operate on more informative representations; 2) Knowledge Memory Module (KMM), which introduces a parametric key-value memory that explicitly stores global knowledge in addressable slots, decoupling storage from computation and enabling direct knowledge retrieval.

为了克服这些限制，我们提出了 LoKiFormer，这是一种新型 LLM 架构，通过两个专用模块增强了标准解码器：1) 局部融合注意力（LFA），它将卷积融合引入注意力机制，显式捕捉局部模式，并使注意力能够在信息量更丰富的表征上运行；2) 知识记忆模块（KMM），它引入了一种参数化键值记忆，通过可寻址槽位显式存储全局知识，实现了存储与计算的解耦，并支持直接的知识检索。

Together, these modules enable LoKiFormer to achieve more efficient and effective integration of information at both levels. Experimental results show that LoKiFormer converges 1.33x faster in pre-training than baseline models, underscoring its superiority over existing LLM architectures.

这些模块共同使 LoKiFormer 能够在两个层面上实现更高效、更有效的信息整合。实验结果表明，LoKiFormer 在预训练中的收敛速度比基准模型快 1.33 倍，凸显了其优于现有 LLM 架构的性能。