[Paper] 合并在不相交数据集上训练的 Kolmogorov‑Arnold 网络

Source: arXiv - 2512.18921v1

概述

论文*“Merging of Kolmogorov‑Arnold networks trained on disjoint datasets”*表明，Kolmogorov‑Arnold 网络（KAN）可以在不同的数据分片上并行训练，然后通过简单的平均步骤进行合并——同时仍然保留 Newton‑Kaczmarz 优化器和分段线性基函数带来的加速优势。这使得 KAN 成为快速、隐私保护的联邦学习以及需要处理海量分布式数据的训练流水线的有力候选方案。

关键贡献

• 证明在不相交子集上训练的 KAN 可以通过朴素的参数平均合并，且不会损失精度。
• 确定 Newton‑Kaczmarz 优化器结合分段线性基函数是当前训练 KAN 的最快配方。
• 提供实证证据表明，将训练集拆分并并行训练可获得超出优化器本身的额外墙钟时间加速。
• 发布完整的开源代码库（训练脚本、合并工具和基准笔记本），以实现可重复性。

方法论

1. 模型选择 – Kolmogorov‑Arnold 网络:
   KANs 是一种最近的神经网络模型，它们用一组单变量函数（“基函数”）对输入的线性组合进行求和，取代了常规的全连接层。它们的结构使得参数在层之间是可加的，这就是在合并各自训练的副本时简单平均能够起作用的原因。

2. 优化 – Newton‑Kaczmarz:
   作者采用了混合的 Newton‑Kaczmarz 方案。Kaczmarz 部分通过迭代投影到超平面上来求解线性子问题（可以看作是随机的、按行的梯度下降）。Newton 校正利用二阶信息对解进行细化，从而在分段线性基函数上实现显著更快的收敛。

3. 在不相交数据上训练:

   • 将完整训练集划分为 k 个不重叠的分片（可以是不同的数据集或随机划分）。
   • 每个分片用于训练一个独立的 KAN 实例，优化器为 Newton‑Kaczmarz。
   • 在固定的 epoch 数后（或每个分片达到局部收敛准则时），将模型参数逐元素平均，得到全局模型。

4. 评估:
   在多个公开的回归和分类任务上进行基准测试（例如 UCI Energy、使用展平特征表示的 CIFAR‑10）。作者比较了三种基线：(i) 单节点使用 Adam 训练，(ii) 单节点使用 Newton‑Kaczmarz 训练，(iii) 所提出的分布式训练加平均流水线。

结果与发现

• Accuracy stays within 0.1 % of the best single‑node baseline, confirming that averaging does not degrade performance.
  准确率保持在最佳单节点基线的 0.1 % 以内，验证了模型平均不会导致性能下降。
• Training time roughly halves when moving from a single Newton‑Kaczmarz run to a 2‑shard setup, and nearly quarters with 4 shards, matching the ideal linear speed‑up predicted by the disjoint‑data assumption.
  训练时间在从单一 Newton‑Kaczmarz 运行切换到 2‑shard 设置时大约减半，使用 4 shard 时几乎缩短至四分之一，符合基于不相交数据假设的理想线性加速。
• The method also shows robustness to heterogeneous data distributions: even when shards are drawn from different domains (e.g., sensor data vs. image features), the merged model still converges to a comparable optimum.
  该方法同样展示了 对异构数据分布的鲁棒性：即使各 shard 来自不同领域（例如传感器数据与图像特征），合并后的模型仍能收敛到相近的最优解。

实际意义

• 联邦学习变得简单：企业可以在边缘设备上部署基于 KAN 的客户端，在本地私有数据上训练，然后在中心服务器上直接对得到的参数进行平均——无需复杂的安全聚合协议。
• 加速模型开发：将海量日志在计算节点之间拆分的数据工程管道，现在可以并行训练 KAN，而无需重写训练循环；唯一额外的步骤是最后的 torch.mean‑式合并。
• 资源受限的环境：由于 Newton‑Kaczmarz 优化器在远少于 Adam 的 epoch 中即可收敛，开发者可以降低 GPU/TPU 使用率，进而降低云端成本。
• 面向表格和分段线性问题的快速原型：KAN 在具有明显阶段变化的回归任务（如金融、物联网传感器校准）上表现出色。所提出的方法让团队能够利用现有的分布式计算集群更快迭代。

限制与未来工作

• 模型类别限制： 平均属性依赖于 KAN 的可加性；它并不能直接迁移到传统的深度 CNN 或 Transformer。
• Newton‑Kaczmarz 步骤的可扩展性： 对于中等规模的 KAN 来说速度很快，但每次迭代的成本会随基函数数量的增加而增长，可能限制非常大规模的部署。
• 异构性处理： 论文的实验使用了相对均衡的分片大小；未来工作可以在分片大小或标签分布差异显著时，探索加权平均或自适应学习率。
• 隐私保证： 简单的平均并未提供正式的差分隐私保护。引入噪声添加机制或安全多方计算将是实现真正隐私保护联邦学习的自然下一步。

如果你想亲自尝试，作者已经发布了一个可直接运行的 Docker 镜像以及一套 Jupyter Notebook，帮助你完成数据划分、使用 Newton‑Kaczmarz 进行训练以及模型合并的全过程。

作者

• Andrew Polar
• Michael Poluektov

论文信息

• arXiv ID: 2512.18921v1
• 分类: cs.LG
• 发布: 2025年12月21日
• PDF: 下载 PDF