Most improvements to large language model architectures target either inference efficiency or downstream task performance. Pretraining efficiency, the question of how quickly a model converges to a useful representation given a fixed compute budget, receives comparatively less architectural attention. LoKiFormer addresses this gap directly, arguing that two structural deficiencies in the standard decoder block create unnecessary redundancy during pretraining: self-attention's indifference to locality, and mixture-of-experts' implicit coupling of knowledge storage with computation.
The paper is worth examining carefully because the proposed fixes are architecturally clean and the empirical claims, while modest in scope, are grounded in systematic ablations rather than cherry-picked benchmarks.
The Core Argument
The authors identify two distinct information processing levels that LLMs must handle. Local-level information covers short-range syntactic and semantic dependencies within a sequence. Global-level information covers world knowledge, commonsense reasoning, and factual content that is independent of any specific input. Their claim is that standard attention handles both levels through the same undifferentiated pairwise interaction mechanism, which is computationally wasteful for short-range patterns. Meanwhile, MoE architectures store knowledge implicitly across expert weights, making targeted retrieval or editing of specific knowledge difficult.
These are not new observations in isolation. Sparse attention work has long noted that many attention heads specialise in local patterns, and retrieval-augmented generation research has argued for separating knowledge from computation. LoKiFormer's contribution is combining both ideas into a single architectural proposal that operates within the decoder block itself, without requiring external retrieval infrastructure at inference time.
Key Contributions
- Local Fusion Attention (LFA): A convolutional fusion layer applied to adjacent token representations before the attention mechanism. The convolution uses grouped convolutions aligned with attention heads (group count g = h), which the ablations show outperforms both ungrouped and fully depthwise variants. Kernel size k = 4 is selected as the practical optimum. The idea is that by pre-computing local context through convolution, the subsequent attention operation receives richer representations and can concentrate on longer-range dependencies rather than redundantly relearning short-range patterns.
- Knowledge Memory Module (KMM): A parametric key-value memory with F learnable slots (F = 64 by default). Tokens project their latent representations into queries, match against fixed-dimension keys, and retrieve aggregated values. This sits alongside the MoE layer, with KMM output combined with MoE output to form the final block representation. Crucially, the memory is trained end-to-end and remains static at inference, so there is no retrieval overhead beyond a matrix multiply against the key bank.
- Pretraining convergence: The combined architecture reaches the same validation loss in 7,500 steps that the baseline requires 10,000 steps to achieve, a 1.33x speedup with only 2.6% throughput overhead per step.
Methodology and Results
The backbone is DeepSeek's Multi-Head Latent Attention (MLA), which compresses key-value representations into a lower-dimensional latent space. LFA and KMM are layered on top of this. The model family spans 1B to 60B parameters, with the 7B variant as the primary evaluation target. Pretraining uses the Matrix Data Pile, a 4.5T token bilingual corpus, on Megatron-LM with H200 and Ascend 910B hardware.
Benchmark results for LoKiFormer-7B are competitive with Llama-3.1-8B, Qwen2.5-7B, and InternLM2-7B across MMLU, CMMLU, C-Eval, HellaSwag, ARC-Challenge, HumanEval, and GSM8K. The paper also compares against frontier 70B+ models, where the 7B model predictably trails but the gap is used to contextualise rather than overclaim.
The knowledge field editability analysis is one of the more interesting empirical contributions. By zeroing out individual KMM slots and measuring performance drops across MMLU domains, the authors show that specific fields specialise in specific domains. Field 32 removal causes a 31.8% drop in algebra performance with limited cross-domain bleed. Politics and history fields show correlated mutual influence, which is semantically plausible. This is a concrete demonstration that the memory structure is doing something interpretable, not just adding parameters.
Ablations on kernel size, grouping strategy, and field count are thorough for a pretraining paper, where each configuration requires training a 5B model to 10k steps. The monotonic improvement in perplexity with increasing field count (F = 32 to 128) suggests the memory capacity is not yet saturated at the chosen operating point.
Limitations and Open Questions
Several aspects of the paper warrant scrutiny. The throughput comparison is reported at 4K context on 8x A100s, but the primary training runs use 2048 context for 7B and 4096 for 33B/60B models. The relationship between context length, LFA kernel size, and relative efficiency gains is not fully explored. Convolution over long sequences with large kernels is not free, and the paper does not report how throughput scales with context length for LoKiFormer versus the baseline.
The KMM is static at inference. This is presented as a feature (no retrieval overhead), but it also means the knowledge encoded during pretraining cannot be updated without retraining. For applications requiring knowledge currency, this is a genuine constraint compared to retrieval-augmented approaches. The paper's framing of KMM as enabling "flexible" knowledge access is somewhat overstated given this limitation, though the editability experiments do show that targeted zeroing of fields is possible post-hoc.
The convergence speedup claim (1.33x) is measured in steps, not wall-clock time adjusted for the 2.6% per-step overhead. The net wall-clock saving is therefore approximately 1.29x, which is still meaningful but slightly less dramatic than the headline figure implies.
The comparison with baseline models uses instructed versions for all evaluations. Since LoKiFormer-7B undergoes SFT without reinforcement learning while some baselines may include RLHF or DPO, direct performance comparisons carry some confounding. The authors note this but do not fully disentangle the architectural contribution from the training recipe difference.
Future directions mentioned include multimodal extension, which is a natural fit: LFA's local inductive bias could benefit vision token sequences, and KMM's structured memory could store cross-modal knowledge abstractions. Whether the current formulation scales cleanly to vision or audio tokens remains an open question.
Overall, LoKiFormer makes a credible case that explicit local inductive bias and decoupled parametric memory are worth incorporating into the standard decoder block. The architectural ideas are not individually novel, but their integration is clean, the ablations are honest, and the pretraining efficiency gains are practically significant. The full paper is available at arxiv.org/abs/2608.12419, with code at the linked GitHub repository.