北大团队提出 Mamba4Net：跨架构蒸馏，让 LLM 网络知识「瘦身」又「提速」

本文来自北京大学 ICNLab（深圳研究生院），通讯作者为雷凯教授。第一作者林晗 Xia，合作者包括杨明瞻、王晶晶、晏紫薇、任亚坤、余果等。

该团队提出了 Mamba4Net——一个跨架构知识蒸馏框架，把 Transformer 大模型的网络领域知识「压缩」到轻量级的 Mamba 学生模型中。核心思路很简单：LLM 懂很多，但跑在网络设备上太重了。Mamba4Net 只蒸馏跟网络任务有关的那部分知识，同时用一种巧妙的权重复用策略让蒸馏过程更稳定。实验结果相当直观：吞吐量是原来的 3.96 倍，模型大小只剩原来的 5.48%。

为什么要把 LLM 的网络知识「搬」到 Mamba 里

大语言模型在自然语言处理领域大放异彩之后，研究者们自然想到：能不能用 LLM 来做网络管理？答案是肯定的。像 NetLLM 这样的工作已经证明，用一个预训练 LLM 配合 LoRA 微调，就能同时处理视口预测（Viewport Prediction）、自适应码率流媒体（ABR）、集群作业调度（CJS）等多个网络任务——一套模型多路复用，工程成本低，泛化性也还行。

但问题也来了。

Transformer 架构的 LLM 有两个硬伤。第一，自注意力机制的时间复杂度是二次方的。网络数据流大且持续，放在路由器或 AR 头显这种边缘设备上根本跑不动。第二，LLM 里塞了大量跟网络毫无关系的通用知识，几百亿参数中真正用到的只是一小部分，白白浪费存储和算力。

Mamba 架构给了一个新的可能。作为选择性状态空间模型（Selective SSM），Mamba 在捕捉长程依赖时的时间复杂度是线性的——这正是网络场景需要的东西。此前 NetMamba 等工作已经用 Mamba 做网络流量分类，效果和资源消耗都比 Transformer 更友好。

但直接从零训练一个 Mamba 模型来做网络任务，工程成本又回到了原点。所以研究团队的设计思路是：用 LLM 当老师，Mamba 当学生，把网络相关的知识「蒸馏」过去。顺便把模型缩小、把推理加速。

不过，Transformer 和 Mamba 的架构差异很大，传统知识蒸馏方法在「跨架构」场景下训练不稳定、收敛慢——这也是 Mamba4Net 重点解决的问题。

方法：蒸馏 + 权重复用，两条腿走路

Mamba4Net 的整体框架如下图所示，包含三个核心组件：多模态编码器、任务特定的网络头、教师-学生蒸馏框架。

Mamba4Net 框架总览

多模态编码器负责把不同类型的网络输入（视频帧、网络吞吐量序列、作业依赖图等）统一编码为 token 嵌入。比如视口预测用 ViT 处理视频帧，集群调度用 GNN 处理 DAG 图。任务特定的网络头则替代了语言模型的默认输出头，直接把模型 logits 映射到标准化的网络决策格式——这样做可以规避 LLM 的「幻觉」问题。

真正的核心在于教师-学生蒸馏框架的设计。Mamba4Net 提出了两个关键技术。

DKO：面向领域知识的跨架构蒸馏

DKO（Domain Knowledge-Oriented Cross-Heterogeneous Distillation）解决的是「怎么蒸」的问题。

对于监督学习任务（比如视口预测），学生模型的训练目标同时包含两部分：常规的监督损失（和标签比较），加上一个蒸馏损失（和教师模型的输出做 KL 散度）。换句话说，学生不仅要学正确的答案，还要尽量模仿老师的「思维方式」。

对于强化学习任务（比如 ABR 和 CJS），框架采用策略蒸馏的方式：先让教师模型用 LoRA 微调获得网络知识，再通过 KL 散度约束学生模型的策略分布去逼近教师。

一个细节值得提：教师模型本身也是通过 LoRA 微调的，不是全参数微调。这样做的好处是把需要训练的参数量从 din×doutd_{in} \times d_{out} 压缩到了 r×(din+dout)r \times (d_{in} + d_{out})，其中 rr 远小于原始维度。

CWR：跨架构权重复用，让蒸馏更稳更快

跨架构蒸馏有个老大难问题：教师是 Transformer，学生是 Mamba，两者的参数结构完全不同。如果学生模型从随机权重开始训练，蒸馏过程往往不稳定，收敛也慢。

CWR（Cross-Heterogeneous Weight Reusing）的思路很巧妙——它不是随机初始化学生模型，而是把教师的注意力权重「翻译」成 Mamba 的初始参数。

具体怎么做？先把教师的注意力权重矩阵 Q,K,VQ, K, V 做截断奇异值分解（SVD），保留最大的 rr 个奇异值对应的分量。然后取前 dSd_S 行（dSd_S 是学生的嵌入维度，比教师小），得到一个维度匹配的低秩矩阵，直接作为学生模型的初始化权重。

Mamba 和 Transformer 的结构对比及权重复用关系

直觉上理解：教师已经学到了「哪些输入特征重要」的表示，通过 SVD 降维后把这些知识传给学生，学生就不用从零开始摸索了。上图中相同颜色的权重表示不同架构之间的复用关系。

实验结果：又快又小，性能不掉

研究团队在三个网络任务上做了全面评测，教师模型统一使用 Llama2-7B，学生模型是 12 层的混合 Mamba 架构（10 层 Mamba + 2 层 Transformer），嵌入维度从 4096 压缩到 512。

三任务性能对比

在视口预测（VP）任务上，Mamba4Net 的 MAE（平均绝对误差）相比 TRACK 等传统方法降低了 2%～35.9%。在集群作业调度（CJS）任务上，作业完成时间（JCT）缩短了 39.6%～62.3%。在自适应码率流媒体（ABR）任务上，Mamba4Net 的 QoE（体验质量）分数只比之前最强的 NetLLM 低 2.4%——差距很小。

换句话说，在预测精度上 Mamba4Net 和全尺寸 LLM 方法基本持平，有些任务还更优。真正的优势在效率上。

模型大小和吞吐量对比

和 NetLLM 相比，Mamba4Net 的模型大小只有后者的 5.48%，吞吐量却是后者的 3.96 倍。这个数字说明 Mamba 的线性复杂度不是白给的——在同样硬件条件下，Mamba4Net 可以处理更多的网络请求。

消融实验：LLM 知识和权重复用都不可或缺

为了验证 LLM 预训练知识的贡献，研究者去掉了教师模型，让学生直接在网络数据上训练（Mamba4Net-S），和有教师蒸馏的版本（Mamba4Net-F）做对比。

LLM 知识消融实验

结果很说明问题：在 VP 任务上，有 LLM 知识蒸馏的版本 MAE 降低了 38.1%；在 ABR 任务上，QoE 提升了 30.1%。大模型的预训练知识对学生模型的性能提升是实质性的。

CWR 消融实验

CWR 的消融实验同样清晰：使用权重复用初始化的模型（Mamba4Net-C）相比随机初始化（Mamba4Net-D），在不同任务上节省了 10.31%～19.77% 的 GPU 训练时间。CWR 提供了一个更好的起点，让学生模型更快收敛。

总结

Mamba4Net 提供了一条将 LLM 网络知识部署到资源受限设备的可行路径，核心贡献可以归结为三点：

  • DKO 蒸馏框架：同时支持监督学习和强化学习场景，通过 LoRA 微调教师 + KL 散度约束学生的方式，高效提取网络领域知识。
  • CWR 权重复用：通过 SVD 降维将教师的注意力权重映射为学生的初始化参数，解决了跨架构蒸馏训练不稳定的问题，节省 10%～20% 的训练时间。
  • 极致压缩：模型大小压缩到原来的 5.48%，吞吐量提升 3.96 倍，同时性能与全尺寸 LLM 方法持平。

未来，这套「跨架构蒸馏」的思路不仅限于网络领域，在可穿戴设备 QoS 优化、无人机集群调度等资源受限场景中也有应用潜力。

论文链接：https://arxiv.org/abs/2510.17147