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 (i.e., our LoKiFormer) are 31.82, 30.57, and 28.50 on the training set, and 31.82, 30.88, and 29.08 on the evaluation set, respectively. Using the baseline’s 10k-step PPL (31.82) as a reference, the LFA-augmented model reaches this level at 9k steps, showing 1.11× faster convergence. With both LFA and KMM, our LoKiFormer achieves the same target at 7.8k (train) and 7.5k (eval) steps, corresponding to 1.28× and 1.33× faster convergence. We also conduct zero-shot downstream ablations on the base model to isolate each module’s effect. On MMLU, both LFA and KMM yield clear gains over the baseline: 17.9 → 21.2 (+LFA) and 22.9 (+KMM). Combining both further improves performance (25.7), indicating complementary effects. These results demonstrate that our proposed LFA and KMM improves convergence efficiency.

We evaluated the scalability of LoKiFormer across model sizes from 1B to 60B parameters. PPL decreases consistently with larger capacity, from 10.05 at 1B to 5.82 at 60B, confirming smooth scaling behavior. In addition, the results indicates that the proposed components, LFA and KMM, remain effective across scales. Pretraining further exhibited stable optimization dynamics, with maximum gradient norms remaining below 1.0 across all models and no collapse or divergence observed. We further compare performance across the base models of LoKiFormer-7B, LoKiFormer-33B, and LoKiFormer-60B. Larger models show significant improvements: LoKiFormer-33B achieves 4.4 points higher on MMLU and 17.9 points on MMLU-Pro compared to LoKiFormer-7B, while LoKiFormer-60B achieves 84.3 on MMLU and 54.6 on MMLU-Pro. The results above demonstrate that LoKiFormer can be scaled reliably to tens of billions of parameters.

We compare LoKiFormer-7B with size-comparable open-source models across four domains. Despite a similar parameter scale, LoKiFormer consistently achieves substantial gains. On language benchmarks, it scores 91.5 on MMLU, 93.4 on CMMLU, and 92.8 on C-Eval, outperforming all dense and MoE baselines by large margins. On reasoning tasks, it attains 89.5 on HellaSwag and 96.2 on ARC-C, again surpassing the strongest size-comparable models. These improvements are mainly attributed to LFA and KMM, which enhance local dependency modeling and explicit knowledge utilization, leading to richer representations during pretraining. Overall, LoKiFormer achieves state-of-the-art performance among size-comparable LLMs.

We compare LoKiFormer-7B with leading open-source LLMs at much larger scales. Despite having only 7B total parameters with 3B activated, our LoKiFormer consistently outperforms these frontier models across most domains. In language understanding, it achieves 91.5 on MMLU, 93.4 on CMMLU, and 92.8 on C-Eval, substantially surpassing the best dense and MoE baselines. Overall, these results demonstrate the efficiency and effectiveness of our LoKiFormer: with less than one-tenth of the activated parameters, it learns more effectively during pretraining and achieves state-of-the-art performance in most domains, highlighting the benefits of LFA and KMM.

Besides, we compare LoKiFormer-7B with frontier closed-source systems, including Claude-3.5-Sonnet and Gemini-1.5-Pro. Despite its much smaller scale, LoKiFormer demonstrates highly competitive performance. In language understanding, it achieves 91.5 on MMLU, outperforming both Claude-3.5-Sonnet (88.3) and Gemini-1.5-Pro (85.9) by large margins. These highlight that LoKiFormer can deliver competitive or superior performance to frontier closed-source systems on several challenging benchmarks, despite operating at a fraction of their parameter scale.

To verify that KMM stores explicit and editable domain-level knowledge, we perform knowledge-field removal experiments. For selected fields in Layer 4, we zero out their parameters and evaluate performance on five representative MMLU domains. Removing a field mainly degrades its corresponding domain (e.g., algebra drops by 31.8% when Field 32 is removed), indicating that KMM organizes knowledge into domain-specialized slots. Cross-domain effects align with conceptual relatedness: politics and history fields affect both domains, while STEM fields exhibit limited mutual influence. These results demonstrate the interpretability and editability of knowledge in KMM.

We compare the training throughput (in tokens per second) of LoKiFormer-7B against the baseline without LFA and KMM under identical hardware configurations (8×NVIDIA A100 GPUs, 80GB memory). At 4K context, the measured throughput for the baseline is 975 tokens/second, while LoKiFormer-7B achieves a throughput of 950 tokens/second. This analysis confirms that the per-step computational overhead introduced by our novel modules is minimal (i.e., 2.6%). Therefore, the significant reduction in the number of training steps required to reach the target loss, translates directly into a substantial reduction in total wall-clock training time. We also report the inference time of our LoKiFormer-7B on Ascend 910B (8 NPUs). At 4K context, compared with th baseline, decoding throughput remains nearly unchanged (1299 vs. 1294 tokens/s/NPU, bs=256), while time-to-first-token increases only slightly (132 → 136 ms). The results above show high computational efficiency of our LoKiFormer.

We ablate the kernel size k of the convolution in the local fusion module with different k (i.e., k = 2 and k = 4). Training PPL decreases from 31.73 (baseline) to 31.06 (k = 2) and 30.59 (k = 4). Within the tested range, we choose k = 4 as a balance between effectiveness and efficiency: it consistently outperforms k = 2, while larger values are not explored due to the high cost of large-scale pretraining. This setting yields strong gains with manageable overhead, indicating that an appropriate local receptive field substantially improves pretraining efficiency.

We vary the number of groups g in the convolution of LFA, comparing g=1 (no grouping), g=d (grouped by feature dimension), and g=h (aligned with attention heads). Under identical settings, we report training PPL at 10k steps. PPL is 31.07 for g=1, 31.49 for g=d, and 30.49 for g=h. Relative to g=1, g=h reduces PPL by 1.87%, while g=d slightly degrades performance; compared to g=d, g=h achieves a 3.18% lower PPL. These results indicate that grouping is beneficial in LFA, with head-aligned grouping yielding the best performance.

We ablate the number of knowledge fields F in KMM. Training PPL on the cc-en dataset decreases monotonically as F increases, from 31.62 (without KMM) to 30.78 (F =32), 30.35 (F =64), and 29.88 (F =128). These correspond to relative reductions of 2.66%, 4.02%, and 5.50%, respectively. While F =128 achieves the lowest PPL, F =64 captures about 73% of the total reduction with lower computational and memory cost. We therefore adopt F =64 as a balance between efficiency and performance.

We proposed LoKiFormer, a large language model architecture that augments the standard Transformer decoder with two complementary modules: LFA for efficient short-range dependency modeling and the KMM for explicit knowledge retrieval. Together, these modules enhance local contextual modeling and global knowledge utilization, resulting in faster pretraining convergence and stronger performance. In the future, we plan to extend LoKiFormer to multimodal reasoning, such as text, images, and audio. We aim to investigate how LFA can improve local dependency modeling across diverse inputs, while KMM can facilitate the integration and retrieval of knowledge across different modalities. This extension will help us address more complex reasoning tasks, allowing LoKiFormer to leverage both richer contextual and cross-modal knowledge for improved performance.

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.