MetaCluster: Enabling Deep Compression of Kolmogorov-Arnold Network¶
会议: NeurIPS 2026
arXiv: 2510.19105
领域: 模型压缩
关键词: Kolmogorov-Arnold网络 / 模型压缩 / 权重共享 / 元学习 / 码本量化
一句话总结¶
MetaCluster 提出一种结合元学习流形塑形与权重共享的三阶段 KAN 压缩框架，通过训练轻量元学习器将高维基函数系数约束在低维聚类友好流形上，随后执行系数向量聚类并剪除元学习器进行码本微调，在视觉与方程建模任务中分别取得高达 80× 与 124.1× 的参数量压缩且完全无精度损失。
研究背景与动机¶
Kolmogorov-Arnold 网络（KAN）作为多层感知机（MLP）的强力替代方案，凭借 Kolmogorov-Arnold 表示定理在科学机器学习、复杂方程建模以及计算机视觉领域展现出极佳的非线性拟合能力。与传统 MLP 在每条连接边上使用标量权重的范式不同，KAN 将可学习激活函数置于边上，通过一组基函数（如 B-样条、径向基函数 RBF 或 Gram 多项式）的加权线性组合来参数化每条边的单变量函数。然而，这种每条边引入基系数向量的设计，导致模型参数量和存储空间随着基函数数量成倍膨胀（同拓扑下参数量大约是 MLP 的 \(|\mathbf{w}|\) 倍），严重阻碍了 KAN 向大规模现代神经网络架构的扩展与部署。
权重共享（Weight Sharing）作为一种经典的神经网络压缩手段，通过 K-means 将连续权重聚类为离散码本并存储紧凑索引，在标量参数网络中被证明能有效降低存储。然而，将传统权重共享直接套用到 KAN 上会遭遇根本性失效。其核心矛盾在于：MLP 的权重共享针对的是标量分布，而 KAN 必须对每条边的高维系数向量 \(\mathbf{w} \in \mathbb{R}^{|\mathbf{w}|}\) 进行聚类。在高维向量空间中，受到“维度诅咒”（Curse of Dimensionality）的影响，向量空间极度稀疏，点与点之间的欧氏距离趋于均等化，导致经典聚类算法无法收敛出紧凑而具有代表性的聚类簇；实验显示直接聚类会导致模型分类精度断崖式下跌，即使经历重新微调也难以恢复。
本文的切入角度是：如果高维系数向量在非约束训练下分布过于弥散，我们是否可以在训练期主动诱导其分布收敛到一个低维连续流形上？核心 idea：引入轻量元学习器将低维隐空间嵌入映射为全量基系数向量，在任务损失驱动下强行对系数空间进行流形塑形（Manifold Shaping），使得高维系数高度聚类友好；完成向量级聚类后彻底剪除元学习器，仅保留紧凑码本与索引并在推理期实现零额外开销。
方法详解¶
整体框架¶
MetaCluster 是一个紧凑的三阶段压缩框架，由「元学习流形塑形训练」、「系数向量级 K-means 聚类」与「剪除元学习器及码本微调」组成。在训练阶段，网络不直接优化各条边的高维系数，而是为每条边分配一个极低维的嵌入向量（如 1 维或 2 维），由一个共享的小型元学习器将其解码为全量系数向量，使系数被强制约束在低维子流形上。训练完成后，对生成的系数向量执行 K-means 聚类并构建索引映射查找表。最后，元学习器与嵌入向量被完全舍弃，网络直接使用码本质心替代原始系数，并通过短周期的质心微调恢复残余精度损失。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["低维边嵌入向量<br/>z_i ∈ R^{d_emb}"] --> B["流形塑形元学习器<br/>M_θ(z_i) → w_i ∈ R^{|w|}"]
B --> C["前向计算与任务损失联合优化<br/>将高维系数约束至低维光滑流形"]
C --> D["系数向量级 K-means 聚类<br/>生成码本质心 c_k 与索引映射 I_j"]
D --> E["剪除元学习器与嵌入<br/>构建紧凑索引查找表"]
E --> F["轻量码本微调<br/>固定拓扑仅更新质心向量"]
F --> G["最终紧凑 KAN 模型<br/>零推理额外开销 + 80x-124x 显存压缩"]关键设计¶
1. 流形塑形元学习器：将高维系数约束至低维流形空间
针对高维系数向量空间由于距离均等化而无法有效聚类的痛点，MetaCluster 引入了一个轻量级两层 MLP 作为全局共享的元学习器 \(M_\theta\)。对于网络中第 \(i\) 条边的低维连续嵌入向量 \(\mathbf{z}_i \in \mathbb{R}^{d_{\text{emb}}}\)（通常 \(d_{\text{emb}} \in \{1, 2\}\)），元学习器将其映射到完整的基系数向量 \(\mathbf{w}_i \in \mathbb{R}^{|\mathbf{w}|}\)：
其中隐层维度 \(d_{\text{hidden}}=32\)，\(\sigma(\cdot)\) 为 ReLU 激活函数。元学习器参数与嵌入向量在主任务损失（如交叉熵或均方误差）的反向传播下联合优化 \(\alpha\) 个 epoch。通过 t-SNE 可视化可以清晰看到：未加元学习器的原始 KAN 系数呈现出杂乱弥散的云团，而在 \(d_{\text{emb}}=1\) 时系数严格收敛于一维流线，在 \(d_{\text{emb}}=2\) 时则形成规则平滑的二维薄片流形。这种结构化的连续流形迫使任务相关的参数变异性高度集中，消除了高维空间的稀疏性，从而赋予了高维系数天然的易聚类特性。
2. 系数向量级 K-means 聚类与紧凑码本索引：利用向量特性摊薄存储开销
在完成流形塑形后，网络收集所有边由元学习器生成的系数向量集合 \(\{M_\theta(\mathbf{z}_j)\}\)，并在向量空间 \(\mathbb{R}^{|\mathbf{w}|}\) 中运行 K-means 算法，将其聚类为 \(k\) 个簇中心 \(\mathcal{C} = \{\mathbf{c}_1, \dots, \mathbf{c}_k\}\)：
每个聚类质心由簇内向量均值更新：\(\mathbf{c}_i = \frac{1}{|C_i|} \sum_{M_\theta(\mathbf{z}_j) \in C_i} M_\theta(\mathbf{z}_j)\)。接着为全网每条边记录其所属簇的紧凑索引 \(I_j = \arg\min_{i \in \{1,\dots,k\}} \|M_\theta(\mathbf{z}_j) - \mathbf{c}_i\|_2^2\)。
该设计的数学精妙之处在于向量聚类对存储开销的摊薄效应（Amortization Effect）。传统 MLP 对标量权重做聚类时，压缩比为 \(r_{\text{MLP}} = \frac{nb}{n \log_2(k) + kb}\)，其分母受制于每条连接必需记录的标量索引项 \(n \log_2(k)\)。而在 KAN 中，每条边原本包含 \(|\mathbf{w}|\) 个标量参数，总比特数为 \(n |\mathbf{w}| b\)；聚类后，每条边依然仅需一个 \(\log_2(k)\) 比特的索引，而每个质心存储 \(|\mathbf{w}|\) 个标量。压缩比变为：
分母中的索引开销项被除以了基函数维度 \(|\mathbf{w}|\)！这意味着基函数阶数越高、每条边原本越臃肿，权重共享带来的相对压缩增益就越显著，使得 KAN 相比 MLP 能从聚类中获得大得多的压缩倍率。
3. 剪除元学习器与码本微调：零推理开销恢复网络精度
已有基于超网络（Hypernetworks）的方法（如 MetaKAN）在推理时依然需要保留生成器并进行动态计算，带来了计算与访存负担。MetaCluster 在聚类完成并固定索引映射 \(I\) 后，彻底丢弃元学习器 \(M_\theta\) 和所有潜在嵌入向量 \(\mathbf{z}\)。此时模型退化为纯粹的“码本 + 查找表”结构。由于离散质心是对流形连续系数的最近邻近似，难免存在残余拟合误差。因此，框架固定索引分配矩阵 \(I\)，仅针对离散码本质心 \(\mathbf{C} = \{\mathbf{c}_1, \dots, \mathbf{c}_k\}\) 进行简短的端到端任务损失微调，微调周期 \(\beta\) 远少于初始训练周期（\(\beta \ll \alpha\)，通常仅需 5-10 个 epoch）。微调使质心直接针对任务目标进行微调补偿，完全抹平了聚类造成的精度下降，而在推理部署阶段没有任何额外的动态计算或内存寻址负担。
一个完整示例¶
以一个包含两层隐层、每层连接数为 \(n=1024\) 的 B-样条 KAN 为例，样条阶数 \(k_{\text{spline}}=3\)、网格数 \(G=5\)，每条边系数维度 \(|\mathbf{w}| = G + k_{\text{spline}} + 1 = 9\)。若采用 32 位浮点数（\(b=32\)）： 1. 训练前：网络原本需要存储 \(1024 \times 9 \times 32 \text{ bits} = 294,912 \text{ bits} \approx 36.86 \text{ KB}\)。在 MetaCluster 第一阶段，为 1024 条边各分配 1 个标量嵌入 \(z_i \in \mathbb{R}^1\)，通过隐层为 32 的元学习器映射生成 9 维系数向量，共同训练 50 个 epoch。 2. 聚类期：设定质心数 \(k=16\)。聚类算法将这 1024 个 9 维向量归纳为 16 个代表质心。每条边仅存储一个 4-bit 索引（\(\log_2(16)=4\)），索引表共占用 \(1024 \times 4 \text{ bits} = 4096 \text{ bits}\)；码本存储 16 个 9 维浮点向量，共占用 \(16 \times 9 \times 32 \text{ bits} = 4608 \text{ bits}\)。 3. 剪除与微调：删去元学习器和 1024 个标量嵌入，总存储量骤降至 \(4096 + 4608 = 8704 \text{ bits} \approx 1.06 \text{ KB}\)，理论压缩比高达 \(34.7×\)。经过 5 个 epoch 的微调，质心完成任务补偿，最终模型以极小内存无损执行前向传播。
损失函数 / 训练策略¶
- 联合训练阶段：使用主任务损失函数 \(\mathcal{L}_{\text{task}}\)（图像分类采用交叉熵损失，方程建模采用 MSE 损失）。优化器采用 AdamW，分别设置元学习器学习率（全连接层 \(lr_h \in [1\times 10^{-4}, 1\times 10^{-3}]\)）与嵌入向量学习率（\(lr_e \in [1\times 10^{-3}, 1\times 10^{-2}]\)），训练至多 50 个 epoch（卷积 KAN 为 150 个 epoch），结合早停机制（patience = 10）。
- 聚类算法：标准 K-means，初始化采用 k-means++，在 CPU/GPU 上运行至收敛（耗时通常小于 0.1 秒）。
- 码本微调阶段：固定所有边的索引分配向量 \(I\)，仅将码本参数 \(\mathbf{C}\) 加入优化器，使用微调学习率 \(lr_c = 1\times 10^{-4}\)（卷积为 \(1\times 10^{-5}\)），在任务损失下快速微调 5 个 epoch（方程建模微调 10 个 epoch，早停 patience = 3）。
实验关键数据¶
主实验¶
论文在全连接与卷积架构下，对 B-spline（KAN）、RBF（FastKAN）及 Gram 多项式（GramKAN）三种基函数进行了全面测试，覆盖 MNIST、CIFAR-10 及 CIFAR-100 数据集。全连接模型聚类数设为 \(k=16\)，卷积模型设为 \(k=256\)。
| 架构 / 模型 | 参数存储 (KB) MNIST / C10 / C100 |
MNIST 准确率 (%) | CIFAR-10 准确率 (%) | CIFAR-100 准确率 (%) | 相比原始 KAN C10 压缩比 |
|---|---|---|---|---|---|
| 全连接 (FC) | |||||
| KAN (B-spline) | 1031.44 / 3064.95 / 3177.45 | 95.37 ± 0.11 | 47.52 ± 0.16 | 17.80 ± 0.10 | 1.0× (Baseline) |
| MetaKAN (超网络) | 141.32 / 410.82 / 422.08 | 96.36 ± 0.09 | 46.28 ± 0.39 | 20.17 ± 0.24 | 7.5× |
| ClusterKAN (直接聚类) | 13.84 / 38.34 / 39.75 | 61.37 ± 2.77 | 27.92 ± 2.88 | 6.20 ± 0.31 | 79.9× (精度崩塌) |
| ClusterKAN + fine-tune | 13.84 / 38.34 / 39.75 | 92.66 ± 0.08 | 43.71 ± 0.27 | 11.04 ± 0.64 | 79.9× (难以恢复) |
| MetaClusterKAN (Ours) | 13.84 / 38.34 / 39.75 | 96.06 ± 0.10 | 45.47 ± 0.44 | 18.97 ± 0.15 | 79.9× |
| MetaClusterKAN + ft | 13.84 / 38.34 / 39.75 | 96.41 ± 0.10 | 45.95 ± 0.33 | 19.16 ± 0.17 | 79.9× (无损恢复) |
| FastKAN (RBF) | 900.56 / 2676.83 / 2778.08 | 96.78 ± 0.08 | 49.10 ± 0.29 | 20.56 ± 0.09 | 1.0× |
| MetaFastKAN | 108.09 / 316.35 / 327.60 | 95.51 ± 0.12 | 47.40 ± 0.46 | 19.14 ± 0.40 | 8.5× |
| ClusterFastKAN + ft | 20.54 / 57.30 / 58.71 | 86.24 ± 1.08 | 38.89 ± 0.35 | 11.35 ± 0.34 | 46.7× (-10.2% 精度) |
| MetaClusterFastKAN + ft | 20.54 / 57.30 / 58.71 | 95.59 ± 0.18 | 47.60 ± 0.37 | 19.05 ± 0.47 | 46.7× (保持 SOTA) |
| GramKAN (多项式) | 497.46 / 1477.48 / 1534.43 | 96.76 ± 0.08 | 49.57 ± 0.43 | 21.89 ± 0.09 | 1.0× |
| ClusterGramKAN + ft | 13.96 / 38.46 / 40.58 | 95.40 ± 0.15 | 45.81 ± 0.66 | 13.98 ± 0.44 | 38.4× (-7.9% 精度) |
| MetaClusterGramKAN + ft | 14.15 / 38.65 / 40.76 | 95.32 ± 0.12 | 48.47 ± 0.30 | 18.96 ± 0.25 | 38.2× (保持性能) |
| 卷积结构 (Conv) | |||||
| KANConv | 13634.0 / 13654.3 / 13744.6 | 99.46 ± 0.02 | 70.66 ± 2.89 | 23.99 ± 0.98 | 1.0× (Baseline) |
| MetaKANConv | 3048.40 / 3052.90 / 3143.25 | 99.32 ± 0.05 | 68.36 ± 2.84 | 35.60 ± 1.78 | 4.5× |
| ClusterKANConv + ft | 429.95 / 430.51 / 520.87 | 98.89 ± 0.09 | 61.80 ± 0.34 | 22.65 ± 0.81 | 31.7× (-8.9% 精度) |
| MetaClusterKANConv + ft | 434.77 / 435.34 / 525.70 | 99.35 ± 0.05 | 66.85 ± 3.11 | 35.11 ± 1.79 | 31.4× (+11.1% C100) |
| FastKANConv | 13632.1 / 13652.4 / 13742.7 | 99.11 ± 0.09 | 76.30 ± 0.11 | 46.36 ± 0.38 | 1.0× |
| ClusterFastKANConv + ft | 428.03 / 428.61 / 518.97 | 96.10 ± 0.09 | 57.45 ± 1.53 | 20.15 ± 0.70 | 31.9× (-18.8% 精度) |
| MetaClusterFastKANConv + ft | 427.97 / 428.55 / 518.91 | 99.08 ± 0.06 | 69.10 ± 1.10 | 39.26 ± 1.10 | 31.9× (大幅优于直接聚类) |
| GramKANConv | 7586.47 / 7597.72 / 7688.07 | 99.40 ± 0.06 | 78.56 ± 1.14 | 48.25 ± 0.59 | 1.0× |
| ClusterGramKANConv + ft | 418.93 / 419.49 / 509.85 | 99.27 ± 0.08 | 75.74 ± 0.63 | 44.01 ± 0.42 | 18.1× |
| MetaClusterGramKANConv + ft | 418.81 / 419.38 / 509.73 | 99.47 ± 0.02 | 80.85 ± 0.41 | 51.20 ± 0.67 | 18.1× (反超 Baseline +2.3%) |
消融实验¶
1. 高维科学方程建模性能与极高压缩比验证
论文在 1000 维输入空间方程拟合任务上评测 MetaCluster，函数定义为：\(f_1(x)=\exp(\frac{1}{n}\sum \sin^2(\frac{\pi x}{2}))\)，\(f_2(x)=\sum(x^2+x^3)\)，\(f_3(x)=\exp(-\frac{1}{n}\sum x^2)\)，聚类簇数仅设为 \(k=4\)。
| 模型配置 | 内存占用 (KB) | \(f_1(x)\) 均方误差 (MSE) | \(f_2(x)\) 均方误差 (MSE) | \(f_3(x)\) 均方误差 (MSE) | 相对压缩比 |
|---|---|---|---|---|---|
| KAN (未压缩基线) | 1303.6 | \(5.06 \times 10^{-5}\) | \(2.70 \times 10^{-1}\) | \(5.30 \times 10^{-7}\) | 1.0× |
| MetaKAN (超网络生成) | 181.6 | \(3.05 \times 10^{-5}\) | \(2.41 \times 10^{0}\) | \(2.76 \times 10^{-6}\) | 7.2× |
| ClusterKAN (直接聚类) | 10.5 | \(9.85 \times 10^{-1}\) | \(1.66 \times 10^{2}\) | \(1.27 \times 10^{-1}\) | 124.1× (误差崩溃) |
| ClusterKAN + fine-tune | 10.5 | \(2.24 \times 10^{-4}\) | \(8.92 \times 10^{-1}\) | \(7.82 \times 10^{-6}\) | 124.1× (残余误差高) |
| MetaClusterKAN (未微调) | 10.5 | \(9.03 \times 10^{-4}\) | \(1.58 \times 10^{0}\) | \(3.75 \times 10^{-4}\) | 124.1× |
| MetaClusterKAN + ft | 10.5 | \(2.70 \times 10^{-5}\) | \(5.04 \times 10^{-1}\) | \(7.55 \times 10^{-7}\) | 124.1× (拟合质量比肩未压缩) |
2. 基函数系数维度与元学习隐空间维度消融
| 配置维度指标 | 实验设定 | 聚类前准确率 (%) | 聚类后准确率 (%) | 微调后准确率 (%) | 相对变化 (%Chg) | 压缩后显存 (KB) |
|---|---|---|---|---|---|---|
| 系数维度消融 (FC) | Coeff = 5 | 47.41 | 47.46 | 47.35 | +0.42% | 57.30 |
| (MetaFastKAN, C10) | Coeff = 8 | 47.62 | 46.96 | 47.80 | +0.02% | 57.68 |
| Coeff = 10 | 46.91 | 47.04 | 47.49 | +0.95% | 57.93 | |
| Coeff = 15 | 45.88 | 45.15 | 45.74 | +0.62% | 58.55 | |
| Coeff = 20 | 43.57 | 43.75 | 44.31 | -0.17% | 59.18 | |
| 隐空间维度消融 (FC) | \(d_{\text{emb}} = 1\) | 47.40 | 46.96 | 47.60 | Baseline | 57.30 |
| (MetaFastKAN, C10) | \(d_{\text{emb}} = 2\) | 48.19 | 45.48 | 47.29 | -0.90% | 57.30 |
| \(d_{\text{emb}} = 3\) | 47.28 | 39.11 | 45.20 | -4.40% | 57.30 | |
| \(d_{\text{emb}} = 4\) | 48.66 | 38.28 | 44.63 | -8.28% | 57.30 |
关键发现¶
- 直接聚类遭遇高维瓶颈：未经流形塑形的朴素聚类（ClusterKAN）在所有任务上均遭遇严重性能崩溃（如 CIFAR-100 上从 21.89% 跌落至 6.34%），即使经过微调，也仅能恢复至 13.98%，落后未压缩基线近 8 个百分点。与之对比，MetaClusterKAN 在微调后达到 18.96%，性能无损。
- 低维嵌入是聚类成功的决定性因素：消融表明，随着元学习器嵌入维度 \(d_{\text{emb}}\) 从 1 增大至 4，聚类后初始准确率显著衰退（全连接网络从 46.96% 暴跌至 38.28%）。这严格印证了流形约束越紧凑（如 1D 或 2D），生成系数在高维空间中的分布结构越紧致，聚类损失越小。
- 高维方程拟合取得 124.1× 极值压缩：在 1000 维方程回归任务中，仅需 4 个质心向量（10.5 KB），MetaClusterKAN 微调后的 MSE 在 \(f_1(x)\) 上甚至优于原始未压缩 KAN（\(2.70 \times 10^{-5}\) vs \(5.06 \times 10^{-5}\)），验证了低维子流形正则化在科学计算中的去噪与泛化增益。
- Hessian 引导量化在 KAN 上并不适用：对比大模型主流的向量量化基线 GPTVQ（Hessian-aware），在 CIFAR-100 上微调后仅取得 8.96% 的精度（而 MetaCluster 达 19.16%），说明曲率加权在 KAN 复杂非线性基函数上容易产生局部过拟合或扰动。
亮点与洞察¶
- 将生成式元学习转化为量化友好正则项：巧妙之处在于打破了“超网络必须在推理阶段运行”的思维定势，仅把超网络当作训练阶段塑造参数几何形态的“模具”，聚类后立即卸磨杀驴，实现了零推理开销与极致压缩的结合。
- 洞察 KAN 参数结构天然享有向量摊薄红利：论文从数学公式上清晰指出，KAN 因每条边携带多维基函数系数，使其比传统标量 MLP 更加适合权重共享——每个质心承载一整组基系数，极大地摊薄了离散索引带来的存储占比。
- 对各类基函数具备普适鲁棒性：无论是局部支撑的 B-样条、全局平滑的 RBF，还是正交的多项式基函数，MetaCluster 均能在几乎不调超参的情况下平滑迁移并取得同等压缩比。
局限与展望¶
- 作者承认的局限：聚类簇数 \(k\) 与隐空间维度 \(d_{\text{emb}}\) 仍需人工设定先验（如 FC 取 16，Conv 取 256）；此外三阶段流程相比端到端单次训练增加了微调步骤。
- 潜在改进方向：本方法与低比特权重整数量化（Bit-width Quantization）在原理上完全正交互补。未来可进一步探索在码本质心矩阵上叠加 4-bit/8-bit 定点量化，或结合自适应非均匀网格剪枝，进一步压榨边缘部署下的显存边界。
相关工作与启发¶
- vs MetaKAN (Zhao et al., 2025)：MetaKAN 通过超网络在推理时根据任务嵌入动态生成系数，虽然减少了存储，但引入了不可忽视的前向推理延迟与计算开销；MetaCluster 仅在训练阶段借助元学习器塑形，推理期完全剥离元网络，速度与普通 KAN 一致且压缩率大幅提高（从约 7.5× 跃升至 79.9×）。
- vs Deep Compression (Han et al., 2015)：经典 Deep Compression 面向标量权重做标量 K-means，其索引项难以被高维摊薄；MetaCluster 针对 KAN 的向量特性定制了向量级聚类，并解决了高维聚类维度诅咒难题。
- vs GPTVQ (Van Baalen et al., 2024)：GPTVQ 依赖二阶 Hessian 矩阵引导向量量化，在大语言模型上表现优异，但在 KAN 上由于非线性基函数曲率剧烈变化而失效；MetaCluster 证明直接流形塑形比二阶曲率加权更适用于 KAN 压缩。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次提出针对 Kolmogorov-Arnold 网络的高维系数向量权重共享方案，构思极其精巧。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 种主流基函数、全连接与卷积架构、3 个视觉数据集以及 1000 维方程建模，消融极其详尽。
- 写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，数学推导与直觉可视化论证严密，结构规范。
- 价值: ⭐⭐⭐⭐⭐ 彻底打破了 KAN 内存难以扩展的致命短板，为科学机器学习与边缘智能的硬件落地铺平了道路。