MeKi：基于记忆的专家知识注入实现高效LLM扩展

通过ROM存储扩展LLM容量，实现零延迟开销的端侧模型性能提升

前置知识

Mixture-of-Experts (MoE)

混合专家架构是一种通过稀疏激活参数子集来扩展模型容量的方法。在MoE中，模型包含多个独立的前馈网络（FFN）作为'专家'，通过路由机制为每个token选择性激活少数专家。这种架构可以在保持每个token的计算量相对较低的同时，大幅增加模型的总参数量，从而提升模型性能。路由机制通常基于门控网络对token特征进行评估，选择最相关的专家组合。

MeKi将自身与MoE进行对比，说明MoE在边缘设备上因动态路由和内存碎片化带来的延迟开销，而MeKi通过静态查找表避免了这些问题

Re-parameterization（重参数化）

重参数化是一种在训练完成后将复杂计算合并到简单结构中的技术。在本文中，训练阶段使用复杂的非线性投影（SwiGLU）来学习丰富的特征表示，推理阶段则将这些投影预先计算并合并到静态嵌入表中。这样做的好处是训练时保持强大的表达能力，推理时却只需简单的查找操作，实现了计算复杂度的零开销转移。

这是MeKi实现零推理延迟的核心技术，理解重参数化才能明白训练时的复杂计算如何在推理时被消除

ROM与RAM的带宽差异

在现代移动SoC中，ROM（只读存储）和RAM（随机存取存储）具有不同的带宽特性。ROM带宽如UFS-4.0可达4.2GB/s，在模型推理期间基本处于闲置状态；而RAM-NPU带宽则相对有限。MeKi正是利用了这种带宽差异，将静态知识卸载到ROM中，既不占用宝贵的RAM空间，又不会因带宽不足而产生瓶颈。每个token在ROM中仅需读取14KB数据（28层乘以256维乘以2字节）。

理解ROM与RAM的带宽差异是MeKi能在边缘设备上实现零延迟开销的硬件基础

低秩门控机制

低秩门控机制是MeKi中将专家知识注入隐藏状态的核心融合技术。给定隐藏状态 $h_t^l \in \mathbb{R}^{d_{model}}$，通过一个低秩线性投影 $W_{gate}^l \in \mathbb{R}^{d_{mem} \times d_{model}}$ 生成门控信号，再经过sigmoid激活。这个门控信号作为通道级偏移量与专家向量相加，实现上下文感知的知识注入。由于 $d_{mem} \ll d_{model}$（如128 vs 2048），这个操作的计算量非常小。

低秩门控机制是MeKi高效性的关键设计，它在保持最小计算开销的同时实现了有效的知识融合

SwiGLU激活函数

SwiGLU是Shazeer在2020年提出的一种门控线性单元变体，它结合了Swish激活函数和GLU门控机制。在FFN层中，SwiGLU将输入先通过两个线性投影，一个经过Swish激活后与另一个相乘，再通过输出投影。这种设计已被证明比传统的ReLU或GELU激活在Transformer架构中表现更好。在本文中，SwiGLU被用作动态投影函数 $G_l(\cdot)$ 的实现。

论文中使用SwiGLU作为动态投影的实现，且实验证明使用更复杂的SwiGLU比简单线性投影能带来更好的最终性能

LogitLens分析技术

LogitLens是一种用于分析Transformer内部表示的可视化技术。它将模型各层的中间隐藏状态直接输入最终的分类头，观察模型在不同层的预测结果。通过计算这些中间预测与最终输出分布之间的KL散度，可以量化模型表示在各层中的收敛速度。较低的KL散度意味着该层的表示已经接近最终预测，说明模型能更快地形成高置信度的内部状态。

论文使用LogitLens证明MeKi能加速预测收敛，这是验证记忆注入有效性的关键分析工具

研究动机

随着大语言模型性能的持续提升，扩大模型参数量或增加测试时计算量已成为主流的扩展策略。然而，这些计算密集型的扩展方法在边缘设备部署时面临严峻挑战。具体而言，在高通骁龙8 Elite移动平台上，增大密集模型的参数会导致浮点运算量（FLOPs）显著增加，带来用户无法接受的延迟和功耗。虽然MoE架构通过稀疏激活专家来降低单token的计算量，但其动态路由机制和大量不连续专家权重的频繁加载，在移动硬件上引入了显著的延迟开销。在边缘设备上，ROM-RAM和RAM-NPU的带宽远低于数据中心的服务器级GPU，MoE的内存访问模式成为主要的延迟瓶颈。这些限制使得现有的计算中心化扩展策略难以在手机等边缘设备上有效部署。

本文的目标是本文的核心目标是提出一种全新的LLM扩展范式，通过利用存储空间（ROM）而非计算量（FLOPs）来提升模型容量。具体而言，作者希望实现以下目标：第一，设计一个能在保持推理速度不变的情况下显著提升模型性能的架构；第二，将1.7B参数的模型性能提升至与4B密集模型相当的水平；第三，在高通骁龙移动平台上验证该方法的实际推理延迟，确保不引入额外的推理开销；第四，通过系统性的实验验证该方法在多个基准测试上的有效性。

与已有工作不同的是，本文的独特切入角度在于观察到现代移动SoC上ROM带宽在模型推理期间基本闲置这一关键事实。与矩阵乘法（FLOPs密集型）相比，内存查找（如从ROM读取）相对廉价且节能。作者提出的核心问题是：能否在不增加推理延迟和FLOPs的情况下，利用存储空间来扩展模型容量？这一问题将研究焦点从计算中心化设计转向内存中心化设计，为边缘设备上的LLM部署开辟了新的研究方向。与现有的检索增强（RAG）和记忆扩展（如PLE、Engram）等方法相比，MeKi通过Token级专家知识注入和重参数化策略，实现了训练容量与推理效率的最佳平衡。

核心方法

MeKi的核心思路是将Transformer每一层的FFN模块视为一个计算密集的容量单元，然后在其旁边并行添加一个内存密集的记忆模块。直觉上，这类似于给模型配备了一个外部知识库——对于每个输入token，模型会从层特定的记忆库中检索一个专用的知识向量，然后将这个向量注入到隐藏状态中。关键的设计在于训练和推理使用不同的参数配置：训练时使用复杂的非线性投影来学习丰富的表示，推理时则通过重参数化将这些复杂计算折叠成一个紧凑的静态查找表。整个系统与Transformer的FFN模块并行运行，实现了隐式的层宽度扩展，却几乎不引入额外的FLOPs开销。

MeKi与已有方法的本质区别在于三个方面。首先，与MoE相比，MeKi消除了在线路由成本，通过直接基于token ID进行查找实现无缝预取，更适合资源受限的环境。其次，与Per-Layer Embedding（PLE）相比，PLE将嵌入记忆串行放置在FFN之后，容易形成信息瓶颈；而MeKi将记忆模块与FFN并行，作为容量扩展器，避免了前向传播中的信息阻塞。第三，与Engram相比，Engram依赖在线哈希进行知识检索，而MeKi利用低秩门控机制动态增强隐藏状态，提供更优的上下文适应性和更低的延迟。核心技术新颖性在于将训练时的复杂非线性投影 $G_l(\cdot)$ 预先计算并合并到嵌入表中，实现 $\tilde{M}_l = \alpha_l \cdot \text{RMSNorm}(M_l) + \beta_l \cdot G_l(E_{global})$ 的重参数化，使得推理时只需简单的查找操作。

方法步骤详情

MeKi方法分为四个关键步骤。第一步是Token级专家构建：给定输入token ID序列 $x \in \{1, ..., |V|\}^T$，从层特定记忆嵌入矩阵 $M_l \in \mathbb{R}^{|V| \times d_{mem}}$ 中检索静态记忆向量 $m_l^{static}(x_t) = M_l[x_t]$，同时通过非线性投影 $G_l(\cdot)$ 从全局词嵌入 $E_{global}$ 生成动态投影 $m_l^{dyn}(x_t) = G_l(E_{global}[x_t])$，最终通过可学习标量 $\alpha_l$ 和 $\beta_l$ 融合得到专家向量 $e_t^l = \alpha_l \cdot \text{RMSNorm}(m_l^{static}(x_t)) + \beta_l \cdot m_l^{dyn}(x_t)$。第二步是知识注入：使用低秩线性投影生成门控信号 $g_t^l = \sigma(W_{gate}^l h_t^l)$，将其与专家向量相加 $v_t^l = e_t^l + g_t^l$，再通过输出投影 $y_t^l = \text{RMSNorm}(W_{out}^l v_t^l)$ 将结果投射回模型维度。第三步是残差连接：MeKi的输出与FFN输出并行相加 $H' = F(H) + \text{MeKi}(H) + A$，实现隐式宽度扩展。第四步是重参数化：将动态投影合并到静态嵌入表 $\tilde{M}_l = \alpha_l \cdot \text{RMSNorm}(M_l) + \beta_l \cdot G_l(E_{global})$，推理时只需查找 $\tilde{e}_t^l = \tilde{M}_l[x_t]$。

技术新颖性

MeKi的技术新颖性体现在以下几个方面。首先，提出了'基于记忆的专家知识注入'这一全新范式，将模型容量与计算成本解耦，通过存储空间而非FLOPs来扩展模型能力。其次，设计了静态-动态双记忆机制，静态记忆学习token级先验知识，动态投影从全局语义空间合成层特定特征，两者通过可学习系数自适应融合，实验证明这种组合比单独使用任一成分平均提升0.7-0.8个点。第三，提出重参数化策略平衡训练容量与推理效率，训练时使用SwiGLU等复杂投影最大化特征可表达性，推理时将这些投影合并到静态嵌入表中实现零成本查找。第四，通过并行架构设计避免了PLE等串行方法的信息瓶颈问题，MeKi模块与FFN并行运行，作为隐式的层宽度扩展。第五，实验验证了记忆大小的对数线性缩放规律，为存储-容量权衡提供了理论指导。

**Figure 2**: General overview of the proposed MeKi architecture

**Figure 4**: Four potential positions to incorporate the MeKi module

**Figure 5**: Validation loss over different feature fusion strategies

**Figure 6**: The final validation loss values with different memory sizes

实验结果

实验结果全面验证了MeKi架构的有效性。在零样本性能方面，MeKi在所有模型规模上都显著超越了密集基线。以1.7B规模为例，MeKi在10个基准测试上的平均得分达到59.7，比密集基线的56.0提升了3.7个点。特别值得注意的是，MeKi-1.7B在ARC-Challenge上达到37.9分，超越1.7B基线（34.4）3.5个点，有效匹配了4B密集模型（38.0）的性能；在SciQ上达到85.4分，超越基线4.8个点。在LAMBADA语言建模基准上，MeKi-1.7B得分45.6，与4B基线模型（45.6）完全相同，表明注入的专家向量为长程依赖预测提供了关键的语义锚定。与同类方法对比，MeKi-1.7B比PLE和Engram分别高出2.7和1.8个点。推理延迟方面，在高通骁龙8 Elite平台上，MeKi保持了与密集基线相同的生成速度（如1.7B规模下13.7 vs 13.8 token/s），实现了接近零的延迟开销。消融实验证明，静态记忆和动态投影捕获了互补信息，其组合使用比单独使用任一成分都能获得更好的性能。记忆注入的最佳位置是与FFN并行，这比其他三个位置平均高出0.4-0.8个点。特征融合策略方面，加法-Sigmoid方法在0.6B和1.7B规模上都取得了最佳训练损失。记忆大小的缩放实验表明，模型性能随记忆大小呈对数线性趋势增长。LogitLens分析显示，MeKi在所有层上都表现出更低的KL散度，证明其能加速预测收敛。

**Table 1**: Zero-shot performance on downstream tasks - MeKi vs dense baselines

**Table 2**: Zero-shot performance comparison with PLE and Engram

**Table 3**: Zero-shot performance with different types of memory

**Table 4**: Zero-shot performance with MeKi module at different positions

**Figure 3**: The comparison of validation loss for different methods

**Figure 7**: Layer-wise KL Divergence Analysis

查看结构化数据

任务	指标	本文	基线	提升
ARC-Easy (ARC-E)	Zero-shot Accuracy	60.2 (0.6B), 66.2 (1.7B), 77.0 (4B)	56.0 (0.6B), 61.7 (1.7B), 72.9 (4B)	+4.2 (0.6B), +4.5 (1.7B), +4.1 (4B)
ARC-Challenge (ARC-C)	Zero-shot Accuracy	33.6 (0.6B), 37.9 (1.7B), 42.2 (4B)	30.5 (0.6B), 34.4 (1.7B), 38.0 (4B)	+3.1 (0.6B), +3.5 (1.7B), +4.2 (4B)
BoolQ	Zero-shot Accuracy	63.0 (0.6B), 62.4 (1.7B), 70.2 (4B)	58.3 (0.6B), 58.9 (1.7B), 66.3 (4B)	+4.7 (0.6B), +3.5 (1.7B), +3.9 (4B)
COPA	Zero-shot Accuracy	72.0 (0.6B), 74.0 (1.7B), 80.0 (4B)	69.0 (0.6B), 69.0 (1.7B), 75.4 (4B)	+3.0 (0.6B), +5.0 (1.7B), +4.6 (4B)
HellaSwag	Zero-shot Accuracy	49.2 (0.6B), 56.6 (1.7B), 62.3 (4B)	45.7 (0.6B), 51.7 (1.7B), 57.9 (4B)	+3.5 (0.6B), +4.9 (1.7B), +4.4 (4B)
LAMBADA	Zero-shot Accuracy	39.8 (0.6B), 45.6 (1.7B), 50.1 (4B)	35.5 (0.6B), 41.3 (1.7B), 45.6 (4B)	+4.3 (0.6B), +4.3 (1.7B), +4.5 (4B)
OBQA	Zero-shot Accuracy	34.6 (0.6B), 39.0 (1.7B), 40.2 (4B)	34.8 (0.6B), 37.6 (1.7B), 39.4 (4B)	-0.2 (0.6B), +1.4 (1.7B), +0.8 (4B)
PIQA	Zero-shot Accuracy	70.2 (0.6B), 71.7 (1.7B), 75.4 (4B)	68.7 (0.6B), 70.8 (1.7B), 72.9 (4B)	+1.5 (0.6B), +0.9 (1.7B), +2.5 (4B)
SciQ	Zero-shot Accuracy	78.4 (0.6B), 85.4 (1.7B), 87.2 (4B)	77.3 (0.6B), 80.6 (1.7B), 84.4 (4B)	+1.1 (0.6B), +4.8 (1.7B), +2.8 (4B)
WinoGrande	Zero-shot Accuracy	53.8 (0.6B), 58.7 (1.7B), 61.6 (4B)	52.6 (0.6B), 53.6 (1.7B), 58.1 (4B)	+1.2 (0.6B), +5.1 (1.7B), +3.5 (4B)
Average (10 benchmarks)	Zero-shot Accuracy	55.5 (0.6B), 59.7 (1.7B), 63.0 (4B)	53.0 (0.6B), 56.0 (1.7B), 60.5 (4B)	+2.5 (0.6B), +3.7 (1.7B), +2.5 (4B)

局限与改进

尽管MeKi取得了显著的成果，但仍存在一些局限性。首先，作者承认ROM带宽虽然在当前移动平台上足够，但随着模型规模进一步增长，可能需要更高效的内存访问模式。其次，实验仅在50B token的训练预算下进行，未验证更大规模训练数据下的表现。第三，当前实验基于Qwen3系列架构，尚未在其他主流架构（如LLaMA、GPT等）上验证通用性。第四，MeKi的存储开销虽然可接受（1.7B模型约1.09B ROM权重），但对于极端资源受限的设备，存储空间仍是一个需要权衡的因素。第五，论文未详细讨论动态投影函数 $G_l(\cdot)$ 的不同选择对最终性能的具体影响，仅简单比较了SwiGLU和线性投影。第六，关于记忆向量的稀疏性和可压缩性未做深入分析，这对于进一步优化存储效率可能有重要价值。此外，当前实验仅在英语基准上进行，未验证多语言场景下的效果。

独立分析的弱点

从独立分析的角度，MeKi存在以下几个可改进的弱点。首先，当前的Token级专家检索完全基于token ID，这意味着同一个token在所有上下文中都检索相同的专家向量，缺乏对上下文语义的敏感性。改进方向可以是引入轻量级的上下文感知路由机制，根据相邻token的语义动态调整检索的专家向量。其次，重参数化策略虽然消除了推理时的FLOPs，但将所有投影合并到静态表中也限制了模型的适应性。未来可以探索部分重参数化方案，在保持大部分效率的同时保留一定的动态适应能力。第三，当前的门控机制使用简单的低秩线性投影，可能无法充分捕获复杂的上下文依赖关系。可以考虑使用更高效的注意力机制或稀疏门控来增强融合效果。第四，实验规模有限，仅在0.6B、1.7B和4B三个规模上验证，缺乏对更大规模模型（如7B、13B）的探索。第五，论文未讨论记忆向量的冗余性和去重潜力，如果不同层的记忆向量存在高度相似性，可以通过共享或压缩进一步优化存储效率。

未来方向

基于MeKi的成果，未来研究可以向以下几个方向延伸。第一，探索自适应记忆大小分配策略，根据各层的重要性动态调整 $d_{mem}$，而非使用统一的记忆维度。第二，研究记忆向量的可迁移性和可组合性，是否可以将特定领域的记忆库作为插件模块加载到基础模型中，实现高效的知识更新。第三，将MeKi与其他效率优化技术（如量化、剪枝、知识蒸馏）结合，探索协同优化的可能性。第四，扩展到更多模态和任务类型，如视觉语言模型、代码生成等，验证记忆注入范式的通用性。第五，研究记忆库的持续学习能力，探索在不重新训练整个模型的情况下更新和扩展记忆内容的机制。作者也提到，未来可以探索更复杂的非线性投影函数在训练阶段的使用，因为实验表明即使这些函数在推理时被消除，其对最终性能仍有重要影响。

复现评估

从复现角度来看，MeKi的复现条件相对友好。作者已在GitHub上开源了项目代码（https://github.com/ningding-o/MeKi），提供了完整的实现细节。训练数据使用公开的FineWeb-Edu-Dedup数据集，随机采样50B token进行训练，数据获取无门槛。模型架构基于开源的Qwen3系列（0.6B、1.7B、4B），这些基础模型均可公开获取。训练框架使用Megatron-LM，这是广泛使用的开源框架。硬件要求方面，虽然论文未明确说明训练所需的具体GPU配置，但基于50B token的训练规模，估计需要多卡A100或同等算力。推理验证需要高通骁龙8 Elite平台或类似的移动设备。整体而言，只要具备足够的算力资源，按照论文提供的配置应该可以复现主要结果。建议复现时重点关注学习率调度、批大小和序列长度等关键超参数的设置。

MeKi：基于记忆的专家知识注入实现高效LLM扩展 MeKi: Memory-based Expert Knowledge Injection for Efficient LLM Scaling