稀疏激活的混合专家算子 Transformer 用于大规模 PDE 预训练
机器学习
2025-11-03 v2 数值分析
数值分析
摘要
预训练在解决神经算子 PDE 问题时已证明有效，但由于 PDE 数据集中方程类型的异构性，导致混合训练误差较大。此外，密集预训练模型通过增加网络宽度或深度来扩大参数规模，推理成本显著。为此，我们提出了一种新的混合专家预训练算子 Transformer（MoE-POT），一种稀疏激活架构，在控制推理成本的同时高效扩展参数规模。具体而言，我们的模型采用层级路由-门控网络，在推理时动态选择 4 个路由专家出 16 个专家网络，以便聚焦方程特定特征。同时，我们还集成 2 个共享专家，以捕捉 PDE 的通用属性并减少路由专家之间的冗余。最终输出由所有激活专家的结果的加权平均计算。我们在 6 个公开 PDE 数据集上预训练了参数从 30M 到 0.5B 的模型。参数为 90M 的模型在零样本误差上较现有参数为 120M 的模型实现了最高可达 40% 的降低。此外，我们进行了可解释性分析，表明数据集类型可从路由-门控网络决策中推断出来，这验证了 MoE 架构的合理性和有效性。
关键词
引用
@article{arxiv.2510.25803,
title = {Mixture-of-Experts Operator Transformer for Large-Scale PDE Pre-Training},
author = {Hong Wang and Haiyang Xin and Jie Wang and Xuanze Yang and Fei Zha and Huanshuo Dong and Yan Jiang},
journal= {arXiv preprint arXiv:2510.25803},
year = {2025}
}