TL-DeepONet：迁移算子学习

前言

产出DeepONet的实验室又来挖坑了，George Em Karniadakis教授将迁移学习和DeeoONet结合起来，将迁移学习的研究内容应用到算子学习这个领域来，解决了诸多问题。迁移学习的核心可以看作是初始化，也可以认为是元学习的一种，利用较小的样本数将相近任务的学习到的全局特征泛化到当前任务来。也就是说，将源模型提取的域不变特征迁移到一个新任务，此外，还利用浅层网络提出的特征具有一般性这一经验观察对深层网络的权重进行调整。

Goswami, S., Kontolati, K., Shields, M. D., & Karniadakis, G. E. (2022). Deep transfer operator learning for partial differential equations under conditional shift. Nature Machine Intelligence, 1-10.

问题描述

源模型，针对特定任务已经训练完成的网络，本文中特指DeepONet，其训练过程我这里就略过去了，具体内容可以参考本系列介绍DeepONet的内容。目标模型，与源模型的任务具有一定相关性但是有明显差异，对于算子学习而言，迁移学习的目标可以是

1. PDE的模型不变，但是参数会根据不同场景发生变化
2. PDE控制输入的采样域发生了变化，也就是采样数据点是不同的
3. 不同的边界条件
4. 不同的系统演化条件
   具体而言，目标模型只有少量的标注数据\( D_t=\left( x_{i}^{tL},y_{i}^{tL} \right) _{i=1}^{N_t}\)和大量的无标注数据\(\left(x_{i}\right)^{tU} \)。

下面介绍迁移学习中的核心问题，也就是如何学习目标域和源域之间的联系，也就是目标域和源域的数据分布之间的偏移问题，源和目标输入数据的边际分布分别是\( P\left( \boldsymbol{x}_s \right)\)和\(P\left( \boldsymbol{x}_t \right) \)，输出的条件分布分别是\( P\left( \boldsymbol{y}_s\mid \boldsymbol{x}_s \right)\)和\( P\left( \boldsymbol{y}_t\mid \boldsymbol{x}_t \right) \)，有两种偏移：

1. 条件性偏移：\( P\left( \boldsymbol{x}_s \right) =P\left( \boldsymbol{x}_t \right), P\left( \boldsymbol{y}_s\mid \boldsymbol{x}_s \right) \ne P\left( \boldsymbol{y}_t\mid \boldsymbol{x}_t \right) \)
2. 协变量偏移：\( P\left( \boldsymbol{x}_s \right) \ne P\left( \boldsymbol{x}_t \right), P\left( \boldsymbol{y}_s\mid \boldsymbol{x}_s \right) = P\left( \boldsymbol{y}_t\mid \boldsymbol{x}_t \right) \)

本文中作者是考虑了条件性偏移，也就是任务输入数据的分布是相同由于参数等原因导致的输出分布差异。

方法与网络架构

本文的核心就是度量条件性偏移，具体而言，根据少量的标注数据比较源模型的输出比较惩罚条件性偏移，使得源模型的输出能够服从标注数据的分布从而实现迁移学习的目标。

核心问题其实如何计算条件性偏移，本文中给出了计算方法，这里只列出计算过程，略去了一些符号说明：

$$ \begin{aligned} D_{CEOD}\left( D_p,D_q \right) &=\left\| \hat{C}_{Y_p\mid}X_p-\hat{C}_{Y_q\mid X_q} \right\| _{HS}^{2}\\ &=\parallel \Phi \left( Y_p \right) \left( \boldsymbol{K}_{X_p}x_p+\lambda N_1\boldsymbol{I} \right) ^{-1}Y^{\top}\left( X_p \right)\\ &-\Phi \left( Y_q \right) \left( \boldsymbol{K}_{X_qX_q}+\lambda N_2\boldsymbol{I} \right) ^{-1}Y^{\top}\left( X_q \right) \parallel _{HS}^{2}\\ &=Tr\left\{ \left( \boldsymbol{K}_{X_pX_p}+\lambda N_1\boldsymbol{I} \right) ^{-1}\boldsymbol{K}_{Y_pY_p}\left( \boldsymbol{K}_{X_pX_p}+\lambda N_1\boldsymbol{I}^{-1}\boldsymbol{K}_{X_pX_p} \right\} \right.\\ &+Tr\left\{ \left( \boldsymbol{K}_{X_qX_q}+\lambda N_2\boldsymbol{I} \right) ^{-1}\boldsymbol{K}_{Y_qY_q}\left( \boldsymbol{K}_{X_qX_q}+\lambda N_2\boldsymbol{I} \right) ^{-1}\boldsymbol{K}_{X_qX_q} \right\}\\ &-2Tr\left\{ \left( \boldsymbol{K}_{X_pX_p}+\lambda N_1\boldsymbol{I} \right) ^{-1}\boldsymbol{K}_{Y_pY_q}\left( \boldsymbol{K}_{X_qX_q}+\lambda N_2\boldsymbol{I} \right) ^{-1}\boldsymbol{K}_{X_qX_p} \right\} ,\\ \end{aligned} $$

下面介绍本文的主要方法，本文的网络结构采用了DeepONet，主要训练流程如下：

首先就是将源模型的权重\( \theta^{S *} \)传递给目标网络，然后冻结大部分网络权重，只训练Trunk Net的最后一层和Branch Net的全连接层的参数。这里有一个很强的先验，也就是计算机视觉领域中公认的是，卷积层是通用的，而全连接层是针对特定任务的。

Loss函数设计如下：

$$ \begin{aligned} L\left( \theta ^T \right) &=\lambda _1\mathrm{L}_r\left( \theta ^T \right) +\lambda _2L_{CEOD}\left( \theta ^T \right)\\ &=\lambda _1\frac{\left\| f_T\left( \boldsymbol{x}^{tL} \right) -\boldsymbol{y}^t \right\| _2}{\left\| \boldsymbol{y}^t \right\| _2}+\lambda _2\left\| \hat{C}_{Y_{tL}\mid X_{tL}}-\hat{C}_{Y_{tU}\mid X_{tU}} \right\| _{HS}^{2}\\ \end{aligned} $$

这里分两项，第一项是针对少量标注数据，第二项是保证标注数据和目标模型的输出数据之间的分布是相同。这里可训练的参数有两个，分别是\( \theta ^T \)和\( \lambda _2 \)。具体而言是分别是最小化第一项和最大化第二项来更新参数，也就是利用梯度下降分别优化\( \theta ^T \)和增大\( \lambda _2 \)，保证了输出分布和预测精度。

实验

实验部分文章里已经进行了大量的说明，这里就略去不谈了。

结论

迁移学习在算子学习中其实很重要，因为针对同一个PDE，大部分任务之间具有很多关联性，也就是都需要提取全局特征，这里信息是可以跨任务共享的。本文中的方法其实很巧妙，并没有引入过多的数学公式，看完之后却让人有种醍醐灌顶的感觉，这也是一个有生命力的领域，非常值得未来的探索。