🔢 mathematics

Scalable Mechanistic Neural Networks for Differential Equations and Machine Learning

本文介绍了可扩展机制神经网络（S-MNN），这是一种增强型框架，在保持原有机制神经网络的准确性和可解释性的同时，将长时序序列建模的计算复杂度从立方级和平方级降低至线性级。

原作者： Jiale Chen, Dingling Yao, Adeel Pervez, Dan Alistarh, Francesco Locatello

发布于 2026-05-15

📖 1 分钟阅读🧠 深度阅读

原作者： Jiale Chen, Dingling Yao, Adeel Pervez, Dan Alistarh, Francesco Locatello

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教计算机理解一首复杂歌曲的节奏，就像理解季节的更替或海浪的律动一样。你希望计算机不仅仅是猜测下一个音符（预测），而是真正学会音乐的规则——即让歌曲运转的底层物理机制。

这就是论文中所称的机制神经网络（MNN）。它是一种智能工具，能够从杂乱的数据中学习这些隐藏的规则。然而，该工具的原始版本存在一个重大问题：它就像试图一次性观察每一块碎片来拼凑一幅巨大的拼图。如果歌曲很短，那还勉强可行。但如果歌曲很长（比如多年的天气数据），计算机就会不堪重负，耗尽内存并崩溃。它过于缓慢且笨重，无法用于现实世界中的长期问题。

本文的作者 Jiale Chen 及其团队构建了一个全新的升级版本，称为S-MNN（可扩展机制神经网络）。以下是他们如何通过一些简单的类比来解决这一问题的：

1. 问题所在：“沉重的背包”

原始的 MNN 就像一位背着装满沉重且多余岩石的背包的徒步者。

那些岩石：这些是额外的数学变量（称为“松弛变量”）以及计算机必须为每一个时间点计算的复杂约束条件。
结果：随着时间线变长（更多的天数、周数或年数），背包变得越来越重。求解拼图所需的时间呈立方级增长（如果时间翻倍，所需时间变为 8 倍），而所需的内存则呈平方级增长。这意味着原始工具只能处理很短的序列，随后就会放弃。

2. 解决方案：“精简的背包”

团队重新设计了该工具，以卸下那些沉重的岩石。

扔掉岩石：他们意识到并不需要那些额外的变量。通过移除它们，并将数学问题从复杂的“二次规划”问题转变为更简单的“最小二乘”问题（将其想象成从解迷宫转变为画一条直线），他们显著减轻了负担。
带状结构：他们发现所进行的数学运算具有特殊的模式，就像一列火车，每节车厢只与紧邻的前后车厢相连。与其尝试将每节车厢与其他所有车厢相连（这既混乱又缓慢），他们组织了数据，使计算机只需关注其直接相邻的部分。

3. 结果：“高速列车”

由于这些改进，新的S-MNN极其高效：

线性速度：随着序列变长，所需时间不再呈指数级恶化，而是呈直线增长。如果时间翻倍，所需时间也仅翻倍。
轻量化：它使用的计算机内存要少得多。
相同的精度：尽管速度更快、更轻量，但其精度与原始笨重缓慢的版本完全一致。它能完美地学习系统的规则。

现实世界测试

作者在三种不同的“歌曲”上测试了他们的新工具，以证明其有效性：

洛伦兹系统（The Lorenz System）：一个著名的混沌天气模型。新工具与旧工具一样好地找到了隐藏的规则，但速度快得多。
KdV 方程（The KdV Equation）：一个描述浅水波的复杂数学问题。新工具准确地解决了它，而旧工具的“稀疏”版本未能完成这项工作。
海表温度（大考）：他们尝试预测4 年（超过 200 周）的海洋温度。
- 原始工具因数据过长而崩溃；它耗尽了内存。
- 新的S-MNN轻松处理了 4 年的时间线，以高精度预测了温度模式。

核心结论

可以将原始的 MNN 想象为一个强大但笨拙的巨人，它只能举起沉重的重量几秒钟就会倒下。而新的S-MNN则是同一个巨人，但现在它成了一名马拉松运动员。它可以承载同样沉重的智力负荷（学习复杂的科学规则），但能够奔跑数英里（处理长时间序列）而毫不费力。

论文得出结论，这个新工具可以“即插即用”地替换旧工具，使科学家能够将这些强大的学习工具应用于长期、现实世界的数据（如气候记录），而这在以前是不可能的。

技术摘要：可扩展机制神经网络（S-MNN）

问题陈述

Pervez 等人（2024）提出的机制神经网络（MNN）通过在含噪观测数据中学习显式的内部常微分方程（ODE）表示，代表了科学机器学习领域的重大进步。与仅专注于预测的传统黑盒方法不同，MNN 支持下游科学分析，如参数识别和因果效应估计。然而，原始 MNN 公式在应用于长时序序列或高分辨率数据（例如气候记录）时，面临严重的可扩展性挑战。

核心瓶颈在于求解器的计算复杂度。原始 MNN 依赖于求解与 ODE 表示相关的大型线性方程组。

稠密求解器：在稠密矩阵上运行，导致相对于序列长度 $T$ 的时间复杂度为立方级 $O(T^3)$ ，空间复杂度为平方级 $O(T^2)$ 。
稀疏求解器：在稀疏矩阵上使用迭代方法（例如共轭梯度法）。虽然这降低了内存使用量，但非结构化的稀疏模式阻碍了对 GPU 并行性的充分利用。此外，迭代方法可能遭受收敛缓慢和数值不精确的问题，特别是在大规模或病态问题中。

这些局限性使得原始 MNN 即使在先进硬件上，对于长期动力学建模而言计算成本也过高。

方法论

作者提出了可扩展机制神经网络（S-MNN），这是对原始框架的重新表述，旨在将相对于序列长度的计算时间和空间复杂度从立方/平方级降低到线性级，同时保持精度。

1. 线性系统重构

S-MNN 通过对原始 MNN 进行三项关键更改，从根本上改变了底层线性系统的公式：

移除松弛变量：消除了原始平滑约束中使用的松弛变量。
扩展平滑约束：将前向和后向平滑约束扩展到最高阶导数（ $r=R$ ），取代了先前使用的中心差分约束。
转向最小二乘法：将优化问题从二次规划转换为最小二乘回归。

这些修改导致了一个线性系统 $Ay = b $，其中矩阵$ A $表现出特定的稀疏模式。关键在于，生成的正规方程矩阵$ M = A^\top W A $（其中$ W $是加权矩阵）具有**带状结构**。任意时间点$ t $的系数仅与时间点$ t-1 $、$ t $和$ t+1$ 直接相关。

2. 高效求解器设计

利用 $M$ 的带状结构，作者开发了一种针对 GPU 执行优化的专用求解器：

带状矩阵分解：求解器不再对稠密矩阵求逆，而是对带状矩阵 $M$ 执行分块 Cholesky 和 LDL 分解。这将 $M$ 分解为 $P L L^\top P^\top$ ，其中 $P$ 是分块下三角矩阵， $L$ 是分块对角矩阵。
前向/后向替换：解 $y$ 通过仅针对非零块（ $M_t, N_t$ ）操作的高效前向和后向替换步骤获得。
解析梯度：对于反向传播，作者推导出了梯度 $\partial \ell / \partial \beta$ 和 $\partial \ell / \partial M$ 的优雅解析解。与标准自动微分相比，这允许进一步加速，因为梯度可以通过使用已分解的带状矩阵求解类似的线性系统来计算。
数值稳定性：与共轭梯度等迭代方法相比，使用直接 Cholesky 分解提供了更优越的数值稳定性，这对于建模误差积累有害的混沌系统至关重要。

3. 复杂度分析

重构实现了以下复杂度降低：

时间复杂度：从 MNN 稠密求解器的 $O(T^3)$ 降低到 S-MNN 的 $O(T)$ 。
空间复杂度：从 MNN 稠密求解器的 $O(T^2)$ 降低到 S-MNN 的 $O(T)$ 。
这两种复杂度现在都随序列长度 $T$ 线性扩展，使得该方法适用于长轨迹。

主要贡献

复杂度降低：通过重构线性系统以消除松弛变量和中心差分约束，作者成功将相对于序列长度的时间和空间复杂度降低到线性级。
高效求解器设计：开发了一种新颖的 GPU 友好型求解器，利用重构系统的带状结构，采用分块分解和解析梯度解，实现了显著的加速。
长期科学任务验证：该框架在多种基准测试中得到了验证，包括控制方程发现（Lorenz 系统）、求解偏微分方程（KdV 方程）以及长期现实世界预测（海表温度），证明了 S-MNN 在大幅降低资源需求的同时，保持了与原始 MNN 相当的精度。

实验结果

作者在三个主要领域评估了 S-MNN：

独立验证：在五个线性 ODE 基准测试（例如 RC 电路、谐波振荡器）上，S-MNN 求解器产生的数值解与闭式解析解高度吻合，误差可忽略不计，证实了求解器的正确性。
Lorenz 系统发现：在发现混沌 Lorenz 系统的系数时，S-MNN 实现了与原始 MNN 相当的收敛率。
- 性能：对于长度为 50 的序列，S-MNN 比 MNN 稠密求解器实现了 4.9 倍的加速，并将 GPU 内存使用量降低了 50%。
- 可扩展性：与原始 MNN 在更长序列（例如长度 500）上出现内存溢出（OOM）或无法收敛的情况不同，S-MNN 保持了高性能和低内存使用量。
KdV 偏微分方程：在求解 Korteweg-de Vries 方程时，S-MNN 实现了 0.00005 的归一化均方误差（NMSE），优于 MNN 稠密求解器（0.00006），并显著超越了 ResNet 和 FNO 基线。MNN 稀疏求解器未能收敛。S-MNN 将训练时间从 38.0 小时减少到 10.1 小时，内存使用量从 3.40 GiB 减少到 2.19 GiB。
海表温度（SST）预测：应用于 4 年（208 周）的全球 SST 数据集，S-MNN 成功建模了原始 MNN 因资源限制而无法处理的长时序序列。
- 可扩展性：随着序列长度增加（高达 1,727 周），S-MNN 的运行时间和内存使用量呈线性增长。相比之下，MNN 稠密求解器在超过 104 周后失效，稀疏求解器在超过 416 周后失效。

意义与主张

该论文声称，S-MNN 为将机制瓶颈集成到复杂动力学系统的神经网络模型中提供了一种实用且高效的工具。通过将计算复杂度降低到线性级，S-MNN 使得在保持原始 MNN 框架固有的准确性和可解释性的同时，对长期动力学进行建模成为可能。

作者断言，S-MNN 可以在应用中直接替换原始 MNN，为解决可扩展性限制提供了方案，这些限制此前阻碍了机制神经网络应用于气候建模等现实世界的长视界问题。虽然作者承认分解算法中的序列依赖关系仍然限制了时间维度上的完全并行化，但他们将 S-MNN 定位为科学机器学习领域的重大进步，弥合了理论机制建模与实际大规模应用之间的差距。