基于积分自编码器网络的离散不变学习

前言

最近机器学习顶刊Journal of Machine Learning Research发表了很多关于PDE深度学习的理论基础，可以看出这个领域将会焕发出生命力，该领域与数学结合较紧密。PDE的深度学习方法很难避开的就是离散化。如何去学习离散过程中产生的不变量是一个重要的话题，这篇文章为理解这个问题打开了全新的视角。

Ong, Y. Z., Shen, Z., & Yang, H. (2022). Integral Autoencoder Network for Discretization-Invariant Learning. Journal of Machine Learning Research, 23(286), 1-45.

离散不变学习的核心就是针对任何方式的离散，神经网络都不需要重新训练，一种现有的方法就是DeepOnet为代表的网格无关的算子学习方法。但是本文中的方法力求去解决离散方法中的离散不变，也即训练好的神经网络可以针对任意离散格式的数据，这样可以提高图网络等方法的泛化能力。

后面的部分我尽量以较为简单易懂的方式去描述，会忽略数学证明等理论性较强的过程。

问题描述

让\( S_x=\left( x_j \right) _{j=1}^{s_x}\subset [0,1]^{d_x}\)和\( S_y=\left(y_j \right) _{j=1}^{S_y}\subset [0,1]^{d_y} \)分别\(\Omega _x\)和 \(\Omega _y\)域上的有限采样点，对应的数值观测值为\( \bar{f}=\left( f\left( x_1 \right) ,...,f\left( x_{s_x} \right) \right)\)和\(\bar{g}=\left( g\left( y_1 \right) ,...,g\left( y_{s_y} \right) \right)\)。特别的，对于不同组的观测值，采样点的数量\( s_x\)和\( s_y\)都可以是不同的。

传统的神经网络或者多层感知机只能处理具有固定输入维度的数据，对于CV任务来说，所有输入的图片都必须先做归一化处理。

这篇文章的核心就是基于神经网络的非线性积分变换，可以将具有任意维度的输入转换进入具有固定维度的空间中，也就是

$$ (\bar{T}f)(z)=\int_{\Omega _x}{\bar{K}}(f(x),x,z)f(x)dx,\quad z\in \Omega _z, \\ (\bar{T}f)\left( z_j \right) =\frac{1}{s}\sum_{i=1}^s{\bar{K}}\left( f\left( x_i \right) ,x_i,z_j \right) f\left( x_i \right) ,\quad z_j\in S_z. $$

这就是本文的中心思想，学习这么一个非线性积分变化，也就是一个积分自编码器(IAE)，进行离散不变学习。

方法与网络架构

这篇文章的主要思想是和FNO差不多，都是以下的流程，首先将输入的观测值 \(\bar{f} \)投影到一个高维特征空间\(a_0\)，后面还是一系列的高维空间\(a_1,...,a_L\)。这些高维空间\(a_1,...,a_L\)是给定的固定维度，并不由随着输入和输出个维度而发生变化，当然了这些高维空间的维度是超参数，在训练过程是固定的。由于中间这些高维空间是具有不变的维度，我们可以利用全连接神经网络进行处理。这里其实是解决了FNO中固定网格的问题，也就是只能针对训练的离散格式进行推理，具体流程如下

$$ \bar{f} \stackrel{F}{\rightarrow} a_0 \stackrel{\mathcal{I A \mathcal { E }}_1}{\longrightarrow} a_1 \stackrel{\mathcal{I A E}_2}{\longrightarrow} \ldots \stackrel{\mathcal{I A \mathcal { E }}_L}{\longrightarrow} a_L \stackrel{G}{\rightarrow} \bar{g} $$

\(a_0\)和\(a_L\)分别是预处理的输入数据特征和等待后处理的输出数据特征，这里的难度就是如何将具有不同离散维度的输入输出数据进行处理，然后利用神经网络进行处理。

预处理：也就是将\( \bar{f} \)转换为\(a_0 \)，应用逐点线性变换（本质是一个一层逐点多层感知机）将每一个离散维度都投影到一个全新的特征空间中去。

后处理：后处理使用了FNO中两个连续的神经网络对输出进行变化（对高频信号进行截断），最后应用一个相同的逐点线性变化将\(a_L\)投影到输出\( \bar{g} \)上。

多通道学习：使用不同特征空间加强不同表征学习的能力，作者将原始数据投影到不同的特征空间上去，如Fourier或Wavelet等，再将学习到的各个通道的特征使用逐点线性变化进行融合。具体流程如下图所示：

IAE模块：下面介绍最重要的部分，也就是每一个IAE模块内部的结构，这篇文章里面说的含糊不清，我参考了源码才看懂，下面的\(a\)和\(b\)是在每一个模块内部的，也就是说这些每一个IAE模块内部包含了如下的结构：

具体而言，\( \phi _1 \)和\( \phi _2 \)分别是非线性积分编码器好非线性积分解码器，\( \phi _0\)是一个多层感知机，这个IAE模块内部是有N层结构，先是进行N次的积分编码，在进行N次积分解码。具体而言，就是

$$ v_i(z)=\int_{\Omega_{i-1}} \phi_1^i\left(v_{i-1}(x), x, z ; \theta_{\phi_1^i}\right) v_{i-1}(x) d x\\ u_i(x)=\int_{\Omega_{u_{i-1}}} \phi_2^i\left(u_{i-1}(z), x, z ; \theta_{\phi_2^i}\right) u_{i-1}(z) $$

这里\(v_0\)是\(a\)，\(u_N\)是\(b\)，这样也就很清楚了，不要和后面的稠密连接技巧搞混了。
这样一个IAE模块就解决了，后面的就是将IAE模块进行堆叠。这个结构是和FNO很像，这里是区别就是编码进行了N次在进行N次的解码，略有不同，需要注意。

稠密链接技巧：稠密连接能够增加每一个IAE的特征多样性，最终实现学习数据的复杂表征并缓解梯度消失的问题

Huang, G., Liu, Z., Van Der Maaten, L., & Weinberger, K. Q. (2017). Densely connected convolutional networks. In IEEE Conference on Computer Vision and Pattern Recognition (pp. 4700-4708).

具体的网络架构如下如所示，具体可以参考作者的代码：

数学表达式为：

$$ a_i=\sigma\left(\mathcal{I} \mathcal{A} \mathcal{E}_i\left(\mathcal{M}_i\left(\left[\mathcal{A}_0\left(a_0 ; \theta_{\mathcal{A}_0}\right), \mathcal{A}_1\left(a_1 ; \theta_{\mathcal{A}_1}\right), \ldots, \mathcal{A}_{i-1}\left(a_{i-1} ; \theta_{\mathcal{A}_{i-1}}\right)\right] ; \theta_{\mathcal{M}_i}\right) ; \theta_{\mathcal{I A} \mathcal{E}_i}\right)\right) $$

随机数据增广：传统的核方法会面临泛化能力较弱的问题，为了解决问题，这里还引入了一个技巧，为了应对不同的离散格式，这篇文章提出利用插值方法得到不同离散格式下的数据，然后分别对其进行训练，来优化以下的的loss函数

$$ \min _{\theta_{\Psi^n}} \mathbb{E}_{(\bar{f}, \bar{g}) \sim p_{\text {data }}} \mathbb{E}_{I_T \sim \mathcal{I}}\left[L\left(\Psi^n\left(\bar{f} ; \theta_{\Psi^n}\right), \bar{g}\right)+\lambda L\left(\Psi^n\left(I_T(\bar{f}) ; \theta_{\Psi^n}\right), I_T(\bar{g})\right)\right] $$

第一项是对原始数据的损失，第二项是进行了随机数据插值增广之后的损失，这样保证了对于不同的离散格式的学习。

实验结果

这里作者对比现有的几种方法：DeepOnet，FNO，ResNet+Interpolation，Unet，Fourier Transformer (FT), Galerkin Transformer (GT)

Cao, S. (2021). Choose a transformer: Fourier or galerkin. Advances in Neural Information Processing Systems, 34, 24924-24940.

实验结果大致如下图

效果还不错，但是参数量、训练时间和推理速度还是有点不尽如人意，

存在很多能提高的方面，如何细化这个自编码积分器可以说是未来一个非常重要的方向。

代码链接为http://github.com/IAE-Net/iae...

结论

这篇文章提出的方法是很具有代表性的，力求统一不同离散格式产生的问题，但是还有很多可以提高的方面，对于DeepOnet而言，已经有人解决了不同传感器位置的问题，将之与本文的思路进行结合也是一个有趣的研究话题。

Prasthofer, M., De Ryck, T., & Mishra, S. (2022). Variable-Input Deep Operator Networks. arXiv preprint arXiv:2205.11404.