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

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.33× faster in pre-training than baseline models, underscoring its superiority over existing LLM architectures.

LLM architectures are expected to incorporate efficient mechanisms for modeling and processing information at different levels: 1) Local-level information, which is internal to the input sequence, encompassing syntactic structures and short-range dependencies between adjacent tokens. Effectively modeling such local patterns is crucial for capturing the finer-grained semantics within a sequence, as these local structures often form the foundation for higher-level meaning extraction. 2) Global-level information, on the other hand, pertains to knowledge that is independent of any specific input sequence and includes broader concepts such as commonsense knowledge, factual information. This type of information can be internalized during pre-training and is vital for tasks that require deep reasoning and generalization beyond the scope of a single sequence. The challenge lies in effectively combining these local patterns with the global knowledge to empower LLMs for complex tasks.

However, we argue that existing LLM architectures are computationally inefficient during pretraining in integrating these two levels of information. For modeling sequence-internal local information, while attention mechanisms can capture both short- and long-range dependencies, it does so through redundant full pairwise interactions, which makes learning local patterns computationally inefficient. For incorporating sequence-external global information, MoE implicitly distribute knowledge across the weights of expert networks. This couples knowledge storage with computational transformation, making that knowledge retrieval becomes indirect, and the process of updating or accessing specific knowledge is inflexible. Thus, designing an architecture that efficiently and effectively integrates local and global information remains an open challenge.

In this paper, we propose LoKiFormer, a novel LLM architecture built upon an enhanced decoder block for efficient pretraining. It extends the vanilla decoder with two new modules: the Local Fusion Attention (LFA) and the Knowledge Memory Module (KMM), which are designed to handle sequence-internal local information and sequence-external global information. For modeling local information, LFA incorporates a convolution operation that fuses adjacent token representations before the attention mechanism. This introduces an explicit local inductive bias, allowing the model to capture short-range dependencies more efficiently in pretraining while enabling the attention layer to concentrate on modeling broader contextual relationships. For incorporating global knowledge, KMM introduces a parameterized key–value memory structure that explicitly stores knowledge in addressable slots. This setup decouples knowledge storage from computational pathways, enabling tokens to directly query and retrieve relevant knowledge, thereby enhancing the transparency and flexibility of knowledge access in pretraining. By integrating LFA and KMM, LoKiFormer achieves a more effective fusion of information at both levels, leading to improved representational capacity and adaptability in tackling complex NLP tasks.

Our novelty and main contributions are as follows:

• Local Inductive Bias for Efficient Attention. We identify that self-attention is inefficient for modeling short-range dependencies. To address this, our LFA introduces a convolutional fusion layer to explicitly capture local patterns before attention, reducing redundant local interactions and allowing attention to focus on more complex and global contextual relationships.

• Decoupled Parametric Memory for Global Knowledge. We argue that implicit knowledge representation in MoE couples knowledge storage with computation, limiting transparent access and flexible manipulation. To address this, our KMM explicitly stores reusable knowledge abstractions in addressable slots, enabling direct knowledge retrieval and interpretable inference.

• Improved Pretraining Efficiency. Our architecture achieves 1.33× faster convergence during pretraining, reaching the same validation loss in 7.5K steps compared to 10K steps for the baseline. This shows superior optimization efficiency, leading to reduced computational cost and energy consumption.

In this paper, we propose LoKiFormer, a novel LLM architecture for efficient pretraining, which enhances the vanilla decoder block with two complementary modules: Local Fusion Attention (LFA) and Knowledge Memory Module (KMM). These together improve the effectiveness and efficiency of information exploitation at both local and global levels. Our LFA extends the Multi-Head Latent Attention (MLA) framework by introducing an explicit local inductive bias. Instead of relying on pairwise interactions, LFA applies convolutional fusion to adjacent tokens before attention. This enables more efficient capture of short-range dependencies while providing the attention with richer contextual representations for modeling broader contextual relationships. At the global level, our KMM provides an explicit key–value memory for knowledge access. It takes the latent representation c as input, projects it into queries, matches them against learnable keys that encode commonsense and domain-specific knowledge fields. Subsequently, we aggregate the corresponding values to produce the output. Finally, the output of KMM is combined with the MoE result to form the final decoder-block representation. In this case, KMM augments the implicit parameterization of MoEs with an explicit, reusable memory pathway, improving transparency and flexibility in knowledge access.

Natural language exhibits strong locality, where adjacent tokens are highly correlated at the semantic levels. Prior works have demonstrated that local modeling improves representation performance. However, both vanilla self-attention and multi-head latent attention (MLA) model short-range dependencies together with long-range ones, requiring global pairwise interactions even when only local patterns are needed. This may be inefficient during pretraining for capturing local fine-grained relationships between adjacent tokens. To alleviate this issue, we propose Local Fusion Attention (LFA), which enhances the attention by introducing a group convolution that explicitly fuses features of neighboring tokens. This provides a strong local inductive bias, enabling more efficient modeling of short-range dependencies while preparing better contextualized representations.

Formally, we split the hidden states U∈RL×d along the feature dimension into h groups, each of size dh = d/h: U= [U(1), U(2), · · ·, U(h)], where U(g) ∈ RL×dh denotes the hidden states for the group g. The number of groups is set equal to the number of attention heads h, so that each head receives its own locally fused representation and learn distinct local fusion patterns. For each group g ∈ [1, h], we apply a causal 1D convolution with a kernel Θ(g) ∈ Rk×dh ×dh along the sequence dimension to aggregate information from its k immediate neighbors within the same group, creating a richer, context-aware representation before being projected into the latent space for attention. At position t, the aggregation involves only tokens in the range [t − k + 1, t] (with left padding if needed). This ensures no access to future tokens, maintaining causal masking in accordance with the standard attention mechanism. The fused representation for group g is computed as: Û(g)t = Σ s=0 k−1 U(g)t−s · Θ(g)[s], where Θ(g)[s] ∈ Rdh ×dh is the slice of the kernel at position s. The index t − s explicitly moves along the sequence axis (L), meaning that each output position aggregates information from its previous k neighboring tokens within the same group. Finally, we concatenate all groups to obtain Û= [Û(1), Û(2), · · ·, Û(h)]. Note that we use a stride of 1 and apply suitable padding so that the sequence length remains unchanged, i.e., Û(g) ∈ RL×dh. The locally fused representation Û is then used for query, key and value construction via Eqn.(1). Our proposed LFA is a lightweight modification only introduces a small number of additional parameters (only 0.41% in our LoKiFormer-7B model). Empirical results show that LFA accelerates convergence in training and consistently enhances downstream performance, demonstrating its effectiveness despite the minimal overhead.

While the proposed LFA focuses on enhancing the modeling of short-range dependencies, effectively leveraging global-level knowledge remains a challenge. Existing approaches such as MoE layers implicitly encode knowledge within distributed parameters, which must be reactivated through linear transformation for every query. This implicit storage lacks transparency and restricts efficient reuse of knowledge patterns. To address this and improve knowledge internalization during pretraining, we propose a Knowledge Memory Module (KMM), which explicitly organizes global knowledge into parameterized key–value fields that can be directly queried during model execution.

Instead of directly using the original hidden states U, KMM takes the latent representation cKV ∈ RL×dc from MLA as input. This compact representation is both computationally efficient and semantically rich, making it well-suited for querying global knowledge fields. We first project it to query space with learnable parameters WH ∈ Rdc ×du: H = cKV WH, where H ∈ RL×du denotes the query features projected from the latent KV representation cKV. Our KMM adopts a multi-group retrieval scheme to increase capacity. Formally, we store the knowledge fields as parameterized keys and values: K ∈ RF ×du, V ∈ RF ×dv, where F denotes the number of fields, du is the query/key dimension, and dv is the value dimension. Note that K and V are randomly initialized and optimized end-to-end during training. We split the projected query H and the knowledge field parameters K, V along the feature dimension into c groups, where each group corresponds to a lower-dimensional subspace. In each group, KMM computes the similarity between the group-specific query H(i) and the knowledge keys K(i) using scaled dot-product, and then aggregates the corresponding values V(i) to obtain the retrieved representation Z(i). The outputs from all groups are finally concatenated to form the overall representation: KMM(H; K, V) = [Z(1),..., Z(c)] ∈ RL×d, Z(i) = softmax(H(i)K(i)⊤/√du)V(i). The knowledge-enhanced output in Eqn. (6) then will be projected back to the model dimension through a linear transformation with learnable parameters WO ∈ Rdv ×d, yielding OK. This representation is then combined with the output of the MLA, denoted as OA, by simple addition: OA + OK. The combined representation is subsequently processed by MoE layer, producing the final output: O = MOE(OA + OK). In this way, the final output integrates the generative expressiveness of MoE with the explicit knowledge reuse enabled by KMM, allowing the model to balance the creation of novel patterns with the efficient retrieval of knowledge fields. During inference, the knowledge fields are fixed but their usage is still query-dependent. Different inputs retrieve different knowledge fields based on Eqn.(6), allowing the model to adapt to a wide range of queries without modifying the memory itself. This decoupling of knowledge from computation enhances the interpretability and modularity.

The roles of KMM and attention are fundamentally different. In attention mechanisms, both keys and values are dynamically computed from the input at each step, allowing the model to capture contextual dependencies within the sequence. This process is flexible but computationally expensive, as it requires full pairwise interactions for each token, making it inefficient for longer sequences. In contrast, KMM uses fixed, parameterized keys and values for knowledge retrieval, which are stored in dedicated memory fields. This allows KMM to provide explicit access to pre-learned knowledge abstractions that are decoupled from the input sequence, enabling efficient knowledge reuse. While attention models dependencies within a sequence, KMM enables the reuse of global knowledge across tasks without recalculating context each time. By decoupling knowledge storage from computation, KMM enhances both the flexibility and interpretability.

To provide an initial validation of our design, we conduct a qualitative analysis by examining the model’s behavior on five diverse domains in MMLU benchmark. The results demonstrates that our KMM facilitates the emergence of semantically specialized knowledge fields. The analysis reveals that, through end-to-end training, individual fields spontaneously evolve to represent concepts from specific domains (e.g., Field 25 is consistently activated for history-related questions, while Field 51 is specialized for physics). The specialization is statistically robust, as evidenced by consistent activation patterns aggregated across multiple layers and a large number of samples. This observed self-organization into a structured memory is a key advantage of our explicit design over implicit approaches like MoE. This enables targeted knowledge retrieval, much like consulting an expert. This makes the model’s predictions more interpretable and reliable, as the knowledge source is directly observable.

Based on the LoKiFormer architecture, we build a family of models with parameter sizes 1B, 5B, 7B, 13B, 33B, and 60B. This enables a systematic evaluation of the effectiveness of LoKiFormer across model capacities. Unless otherwise specified, our main experiments are conducted on LoKiFormer-7B. We evaluate our model across four core dimensions using standard benchmarks. For language understanding, we report results on MMLU for English, and CMMLU, C-Eval for Chinese. Reasoning is assessed with HellaSwag and ARC-Challenge. We measure code generation performance via pass@1 on HumanEval, and evaluate mathematical reasoning using GSM8K.

We pretrain our model on the publicly released Matrix Data Pile, a comprehensive bilingual corpus of 4.5T high-quality tokens curated for the MAP-Neo series, comprising re-processed high-quality English datasets (e.g., RedPajama, Dolma) and Chinese datasets (e.g., Skypile, ChineseWebText), along with a large volume of newly crawled Chinese web content. For validation, we randomly sample 1% from the English Common Crawl portions of the Matrix Data Pile to form validation sets of approximately 18B. For SFT, we employ a curated dataset of approximately 10B high-quality tokens, collected and filtered from existing public instruction-following corpora.

We investigate the effect of different components by progressively add our proposed modules in the baseline. At 10k training steps, the PPL of the baseline, baseline+LFA, and baseline+LFA+KMM…

Based on the paper, here are the specific improvements I can make to AI systems:

1. Add a convolutional fusion layer before self-attention to explicitly capture local token dependencies (syntactic structures, short-range relationships) rather than relying on redundant full pairwise attention interactions. This reduces computational waste on local patterns and lets attention focus on long-range context.

2. Introduce a decoupled parametric key-value memory module that stores global knowledge (commonsense, factual information) in addressable slots, separate from computational pathways. This allows direct, interpretable retrieval of knowledge instead of implicitly encoding it in network weights like MoE does.

3. Enable explicit knowledge field specialization where memory slots self-organize into domain-specific experts (e.g., one slot for history, another for physics) through end-to-end training, making knowledge access transparent and editable—removing a slot degrades only its corresponding domain.

4. Combine local fusion attention with the knowledge memory module to achieve complementary gains: local patterns are captured efficiently, while global knowledge is retrieved explicitly, leading to 1.33× faster pretraining convergence and state-of-the-art performance at 7B scale (91.5 on MMLU, outperforming models 10× larger).

5. Use head-aligned grouping in the convolution (number of groups = number of attention heads) to maximize local fusion effectiveness, reducing training perplexity by 1.87% compared to no grouping.

6. Set the convolution kernel size to 4 for optimal local receptive field, and use 64 knowledge fields as a balance between performance and computational cost, capturing 73% of the improvement from using 128 fields.

The improved AI system can:

• Pretrain significantly faster (1.33×) while achieving lower validation loss

• Perform targeted knowledge retrieval and editing (e.g., remove or update specific domain knowledge without retraining)

• Achieve superior language understanding, reasoning, code generation, and math performance at a fraction of the parameter scale of frontier models

• Provide interpretable inference by showing which knowledge fields are activated for given inputs

• Scale reliably from 1B to 60B parameters without optimization instability

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.