Transfer Learning with Physics-Informed Neural Networks for Efficient Simulation of Branched Flows

This work uses PINNs that are trained in an unsupervised manner, as detailed below. We compare the performance achieved by feed forward neural networks (FFNN) and GAN models.

FFNN: The standard unsupervised neural network approach was introduced by Lagaris et al., [7] and can be used to solve differential equations of the form

$$\mathcal{L}\left(u(t,\mathbb{x})\right)=0,$$

(1)

where $\mathbb{x}=(x_{1},\dots,x_{n})$, $u:\mathbb{R}^{n+1}\rightarrow\mathbb{R}$ and $\mathcal{L}$ is a differential operator. During training, we sample $(t,\mathbb{x})$ from the domain of the equation $D$ and use this vector as input to a FFNN, which outputs the neural solution $u_{\theta}(t,\mathbb{x})$. We re-parametrise this output into $\tilde{u}_{\theta}(t,\mathbb{x})$ to satisfy initial and boundary conditions exactly. Using automatic differentiation, we can compute the derivatives of this output with respect to each of the independent variables and build the loss by summing the squared residual over $M$ training points,

$$\frac{1}{M}\sum_{(t,\mathbf{x})\in D}\mathcal{L}\left(\tilde{u}_{\theta}(t,\mathbb{x})\right)^{2}.$$

(2)

Note that if the network $\tilde{u}_{\theta}$ perfectly satisfies Equation 1, then Equation 2 will be zero.

DEQGAN: Bullwinkel et al., [1] noted that there is no theoretical reason to use the $L_{2}$ norm of the residuals over any other loss function and proposed DEQGAN, which extends the FFNN method to GANs and can be thought of as “learning the loss function.” Rather than computing a loss over the equation residuals, DEQGAN labels these vectors “fake” data samples and zero-centered Gaussian noise as “real” data samples. As the discriminator gets better at classifying these samples, the generator $\tilde{u}_{\theta}$ is forced to propose solutions such that the equation residuals are increasingly indistinguishable from a vector of zeros, thereby approximating the solution to the differential equation.

Figure 1 illustrates the multi-head architecture that we apply to FFNN and DEQGAN models to perform transfer learning. The output of the base neural network is passed to heads $h_{1},...,h_{L},$ each of which corresponds to the solution to the system at a particular initial condition. Importantly, this architecture can be used to apply transfer learning to non-linear problems; in this work, we consider one such system of ODEs that is used to model particles moving through a weak random potential.

Transfer learning with multi-head PINNs is performed in two stages. First, we train the multi-head model on a given set of initial conditions until convergence. Next, we freeze the weights of the base network and fine-tune only the output heads on a second set of initial conditions. As detailed below, we use this procedure to perform two transfer learning tasks: 1) Initial Condition Transfer, which involves fine-tuning the heads on initial conditions that were not used to train the base. 2) Potential Transfer Learning, an even more challenging task that allows us to obtain solutions for new initial conditions and a different potential than the one used to train the base. Our results on these tasks suggest that the base is able to learn highly general properties of the system.

Stochastic branched flow is a universal wave phenomenon that occurs when waves propagate in random environments. Branching has been observed in tsunami waves [2], electronic flows in graphene [12], and electromagnetic waves in gravitational fields [8]. We can model a two dimensional branched flow by considering a particle with position $\mathbb{x}=(x,y)$ and velocity $\mathbb{p}=(p_{x},p_{y})$, both functions of time $t$, traveling through a weak random potential $V(x,y)$. With the Hamiltonian $H(\mathbb{x},\mathbb{p})=||\mathbb{p}||^{2}_{2}/2+V(\mathbb{x})$ we obtain Hamilton’s equations, given by the following system of ODEs

$$\left\{\begin{aligned} \dot{x}(t)&=p_{x}(t)\\ \dot{y}(t)&=p_{y}(t)\\ \dot{p_{x}}(t)&=-\frac{\partial V(x(t),y(t))}{\partial x}\\ \dot{p_{y}}(t)&=-\frac{\partial V(x(t),y(t))}{\partial y},\end{aligned}\right.$$

(3)

For a plane wave, the initial conditions at $t=0$ can be chosen as $(x(0),y(0),p_{x}(0),p_{y}(0))=(0,y(0),1,0)$ [12].

We build random potentials by summing $K$ randomly distributed Gaussian functions with covariance matrix $\sigma^{2}I_{2}\in\mathbb{R}^{2\times 2}$ and means $\mu_{i}\in\mathbb{R}^{2}$, $i=1,\dots,K$, and scaling the result by $-A$, where $A\in\mathbb{R}^{+}$, as in [12]. That is,

$$V(\mathbb{x})=-\frac{A}{2\pi\sigma^{2}}\sum_{i=1}^{K}\exp\left(-\frac{1}{2\pi\sigma^{2}}||\mathbb{x}-\mathbb{\mu}_{i}||_{2}^{2}\right).$$

(4)

In the experiments presented below, we use $K=10,~{}A=0.1,~{}\sigma=0.1$.

For both the FFNN and DEQGAN models, we use networks with a base consisting of 5 hidden layers and 40 nodes. We then use $L$ linear layers for the heads. Each head is responsible for the solution to one ray, i.e., one initial condition $\mathbb{z}_{l}(0)=(0,y_{l}(0),1,0)$, where $l=1,\dots,L$, and has four outputs corresponding to $x$, $y$, $p_{x}$ and $p_{y}$. For head $l$, we denote the outputs as $u_{l,1}$, $u_{l,2}$, $u_{l,3}$ and $u_{l,4}$. We use the initial value re-parameterization proposed by Mattheakis et al., [10]

$$\tilde{u}_{l,i}(t)=\left[\mathbb{z}_{l}(0)\right]_{i}+\left(1-e^{-t}\right)u_{l,i}(t),\quad l=1,\dots,L;i=1,\dots,4$$

(5)

which forces the proposed solution to be exactly $\mathbb{z}_{l}(0)$ when $t=0$ and decays this constraint exponentially in $t$.

Our first transfer learning task, Initial Condition Transfer, allows us to efficiently obtain solutions to the system for many initial conditions. We used multi-head models to train the base networks on $L=11$ initial conditions $y_{l}(0)=0.0,0.1,\dots,1.0$ and performed single-head transfer learning on $100$ evenly-spaced initial conditions in $[0,1]$ while keeping the potential fixed. All experiments were performed on a Microsoft Surface laptop with Intel i7 CPU. Figure 2(a) shows the ray trajectory solutions corresponding to the initial conditions used for base training (black) and transfer learning (blue) obtained with the FFNN model. These trajectories also illustrate branched flows.

In Figure 3, we compare the losses achieved by the FFNN and DEQGAN models during base training and transfer learning. We also show the residuals for classical models that do not leverage transfer learning. Notably, we see that single-head models that use transfer learning converge more rapidly than those trained from scratch. Further, Table 1 shows that each epoch of transfer learning (bold) is also significantly faster. This is to be expected because transfer learning involves only fine-tuning linear heads, rather than training an entire base network.

Our second transfer learning task, Potential Transfer Learning, utilizes the same pre-trained base described above. This task, however, changes not only the initial conditions, but also the potential (Equation 4). More specifically, we constructed a new potential by randomly sampling ten new Gaussian means. To avoid significantly altering the statistical properties of the system, we used the same values of $\sigma$ and $A$. Figure 2(b) visualizes the ray trajectories obtained using this method and suggests that the multi-head models are, indeed, able to learn highly general bases for the system.