CALM-PDE：Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs

1. 论文

摘要：

本文从压缩角度叙事。

引入了连续的自适应卷积，用于将偏微分方程数据从高分辨率的物理空间压缩到低分辨率的潜在空间。其中query points 是可学习的（可移动），提取 epsilon neighborhood 的信息。 NeurIPS

---

2. 背景：

1. PDEs 数据通常是由一组网格 \(M\) （均匀或者非均匀）离散区域 \(\Omega\) 得到的。在 2D 问题中，\(M\) 可能是 \(211 \times 211\)；在 3D 问题中，\(M\) 可能是 \(64 \times 64 \times64\)。
2. 然而，直接在物理域中计算会带来大量的计算成本，尤其是基于 transformer 的模型。现有的方法通过将输入压缩到潜在空间 \(L\) 上来缓解这个问题，其中潜在空间的大小 \(|L| << |M|\)。

---

3. 核心思想：

本文遵循了 encoder-process-decode 的框架。其中，encoder 压缩空间维度并且增加通道维数，由多个 CALM layers 组成。

下面介绍 CALM layer的设计。

Continuous Convolutional Neural networks [2]：

令 \(f, k: \mathbb{R} \to \mathbb{R}\) 是两个实值函数，\(a \in \mathbb{R}\)，那么 \(f\) 和 \(k\) 的卷积为：

\[(f * k)(a) \coloneqq \int_{-\infty}^{+\infty} f(\alpha)\cdot k(a - \alpha)d\alpha. \tag{1} \]

在 [2] 中使用 MLP 近似 learnable kernel \(k\) ，如下：

\[(f * k)(a) = \int_{-\infty}^{+\infty} f(a) \cdot k(a-\alpha)d\alpha \approx \frac{1}{N} \sum_{n=1}^{N} f(\alpha_n) \cdot k(a-\alpha_n), \tag{2} \]

其中，N 个输入点 \(\alpha_n\) 从区域中采样。核函数由一个MLP \(k_{\theta}(a-\alpha)\) 构建得到。

类似离散卷积，连续卷积层包含针对 \(N_i\) 个输入通道的滤波核：

\[(f * k)_o(a_j) = \sum_{i=1}^{N_i} \sum_{n=1}^{N} f_i(\alpha_n) \cdot k_{i,o}(a_j-\alpha_n), \tag{3} \]

其中，\(o\) 表示 \(o^{th}\) 输出通道。

定义输出点集合为 \(A=\{ a_j \}_{j=1}^{N'}\)，输入点集合 \(\mathcal{A} =\{ \alpha_n \}_{n=1}^{N}\)。与离散卷积不同的是，输出点的数量不一定等于连续点的数量，从而使得连续卷积可以压缩信息以减少点的数量（或上采样）。

类似于离散卷积，每一个输出点的感受野被限制到 M 个点，使得：

\[(f * k)_o(a_j) = \sum_{i=1}^{N_i} \sum_{m\in \mathrm{RF} (a_j)} f_i(\alpha_m) \cdot k_{i,o}(a_j-\alpha_m), \tag{4} \]

\(\rm{RF}(a_j)\) 表示 \(a_j\) 感受野内的点的指标。

CALM Layer：

有了连续卷积，下面来构造 CALM 层。

1. 参数化核函数 \(k\)：使用两层的 MLP 接收带来随机傅里叶特征的输入 \(a-\alpha\)，即 \(k_{i,o}(a-\alpha) = \mathrm{MLP}(\mathrm{RFF}(a-\alpha))_{i,o}\)，其中 \(i,o\) 分别表示输入和输出通道。

2. Epsilon 最近邻和距离加权：epsilon 近邻表示仅考虑所有输入点 \(\alpha_i\) 中，距离 query 点 \(a_j\) 最近的 \(p\%\) 个点，\(p\) 依赖于问题和数据。因此，感受野可以看作：\(\mathrm{RF}(a_j) =\{ \alpha_n\in \mathcal{A} | \| a_j - \alpha_n \|<\epsilon(a_j) \}\)，\(\epsilon(a_j)\) 表示最近 \(p\) 个点。因此，核函数被修正为：

   \[$k_{i,o}(a-\alpha) = \mathrm{MLP}(\mathrm{RFF}(a-\alpha))_{i,o}\frac{\exp\!\left(-\dfrac{\|a - \alpha\|^2 - \min(a)}{T\,(\max(a) - \min(a))}\right) }{\displaystyle\sum_{\alpha_j \in \mathrm{RF}(a)}\exp\!\left(-\dfrac{\|a - \alpha_j\|^2 - \min(a)}{T\,(\max(a) - \min(a))}\right) }, \quad min(a) = min_{\alpha_j \in \mathrm{RF}(a)}\| a-\alpha_j\|^2. \]

3. 可学习的query点：初始化 query 点为均匀采样，或者从底层网格中采样。在训练期间，query 点通过梯度下降来优化：

   \[(x_i, y_j)^{\top} \gets (x_i, y_j)^{\top} - \eta \frac{\partial \mathcal{L_\theta}}{\partial (x_i, y_j)^{\top}}. \]

   注意，decoder 的最后一层使用输出点，不更新。

4. 核调制：为每一个 query 点设计特有的过滤核：

   \[k_{i,o}(a - \alpha) = \left(W_2 \cdot \sigma\!\left((W_1 \cdot \mathrm{RFF}(a - \alpha) + b_1)\odot \gamma_a + \beta_a\right)+ b_2 \right)_{i,o}, \]

   \(\gamma_a, \beta_a\) 分别表示 query 点 \(a\) 的缩放和偏移。

经过1，2，3，4，CALM layer如下所示：

\[l(f, \mathcal{A},A) = \{ (f * k)(a_j) \}_{j=1}^{N'} ,\quad (f * k)_o(a_j) = \sigma\!\left(\sum_{i=1}^{N_i}\sum_{m \in \mathrm{RF}(a_j)}f_i(\alpha_m) \cdot k_{i,o}(a_j - \alpha_m)+ b_o \right), \]

其中 \(f, \mathcal{A}\) 表示输入，\(A\) 表示 query 点。当 query 点数小于输出点时，执行下采样（压缩）；反之执行上采样。

Model：encoder 和 decoder 由多个 CALM层构建，processor 可以使用显式欧拉法，NODE，自注意力等。

Remark: 本文发现使用课程学习方法能够有效的提升精度，即从短的轨迹逐渐增加到长轨迹。此外，一个自增强被引入，即预测初始条件。

---

4. 实验：

对于图3，可以看到解码器学习到的query points 捕捉到了物理规律分布。

内存和训练开销可以接受。

5. 讨论：

在训练完成后，query 点被保留下来用于推断。所获得的query points 是符合训练集的整体规律的，然后对于单个样本（有可能会偏离总体规律）并不是最优的。这主要是因为，query points 的移动，与输入内容无关。

---

6. 参考文献

1. Jan Hagnberger, Daniel Musekamp, and Mathias Niepert. CALM-PDE: Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs. In Advances in Neural Information Processing Systems (NeurIPS), 2025. [PDF] [Code]
2. Shenlong Wang, Simon Suo, Wei-Chiu Ma, Andrei Pokrovsky, and Raquel Urtasun. Deep Parametric Continuous Convolutional Neural Networks. In Conference on Computer Vision and Pattern Recognition (CVPR), 2018.