知识卸载：将大语言模型分解为稀疏骨干网络与记忆模块

# 知识卸载：将大语言模型分解为稀疏骨干网络和记忆模块
来源: https://arxiv.org/html/2605.29075
Karim Galliamov1Rochelle Choenni1Ivan Titov2,1 1阿姆斯特丹大学2爱丁堡大学 karim\.galliamov@student\.uva\.nl\{r\.m\.v\.k\.choenni, titov\}@uva\.nl

###### 摘要

大语言模型（LLM）将通用能力和领域特定知识编码在同一组参数中。我们探究能否对这种能力进行重组：将广泛适用的计算保留在共享骨干网络中，同时将专门化知识迁移到外部记忆模块中。我们提出*知识卸载*（KOFF）框架，用于将预训练LLM分解为稀疏共享骨干网络和领域特定记忆。从冻结的基础模型出发，我们联合学习结构化剪枝掩码和轻量级恢复模块，这些模块以LoRA适配器和可学习键值缓存的形式实现。在Llama和Qwen（3B到8B）模型上的实验表明，非平凡的能力可以从共享骨干网络中移出，而不会造成模型能力的大幅下降。在大约12%的全局稀疏度下，KOFF保留了未剪枝模型的大部分性能，而在无记忆模块的情况下对相同的冻结模型进行剪枝则会导致性能急剧下降。消融实验显示，LoRA和可学习KV记忆是互补的。专门化分析表明，学习到的分解在内部是有意义的：语言特定神经元被优先移除，而语言通用神经元主要保留在骨干网络中。这些结果表明，知识可以在共享核心和可交换的外部记忆之间重新分配。

知识卸载：将大语言模型分解为稀疏骨干网络和记忆模块

Karim Galliamov1Rochelle Choenni1Ivan Titov2,11阿姆斯特丹大学2爱丁堡大学 karim\.galliamov@student\.uva\.nl\{r\.m\.v\.k\.choenni, titov\}@uva\.nl

## 1 引言

大语言模型（LLM）被期望能在多种场景中发挥作用：一个模型可能需要服务多种语言的用户，并在不同领域进行推理。这些用例对模型能力提出了高要求。有些知识具有广泛适用性，支持通用推理和语言能力；而另一些知识则高度专门化，仅与特定主题或语言相关。尽管存在这种多样性，LLM缺乏明确的管理机制来分配不同类型知识的能力。所有信息，无论是通用的还是高度专门化的，都存储在同一组参数中。然而，先前的研究表明，不同子集的神经元专门处理不同功能，包括语言和领域特定的处理（Tang等人，2024 (https://arxiv.org/html/2605.29075#bib.bib4)；Hendy等人，2022 (https://arxiv.org/html/2605.29075#bib.bib5)；Foroutan等人，2022 (https://arxiv.org/html/2605.29075#bib.bib3)），这表明对于给定的输入，只有部分参数是相关的。这提出了一个基本问题：我们能否重组模型能力，使广泛适用的知识保留在共享核心中，而将专门化知识迁移到外部模块中？

参见图注

图1：知识卸载（KOFF）将预训练LLM分解为稀疏骨干网络和可分离的可交换记忆模块。

一种自然的方法是向LLM增加外部记忆。先前的工作探索了检索增强生成（Lewis等人，2020 (https://arxiv.org/html/2605.29075#bib.bib65)）、可学习KV缓存（Eyuboglu等人，2025 (https://arxiv.org/html/2605.29075#bib.bib86)）和参数高效适配器（Hu等人，2021 (https://arxiv.org/html/2605.29075#bib.bib70)）作为扩展模型能力的机制。这些方法允许将专门化知识存储在主参数之外，并在需要时访问，从而减轻共享模型的负担。然而，这类方法通常假设一个固定的骨干网络，并不决定哪些内容应保留在模型中、哪些应外部化。因此，它们解决的是如何添加记忆的问题，而不是如何在共享计算和专门化知识之间重新定位能力。相比之下，结构化剪枝方法通过移除被认为不那么重要的参数，提供了一种重塑模型能力的方式（Ma等人，2023 (https://arxiv.org/html/2605.29075#bib.bib51)）。然而，剪枝方法旨在生成一个服务于所有输入的单一压缩网络，迫使不同领域的专门化知识在有限的能力中竞争。

在这项工作中，我们提出*知识卸载*（KOFF），一个将记忆增强建模和结构化剪枝相统一的框架，通过学习如何在共享计算和外部记忆之间重新定位能力。KOFF不是将模型压缩成一个更小的单一网络，而是将预训练LLM分解为通用目的骨干网络和领域特定记忆模块。它联合学习两个具有不同角色的组件：(1) 一个保留广泛有用信息的共享稀疏骨干网络，以及 (2) 存储和恢复领域特定知识的轻量级*记忆模块*。

具体来说，我们使用基于Hard Concrete门控（Louizos等人，2018 (https://arxiv.org/html/2605.29075#bib.bib57)）的结构化剪枝机制来决定哪些神经元保留在共享骨干网络中。同时，记忆模块被训练来通过两种恢复机制补偿被移除的能力：(a) 一个在参数空间操作的LoRA适配器（Hu等人，2022 (https://arxiv.org/html/2605.29075#bib.bib69)），以及 (b) 一个注入到注意力机制中的可学习键值（KV）缓存（Eyuboglu等人，2025 (https://arxiv.org/html/2605.29075#bib.bib86)），充当外部化的知识存储。通过联合训练这些组件，模型可以以协调的方式重组能力：骨干网络保留共享功能，而记忆模块专门处理已卸载的内容。在推理时，我们使用一个小型的训练分类器来为给定输入选择要附加到骨干网络的记忆模块。

我们在Llama和Qwen（3B到8B）模型上评估知识卸载，使用主题和语言两种领域。实验表明，能力可以从共享骨干网络中移出而不会造成模型能力的大幅下降：在大约12%的全局稀疏度下，卸载保留了未剪枝模型的大部分性能，而在无记忆模块的情况下对相同的冻结模型进行剪枝则会导致性能急剧下降。消融实验显示，LoRA适配器和可学习KV缓存是互补的。最后，专门化分析表明，学习到的分解在内部是有意义的：语言特定神经元被优先移除，而语言通用神经元被保护，骨干网络保留了原始的表示几何结构。据我们所知，这是首次将剪枝视为学习记忆分配的机制的工作。

## 2 相关工作

##### 记忆增强模型。

先前的工作通过将知识外部化到激活空间或参数空间来扩展LLM的能力。前缀和提示微调将可学习表示注入到模型的计算中（Li和Liang，2021 (https://arxiv.org/html/2605.29075#bib.bib66)；Lester等人，2021 (https://arxiv.org/html/2605.29075#bib.bib67)），有效地将任务特定知识存储在激活空间而非基础模型参数中。互补的方法，如LoRA（Hu等人，2022 (https://arxiv.org/html/2605.29075#bib.bib69)），则引入轻量级参数更新，使模型能够将任务特定知识存储在参数空间中。我们从这些工作中汲取灵感，让模型学习如何将领域特定知识卸载到这样的外部记忆中。与我们的工作最接近的是参数记忆方法，如记忆网络（Sukhbaatar等人，2015 (https://arxiv.org/html/2605.29075#bib.bib64)）和层次参数记忆（Pouransari等人，2025 (https://arxiv.org/html/2605.29075#bib.bib68)），它们与模型一起或从头开始学习结构化记忆组件。相反，我们从预训练的密集模型开始，对其剪枝，并使用记忆模块来恢复在剪枝过程中被移除的领域特定知识。与先前专注于冻结模型适应的工作不同，我们的方法将记忆视为协同适应的剪枝-恢复系统中的恢复机制：LoRA适配器提供参数空间记忆，而KV缓存充当激活空间记忆。

##### 模型剪枝。

最近的研究表明，LLM可以通过在训练后应用的结构化剪枝方法进行压缩。流行的方法如SliceGPT（Ashkboos等人，2024 (https://arxiv.org/html/2605.29075#bib.bib49)）、FASP（Hu等人，2025 (https://arxiv.org/html/2605.29075#bib.bib50)）和LLM-Pruner（Ma等人，2023 (https://arxiv.org/html/2605.29075#bib.bib51)）移除了整个神经元或通道，同时保留了基础模型的大部分性能。后者还推广了“先剪枝后适应”范式，即使用参数高效微调方法（如LoRA）来后续恢复剪枝后的模型。然而，这些方法产生具有固定质量-效率权衡的单一模型。我们则使用结构化剪枝作为重组模型能力的机制。通过学习共享的稀疏骨干网络和领域特定记忆模块，我们允许对于不同的输入恢复不同的被剪枝知识。这重新定义了剪枝，从纯粹的压缩操作转变为更广泛的记忆分配策略的组成部分。

##### 条件计算和专家模型。

我们的工作也与稠密到MoE转换和专家LM方法相关，包括MoEfication（Zhang等人，2022 (https://arxiv.org/html/2605.29075#bib.bib82)）、upcycling（Komatsuzaki等人，2023 (https://arxiv.org/html/2605.29075#bib.bib83)）、Branch-Train-Merge（BTM）（Li等人，2022 (https://arxiv.org/html/2605.29075#bib.bib84)）和Branch-Train-MiX（BTX）（Sukhbaatar等人，2024 (https://arxiv.org/html/2605.29075#bib.bib85)）。这些方法通过构建、训练或路由专家计算来引入条件专门化。相反，我们涉及剪枝的目标使得能力分配本身成为学习的对象：从预训练的密集模型开始，我们学习哪些内容应保留在共享骨干网络中，哪些可以外部化到领域特定记忆模块中。在架构上，我们的模块不是路由的FFN专家，而是附加到共享骨干网络上的记忆组件（KV缓存或LoRA）。

## 3 方法

### 3.1 总体框架

KOFF依赖于角色的分离。共享骨干网络应保留跨领域广泛适用的计算，而领域模块则为特定领域的信息提供能力。我们通过联合训练剪枝掩码和领域模块来学习这种分离。这使剪枝目标能够访问剪枝后同样可用的恢复机制：可以通过记忆模块恢复的能力不需要保留在共享骨干网络中，而难以恢复或跨多个领域有用的能力则应保持共享。每个领域模块Md\mathcal{M}_{d}包含两种互补的记忆形式：一个参数空间组件，通过LoRA实现；以及一个激活空间组件，实现为可学习键值记忆。这些组件在下一小节中描述，并且在消融实验（第4.4节）中均显示出益处。

##### 学习目标。

设fθf_{\theta}为原始预训练模型，带有冻结参数θ\theta。我们学习共享的结构化剪枝掩码zz，以及领域模块{Md}d∈D\{\mathcal{M}_{d}\}_{d\in\mathcal{D}}。掩码zz跨领域共享，而每个模块Md\mathcal{M}_{d}是领域特定的。我们将fθ,zf_{\theta,z}记为带掩码的骨干网络，将fθ,z,Mdf_{\theta,z,\mathcal{M}_{d}}记为附加领域dd模块后得到的模型。

我们的目标是保持性：卸载后的模型应保留原始模型的行为，同时将可恢复的能力迁移到外部记忆中。因此，我们使用来自fθf_{\theta}的蒸馏，而非标准的下一词元交叉熵损失，既在领域数据上，也在通用保持数据上。在高层面上，目标函数为

minz,{Md}\displaystyle\min_{z,\{\mathcal{M}_{d}\}}∑d∈DEx∼Dd[Ldistill(fθ,z,Md,fθ;x)]\displaystyle\sum_{d\in\mathcal{D}}\;\mathbb{E}_{x\sim\mathcal{D}_{d}}\bigl[\mathcal{L}_{\mathrm{distill}}(f_{\theta,z,\mathcal{M}_{d}},f_{\theta};x)\bigr] (1)+λΩ(z)+μLretain(z,{Md}).\displaystyle+\;\lambda\,\Omega(z)+\mu\,\mathcal{L}_{\mathrm{retain}}(z,\{\mathcal{M}_{d}\}).

第一项训练每个领域模块在其自身的领域数据上恢复教师模型的行为，同时更新相同的共享掩码zz。惩罚项Ω(z)\Omega(z)鼓励共享骨干网络使用更少的活跃单元；在实践中，我们将其实现为对结构化Hard Concrete门控（第3.3节）的期望L0L_{0}惩罚。

保持项在通用领域数据上应用相同的蒸馏损失：

Lretain\displaystyle\mathcal{L}_{\mathrm{retain}}(z,{Md})=\displaystyle(z,\{\mathcal{M}_{d}\})=(2)Ex∼Dret,d~∼Unif(D)[Ldistill(fθ,z,Md~,fθ;x)].\displaystyle\mathbb{E}_{x\sim\mathcal{D}_{\mathrm{ret}},\;\tilde{d}\sim\mathrm{Unif}(\mathcal{D})}\Bigl[\mathcal{L}_{\mathrm{distill}}(f_{\theta,z,\mathcal{M}_{\tilde{d}}},f_{\theta};x)\Bigr].

这阻止共享掩码移除在领域模块训练分布之外所需的能力。换句话说，领域批次教会模块哪些内容可以在本地恢复，而保持批次则推动共享骨干网络保留那些应保持广泛可用的能力。

##### 推理。

带掩码的骨干网络不需要作为动态门控模型来服务。训练后，zz被固定并物化：被剪枝的神经元及其关联的权重行和列被移除，得到一个紧凑的密集骨干网络。在推理时，领域标签不被假设为已知。因此，我们在用于训练记忆的相同领域标记数据上训练一个轻量级路由器rϕ(d∣x)r_{\phi}(d\mid x)。给定输入xx，路由器选择

d^=argmaxd∈Drϕ(d∣x),\hat{d}=\arg\max_{d\in\mathcal{D}}r_{\phi}(d\mid x),

然后使用fθ,z,Md^f_{\theta,z,\mathcal{M}_{\hat{d}}}进行生成。路由器仅选择要附加哪个记忆；它不改变共享掩码或紧凑骨干网络。我们在附录A.1中进一步描述它。

### 3.2 领域特定记忆模块

每个模块Md\mathcal{M}_{d}通过两种互补机制向共享稀疏骨干网络添加领域特定能力：一个LoRA适配器和一个可学习KV缓存。

##### LoRA适配器。

对于每个领域（或语言）d∈Dd\in\mathcal{D}，我们引入一个LoRA适配器（Hu等人，2022）。适配器为剪枝后的骨干网络提供轻量级的参数空间修正。在我们的实验中，LoRA适配器针对选定的投影矩阵，并且可以放置在所有Transformer块中，或仅在连续的子集层中，例如顶部NN层。我们在第4.4节讨论选择。

##### 可学习KV缓存。

除了LoRA适配器，每个领域dd在模块活跃的每个注意力层配备一组NkvN_{kv}个可学习键值记忆表示。如Eyuboglu等人（2025