PaperlayerView exact revision on arXiv

Source-linked AI summary

MeKi: Memory-based Expert Knowledge Injection for Efficient LLM Scaling

Ning Ding, Fangcheng Liu, Kyungrae Kim, Linji Hao, Kyeng-Hun Lee, Hyeonmok Ko, Yehui Tang

arXiv:2602.03359v1cs.LGcs.AIcs.CL

TL;DR

Edge deployment makes conventional compute-based LLM scaling costly because limited RAM and NPU resources constrain latency and power. MeKi scales capacity through token-level ROM-resident memory experts and re-parameterization, and reports higher performance with identical decoding speed on Qualcomm Snapdragon hardware.

  • Problem

    Compute-heavy LLM scaling is difficult to deploy on edge devices because larger models increase FLOPs, latency, and power consumption.

  • Method

    MeKi adds token-level memory experts to Transformer layers and folds training-time projections into static ROM-resident lookup tables for inference.

  • Results

    MeKi-1.7B rivals a 4B dense model while maintaining identical decoding speed on Qualcomm Snapdragon hardware.

  • Takeaways & Limitations

    The results support shifting the scaling bottleneck from computation to memory for on-device LLMs.

Abstract

from arXiv · show

Scaling Large Language Models (LLMs) typically relies on increasing the number of parameters or test-time computations to boost performance. However, these strategies are impractical for edge device deployment due to limited RAM and NPU resources. Despite hardware constraints, deploying performant LLM on edge devices such as smartphone remains crucial for user experience. To address this, we propose MeKi (Memory-based Expert Knowledge Injection), a novel system that scales LLM capacity via storage space rather than FLOPs. MeKi equips each Transformer layer with token-level memory experts that injects pre-stored semantic knowledge into the generation process. To bridge the gap between training capacity and inference efficiency, we employ a re-parameterization strategy to fold parameter matrices used during training into a compact static lookup table. By offloading the knowledge to ROM, MeKi decouples model capacity from computational cost, introducing zero inference latency overhead. Extensive experiments demonstrate that MeKi significantly outperforms dense LLM baselines with identical inference speed, validating the effectiveness of memory-based scaling paradigm for on-device LLMs. Project homepage is at https://github.com/ningding-o/MeKi.

1 Introduction

MeKi addresses the difficulty of scaling LLMs on edge devices by shifting capacity growth from computation to storage. It uses token-level memory experts and re-parameterization to improve performance without increasing inference FLOPs or latency.

  • Motivation: Edge deployment makes compute-heavy scaling impractical because larger dense models increase FLOPs, latency, and power consumption, while MoE introduces weight-loading overhead.These constraints are especially relevant on mobile hardware, where latency and power are user-sensitive.
  • Motivation: MeKi asks whether model capacity can be increased through storage space without increasing inference latency or FLOPs.The approach is motivated by relatively cheap, energy-efficient ROM lookups and underused ROM bandwidth on mobile SoCs.
  • MeKi: MeKi retrieves token-level dedicated knowledge vectors from a layer-wise memory bank and injects them into hidden states, decoupling model capacity from computational cost.The memory experts are designed to provide specialized knowledge for different tokens.
  • MeKi: During training, MeKi uses embedding memories, nonlinear projections, and low-rank gating; after training, re-parameterization merges the projections into a static ROM-resident memory table.This separates training-time representational capacity from inference-time efficiency.
  • Evaluation: MeKi is evaluated at 0.6B, 1.7B, and 4B scales across 10 benchmarks, outperforming dense baselines while maintaining the same inference speed on Qualcomm Snapdragon Android hardware.The reported evaluation covers both model quality and on-device latency.

2 Related Works

Related work has mainly pursued scaling through sparse activation, retrieval, or external lookup structures. MeKi belongs to the storage-centric direction while integrating memory experts into Transformer layers.

  • Evolution of LLM Architectures: Scaling laws and MoE architectures increase capacity through larger training resources or sparse activation, but edge deployment remains comparatively underexplored.MoE reduces per-token computation by activating subsets of weights, yet its expert-weight access pattern creates deployment challenges.
  • New Scaling Paradigm Beyond RAM: External-storage scaling includes retrieval-augmented methods and lookup-table approaches that replace all or part of the FFN.MeKi extends this memory-centric line by using layer-wise token-level experts alongside the Transformer backbone.

3 Architecture

MeKi adds a token-level memory-expert branch parallel to the FFN, injects retrieved knowledge through gated fusion, and re-parameterizes training projections into static lookup tables for efficient inference.

  • Overview: MeKi integrates a layer-specific memory branch parallel to the Transformer FFN and retrieves one knowledge vector per input token.The retrieved vector is projected back to the model dimension and added to the residual stream.
  • Token-Level Expert: For layer l, MeKi stores a memory embedding matrix M_l and obtains each token’s static knowledge vector through an embedding lookup.The memory dimension d_mem is smaller than the model dimension d_model.
  • Token-Level Expert: During training, a nonlinear projection of the global word embedding enriches the token expert representation without adding inference-time cost.The projection is implemented as SwiGLU in practice and is combined with static prior knowledge.
  • Knowledge Injection: MeKi uses additive gated fusion to modulate retrieved knowledge with the hidden state before projecting it into the Transformer’s model dimension.The gate is generated by a low-rank projection and sigmoid, while the projected vector is added to the residual stream.
  • Re-Parameterization: After training, MeKi pre-computes the nonlinear projection over static embedding weights and absorbs it into a new memory lookup table.Inference therefore removes the heavy projection and retains only lookup, lightweight gating, projection, and addition operations.
  • Efficiency: Re-parameterized inference reduces the remaining MeKi computation to O(d_model · d_mem), with d_mem much smaller than d_model.For a 28-layer model with d_mem = 256, ROM traffic is 14KB per token in float16 format.

4 Experiments

Experiments evaluate MeKi against dense and sparse-memory baselines using matched training and model-scale settings. MeKi achieves stronger downstream performance and validation loss while maintaining near-zero inference latency overhead.

  • Dense Baselines: MeKi achieves the best average zero-shot performance across all evaluated model scales against dense baselines.All models are trained from scratch on the same 50B-token subset for fair comparison.
  • Dense Baselines: On ARC-Challenge, MeKi-1.7B scores 37.9 versus 34.4 for the 1.7B baseline and 38.0 for the 4B dense model.On SciQ, MeKi-1.7B reaches 85.4, exceeding its baseline by 4.8 points.
  • Sparse Memory Architectures: MeKi-1.7B averages 59.7, outperforming PLE by 2.7 points and Engram by 1.8 points.Figure 3 also reports lower validation loss for MeKi than the other methods at 0.6B and 1.7B scales.
  • Inference Latency: MeKi maintains the dense baseline’s active RAM parameters and achieves nearly zero latency overhead by offloading re-parameterized memory weights to ROM.Inference uses asynchronous prefetching directly via token ID on a Qualcomm Snapdragon 8 Elite smartphone platform with a 10K KV-cache length.

5 Ablation Study and Analysis

The ablations examine MeKi’s memory components, placement, fusion strategy, memory size, and prediction convergence. Results favor combining static and dynamic memory, placing MeKi parallel to the FFN, using Additive-Sigmoid fusion, and scaling memory dimensions under a storage–capacity trade-off.

  • 5.1 Effectiveness of Static and Dynamic Memory: Static memory alone raises the average score from 53.0 to 54.8 over the 0.6B baseline.This indicates that the static memory learns token-level priors and factual knowledge in embedding space.
  • 5.1 Effectiveness of Static and Dynamic Memory: Dynamic-Only performs similarly to Static-Only, suggesting nonlinear projections synthesize expressive layer-specific features from global embeddings.The projection provides enhanced representability without layer-specific memory.
  • 5.1 Effectiveness of Static and Dynamic Memory: Combining static memory and dynamic projection gains +0.7 over Static-Only and +0.8 over Dynamic-Only across most benchmarks.MeKi-0.6B combines both components through learnable coefficients α_l and β_l.
  • 5.2 Optimal Position for Memory Injection: Parallel-to-FFN placement achieves the best average scores across 10 downstream tasks and the best validation loss.The other placements are reported as suboptimal or worst-performing, with the authors conjecturing information bottlenecks for placements after Attention or FFN.
  • 5.3 Optimal Feature Fusion: Additive-Sigmoid achieves the best training-loss performance among the four feature-fusion strategies.The variants are trained from scratch on MeKi-0.6B and MeKi-1.7B.
  • 5.4 Scaling Law of Memory Size: Validation performance follows a log-linear trend with memory size, leading to selected dimensions d_mem = 128 for MeKi-0.6B and d_mem = 256 for MeKi-1.7B.These dimensions are chosen as storage-cost and model-capacity trade-offs.
  • 5.5 Layer-wise KL Divergence Analysis: MeKi shows systematically lower KL divergence across all layers than the baseline, indicating faster alignment of intermediate predictions with final outputs.The analysis averages results from 10K validation tokens and attributes the observation to layer-specific memory retrieval.
  • 5.6 Projection Analysis: Replacing the nonlinear projector with a simple linear projection degrades both training loss and downstream performance, although re-parameterization removes its inference computation.The projector is absorbed into the final embedding memory during inference.

6 Conclusion

MeKi scales LLM capacity through ROM-resident memory rather than computation. Re-parameterized lookup tables provide zero inference latency overhead while a 1.7B model rivals a 4B dense model at identical decoding speed on Qualcomm Snapdragon hardware.

  • 6 Conclusion: MeKi decouples model capacity from computational cost by re-parameterizing training computation into static memory lookup tables.The memory is offloaded to abundant ROM space for on-device scaling.
  • 6 Conclusion: MeKi achieves zero inference latency overhead while expanding on-device model capacity through storage.The conclusion frames this as shifting the scaling bottleneck from computation to memory.
  • 6 Conclusion: A 1.7B MeKi model rivals a 4B dense model while maintaining identical decoding speed on Qualcomm Snapdragon hardware.

A.1 Detailed Model Architecture and Hyper Parameters

The supplied passage identifies Table 5 as containing detailed model architecture information and training hyperparameters.

  • A.1 Detailed Model Architecture and Hyper Parameters: Table 5 presents detailed model architecture information and training hyperparameters.

A.2 Detailed Statistics for NPU Inference

The NPU inference statistics measure component time consumption on a Qualcomm Snapdragon 8 Elite with a 10K KV-cache length. FFN parameter movement is a major bottleneck for Qwen3-4B, while MoE requires still more DMA time.

  • A.2 Detailed Statistics for NPU Inference: The analysis measures component time consumption during Qwen3 inference on Qualcomm Snapdragon 8 Elite with KV cache length 10K.
  • A.2 Detailed Statistics for NPU Inference: 38% of Qwen3-4B inference time is spent moving FFN parameters from RAM to NPU, versus 32% for FFN computation.DMA denotes moving corresponding parameters from RAM to NPU.
  • A.2 Detailed Statistics for NPU Inference: MoE deployment requires much more DMA time than dense LLM deployment.
Loading 2602.03359v1…