LoKiFormer: Locality-aware Attention with Decoupled Knowledge Memory for Efficient Large Language Model Pretraining
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.