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

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. 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. 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.

Large language models (LLMs) are very powerful at understanding and generating text, but training them is slow and resource-intensive. One reason is that they often process local patterns in text inefficiently, repeatedly learning the same small-scale structures. Another is that the way they store knowledge is tied to their computation, making it harder to access global information when needed. To address these challenges, we designed LoKiFormer, a new model that adds two specialized components: Local Fusion Attention (LFA), which helps the model focus on and combine nearby information more effectively, and Knowledge Memory Module (KMM), which provides a separate memory for storing and retrieving global knowledge. Together, these modules allow the model to handle both local details and global knowledge more efficiently. In experiments, LoKiFormer achieved 1.33× faster training than standard models, demonstrating a more effective and resource-efficient approach to building large language models.