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

2026-08-12

AI Summary

This work addresses the inefficiencies in large language model pretraining stemming from the self-attention mechanism’s lack of locality inductive bias and the entanglement of knowledge storage with computation paths in mixture-of-experts architectures, which hinders flexible access to global knowledge. To overcome these limitations, the authors propose LoKiFormer, a novel architecture that integrates Local Fusion Attention (LFA) into the standard decoder to explicitly model local patterns and introduces a Knowledge Memory Module (KMM) that decouples knowledge storage from computation via addressable parametric memory. This approach is the first to explicitly incorporate locality inductive bias into the attention mechanism, enabling efficient synergy between local and global information. Experiments demonstrate that LoKiFormer achieves a 1.33× faster convergence rate during pretraining and significantly outperforms existing large language model architectures.

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

Problem

• large language models
• pretraining efficiency
• locality
• knowledge memory
• self-attention

Innovation

• Locality-aware Attention
• Decoupled Knowledge Memory
• Local Fusion Attention
• Knowledge Memory Module
• Efficient Pretraining