SwitchNet: a neural network model for forward and inverse scattering problems

Introduction

In this paper, we study the forward and inverse scattering problems via the use of artificial neural networks (NNs). In order to simplify the discussion, we focus on the time-harmonic acoustic scattering in two dimensional space. The inhomogeneous media scattering problem with a fixed frequency $\omega$ is modeled by the Helmholtz operator

where $c(x)$ is the velocity field. In many settings, there exists a known background velocity field $c_{0}(x)$ such that $c(x)$ is identical to $c_{0}(x)$ except in a compact domain $\Omega$. By introducing the scatterer $\eta(x)$ compactly supported in $\Omega$

we can equivalently work with $\eta(x)$ instead of $c(x)$. Note that in this definition $\eta(x)$ scales quadratically with the frequency $\omega$. However, as $\omega$ is assumed to be fixed throughout this paper, this scaling does not affect any discussion below.

In many real-world applications, $\eta(\cdot)$ is unknown. The task of the inverse problem is to recover $\eta(\cdot)$ based on some observation data $d(\cdot)$. The observation data $d(\cdot)$ is often a quantity derived from the Green’s function $G=L^{-1}$ of the Helmholtz operator $L$ and, therefore, it depends closely on $\eta(\cdot)$. This paper is an exploratory attempt of constructing efficient approximations to the forward map $\eta\rightarrow d$ and the inverse map $d\rightarrow\eta$ using the modern tools from machine learning and artificial intelligence. Such approximations are highly useful for the numerical solutions of the scattering problems: an efficient map $\eta\rightarrow d$ provides an alternative to expensive partial differential equation (PDE) solvers for the Helmholtz equation; an efficient map $d\rightarrow\eta$ is more valuable as it allows us to solve the inverse problem of determining the scatterers from the scattering field, without going through the usual iterative process.

In the last several years, deep neural network has become the go-to method in computer vision, image processing, speech recognition and many other machine learning applications . More recently, methods based on NN have also been applied to solving PDEs. Based on the way that the NN is used, these methods for solving the PDE can be roughly separated into two different categories. For the methods in the first category , instead of specifying the solution space via the choice of basis (as in finite element method or Fourier spectral method), NN is used for representing the solution. Then an optimization problem, for example an variational formulation, is solved in order to obtain the parameters of the NN and hence the solution to the PDE. Similar to the use of an NN for regression and classification purposes, the methods in the second category such as use an NN to learn a map that goes from the coefficients in the PDE to the solution of the PDE. As in machine learning, the architecture design of an NN for solving PDE usually requires the incorporation of the knowledge from the PDE domain such that the NN architecture is able to capture the behavior of the solution process. Despite the abundance of the works in using the NN for solving PDE, none of the above mentioned methods have tried to obtain the solution to the wave equation.

This paper takes a deep learning approach to learn both the forward and inverse maps. For the Helmholtz operator (1), we propose an NN architecture for determining the forward and inverse maps between the scatterer $\eta(\cdot)$ and the observation data $d(\cdot)$ generated from the scatterer. Although this task looks similar to the computer vision problems such as image segmentation, denoising, and super-resolution where the map between the two images has to be determined, the nature of the map in our problem is much more complicated. In many image processing tasks, the value of a pixel at the output generally only depends on a neighborhood of that pixel at the input layer. However, for the scattering problems, the input and output are often defined on different domains and, due to wave propagation, each location of the scatterer can influence every point of the scattered field. Therefore, the connectivity in the NN has to be wired in a non-local fashion, rendering typical NN with local connectivity insufficient. This leads to the development of the proposed SwitchNet. The key idea is the inclusion of a novel low-complexity switch layer that sends information between all pairs of sites effectively, following the ideas from butterfly factorizations . The same factorization was used earlier in the architecture proposed , but the network weights there are hardcoded and not trainable.

The paper is organized as followed. In Section 2, we discuss about some preliminary results concerning Helmholtz equation. In Section 3, we study the so called far field pattern of the scattering problem, where the sources and receivers can be regarded as placed at infinity. We propose SwitchNet to determine the maps between the far field scattering pattern and the scatterer. In Section 4, we turn to the setting of a seismic imaging problem. In this problem, the sources and receivers are at a finite distance, but yet well-separated from the scatterer.

Preliminary

Using the background velocity field $c_{0}(x)$, we first introduce the background Helmholtz operator $L_{0}=-\Delta-\omega^{2}/c_{0}^{2}$. With the help of $L_{0}$, one can write $L$ in a perturbative way as

where $E$ is viewed as a perturbation. By introducing the background Green’s function

one can write down a formal expansion for the Green’s function $G=L^{-1}$ of the $\eta$-dependent Helmholtz operator $L$:

which is valid when the scatterer field $\eta(x)$ is sufficiently small. The last line of the above equation serves as the definition of the successive terms of the expansion ($G_{1}$, $G_{2}$, and so on). As $G_{0}$ can be computed from the knowledge of the background velocity field $c_{0}(x)$, most data gathering processes (with appropriate post-processing) focus on the difference $G-G_{0}=G_{1}+G_{2}+\cdots$ instead of $G$ itself.

A usual experimental setup consists of a set of sources $S$ and a set of receivers $R$:

The data gathering process usually involves three steps: (1) impose an external force or incoming wave field via some sources, (2) solve for the scattering field either computationally or physically, (3) gather the data with receivers at specific locations or directions. The second step is modeled by the difference of the Green’s function $G-G_{0}$, as we mentioned above. As for the other steps, it is convenient at this point to model the first step with a source-dependent operator $\Pi_{S}$ and the third one with a receiver-dependent operator $\Pi_{R}$. We shall see later how these operators are defined in more concrete settings. By putting these components together, one can set the observation data $d$ abstractly as

In this paper, we focus on two scenarios: far field pattern and seismic imaging. We start with far field pattern first to motivate and introduce SwitchNet. We then move on to the seismic case by focusing on the main differences.

SwitchNet for far field pattern

In this section, we consider the problem of determining the map from the scatterer to the far field scattering pattern, along with its inverse map. Without loss of generality, we assume that the diameter of the domain $\Omega$ is of $O(1)$ after appropriate rescaling. The background velocity $c_{0}(x)$ is assumed to be 1 since the far field pattern experiments are mostly performed in free space.

In this limiting setting, one redefine the observation data as

Now taking the two limits (10) and (14) under consideration, one arrives at the following representation of the observation data $d(r,s)$ for $r\in R$ and $s\in S$

2. Low-rank property

The intuition behind the proposed NN architecture comes from examining (19) when $E$ (or $\eta$) is small. In such a situation, we simply retain the term that is linear in $E$. Using the fact that $E=\text{diag}(\eta)$, (19) becomes

For any $D_{i}$ and $X_{j}$, the submatrix

The proof of this theorem follows the same line of argument in and below we outline the key idea. Denote the center of $D_{i}$ by $(r_{i},s_{i})$ and the center of $X_{j}$ by $x_{j}$. For each $(r,s)\in D_{i}$ and $x\in X_{j}$, we write

Note that for fixed $D_{i}$ and $X_{j}$ each of the last three terms is either a constant or depends only on $x$ or $(r,s)$. As a result, $\exp(i\omega(s-r)\cdot x)$ is numerically low-rank if and only if the first term $\exp(i\omega((s-r)-(s_{i}-r_{i}))\cdot(x-x_{j}))$ is so. Such a low-rank property can be derived from the conditions concerning the side-lengths of $D_{i}$ and $X_{j}$. More precisely, since $(r,s)$ resides in $D_{i}$ with center $(r_{i},s_{i})$, then

Similarly as $x$ resides in $X_{j}$ with center $x_{j}$, then

Multiplying these two estimates results in the estimate

for the phase of $\exp(i\omega((s-r)-(s_{i}-r_{i}))\cdot(x-x_{j}))$. Therefore,

for $(r,s)\in D_{i}$ or $x\in X_{j}$ is non-oscillatory and hence can be approximated effectively by applying, for example, Chebyshev interpolation in both the $(r,s)$ and $x$ variables. Since the degree of the Chebyshev polynomials only increases poly-logarithmically with respect to the desired accuracy, $\exp(i\omega((s-r)-(s_{i}-r_{i}))\cdot(x-x_{j}))$ is numerically low-rank by construction. This proves that the submatrix $A_{ij}$ defined in (22) is also numerically low-rank. ∎

3. Matrix factorization

By applying (28) to each block $A_{ij}$, $A$ can be approximated by

The next step is to write (29) into a factorized form. First, introduce $U_{i}$ and $V_{j}$

With the above definitions for $U$, $V$, and $\Sigma$, the approximation in (29) can be written compactly as

Notice that although $A$ has $M^{2}\times N^{2}$ entries, using the factorization (34), $A$ can be stored using $tP(M^{2}+P+N^{2})$ entries. In this paper, $P\approx\max(M,N)$ and $M$ and $N$ are typically on the same order. Therefore, instead of $O(N^{4})$, one only needs $O(N^{3})$ entries to parameterize the map $A$ approximately using (34). Such a factorization is also used in for the compression of Fourier integral operators.

We would like to comment on another property that may lead to further reduction in the parameters used for approximating $A$. Let us focus on any two submatrices $A_{ij}$ and $A_{ik}$ of $A$. For two regions $X_{j}$ and $X_{k}$, where the center of $X_{j}$ and $X_{k}$ are $x_{j}$ and $x_{k}$ respectively, $X_{k}=X_{j}+(x_{k}-x_{j})$. Let $(r,s)\in D_{i}$. For $x\in X_{j}$ and $x^{\prime}=x+(x_{k}-x_{j})\in X_{k}$, we have

Therefore, the low-rank factorizations of $A_{ij}$ and $A_{ik}$ are solely determined by the factorization of $h(rs,x)$. This implies that it is possible to construct low-rank factorizations for $A_{ij}$ and $A_{ik}$:

such that $V^{*}_{ij}=V^{*}_{ik}$. Since this is true for all possible $j,k$, one can pick low-rank factorizations so that $V_{0}=V_{1}=\cdots=V_{P-1}$.

As a final remark in this section, this low complexity factorization (34) for $A$ can be easily converted to one for $A^{*}$ since

where $U,\Sigma,V$ are provided in (30), (31), and (32).

4. Neural networks

Based on the low-rank property of $A$ in Section 3.2 and its low-complexity factorization in Section 3.3, we propose new NN architectures for representing the inverse map $d\rightarrow\eta$ and the forward map $\eta\rightarrow d$.

As pointed out earlier, $d\approx A\eta$ when $\eta$ is sufficiently small. The usual filtered back-projection algorithm solves the inverse problem $d\rightarrow\eta$ via

where $\epsilon$ is the regularization parameter. In the far field pattern problem, $(A^{*}A+\epsilon I)^{-1}$ can be understood as a deconvolution operator. To see this, a direct calculation reveals that

for $x,y\in X$. (41) shows that $A^{*}A$ is a translation-invariant convolution operator. Therefore, the operator $(A^{*}A+\epsilon I)^{-1}$, as a regularized inverse of $A^{*}A$, simply performs a deconvolution. In summary, the above discussion shows that in order to obtain $\eta$ from the scattering pattern $d$ in the regime of small $\eta$, one simply needs to apply sequentially to $d$

a translation-invariant filter that performs the deconvolution $(A^{*}A+\epsilon I)^{-1}$.

Although these two steps might be sufficient when $\eta$ is small, a nonlinear solution is needed when $\eta$ is not so. For this purpose, we propose a nonlinear neural network SwitchNet for the inverse map. There are two key ingredients in the design of SwitchNet.

The first key step is the inclusion of a Switch layer that sends local information globally, as depicted in Figure 4. The structure of the Switch layer is designed to mimic the matrix-vector multiplication of the operator $A^{*}\approx V\Sigma^{*}U^{*}$ in (29). However unlike the fixed coefficients in (29), as an NN layer, the Switch layer allows for tunable coefficients and learns the right values for the coefficients from the training data. This gives the architecture a great deal of flexibility.

The second key step is to replace the linear deconvolution in the back-projection algorithm with a few convolution (Conv) layers. This enriches the architecture with nonlinear capabilities when approximating the nonlinear inverse map.

These basic building blocks of SwitchNet are detailed in the following subsection. We also take the opportunity to include the details of the pointwise multiplication PM layer that will be used in later on.

4.2. Layers for SwitchNet

In this section we provide the details for the layers that are used in SwitchNet. Henceforth, we assume that the entries of a tensor is enumerated in the Python convention, i.e., going through the dimensions from the last one to the first. One operation that will be used often is a reshape, in which a tensor is changed to a different shape with the same number of entries and with the enumeration order of the entries kept unchanged.

Swap the second and the third dimensions to get a ${\sqrt{P}}\times{\sqrt{P}}\times\frac{n}{\sqrt{P}}\times\frac{n}{\sqrt{P}}$ tensor.

Swap the second and the third dimensions to get a ${\sqrt{P}}\times\frac{n}{\sqrt{P}}\times{\sqrt{P}}\times\frac{n}{\sqrt{P}}$ tensor.

Here the non-zero entries of $U,V$ are the trainable parameters. The Switch layer is illustrated in Figure 4.

We remark that, among these layers, the Switch layer has the most parameters. If the input and output to the Switch layer both have size $n\times n$, the number of parameter is $2tPn^{2}$ where $P$ is the number of squares that partition the input field and $t$ is the rank of the low-rank approximation.

4.3. NN for the forward map η→d→𝜂𝑑\eta\rightarrow d.

We move on to discuss the parameterization of the forward map $\eta\rightarrow d$. The proposal is based on the simple observation that the inverse of the inverse map is the forward map.

More precisely, we simply reverse the architecture of the inverse map proposed in Algorithm 1. This results in an NN presented Algorithm 2. The basic architecture of this NN involves applying a few layers of Conv first, then followed by a Switch layer that mimics $A\approx U\Sigma V^{*}$.

We would also like to mention yet another possibility to parameterize the forward map $\eta\rightarrow d$, via a recurrent neural network . Let

5. Numerical results

In this section, we present numerical results of SwitchNet for far field pattern at a frequency $\omega\approx 60$. The scatterer field $\eta(x)$ supported in $\Omega=[-0.5,0.5]^{2}$ is assumed to be a mixture of Gaussians

In Algorithm 1, the parameters are specified as $t=3$ (rank of the low-rank approximation), $P_{X}=8^{2}$, $P_{D}=4^{2}$, $w=10$ (window size of the convolution layers), $\alpha=18$ (channel number of the convolution layers), and $L=3$ (number of convolution layers), resulting 3100K number of parameters. The parameters for Algorithm 2 are chosen to be $t=4$, $P_{X}=8^{2}$, $P_{D}=4^{2}$, $w=10$, $\alpha=24$, and $L=3$, with a total of 4200K parameters. Note that for both algorithms the number of parameters is significantly less than the one of a fully connected NN, which has at least $80^{4}=40960$K parameters.

SwitchNet is trained with the ADAM optimizer in Keras with a step size of 0.002 and a mini-batch size of size 200. The optimization is run for 2500 epochs. Both the training and testing data sets are obtained by numerically solving the forward scattering problem with an accurate finite difference scheme with a perfectly matched layer. In the experiment, 12.5K pairs of $(\eta,d)$ are used for training, and another 12.5K pairs are reserved for testing. The errors are reported using the mean relative errors

Table 1 summarizes the test errors for Gaussian mixtures with different choices of $n_{s}$. For the purpose of illustration, we show the predicted $d$ and $\eta$ by SwitchNet along with the ground truth in Figure 5 for one typical test sample.

SwitchNet for seismic imaging

This section considers a two-dimensional model problem for seismic imaging. The scatterer $\eta(x)$ is again assumed to be supported in a domain $\Omega$ with an $O(1)$ diameter, after appropriate rescaling. $\Omega$ is discretized with a Cartesian grid $X=\{x\}_{x\in X}$ at the rate of at least a few point per wavelength. Compared to the source and receiver configurations in Section 3.1, the experiment setup here is simpler. One can regard both $S=\{s\}_{s\in S}$ and $R=\{r\}_{r\in R}$ to be equal to a set of uniformly sampled points along a horizontal line near the top surface of the domain. The support of $\eta$ is at a certain distance below the top surface so that it is well-separated from the sources and the receivers (see Figure 6 for an illustration of this configuration).

The source and receiver operators in (9) take a particularly simple form. For the sources, the operator $(G_{0}\Pi_{S})$ is simply given by sampling:

Similarly for the receivers, the operator $(\Pi_{R}^{T}G_{0})$ is given by

After plugging these two formulas back into (9), one arrives at the following representation of the observation data $d(r,s)$ for $r\in R$ and $s\in S$

2. Low-rank property

Following the approach taken in Section 3.2, we start with the linear approximation under the assumption that $\eta(x)$ is weak. Since $E=\text{diag}(\eta)$, the first order approximation is

where the element $A$ at $(r,s)\in R\times S$ and $x\in X$ is given by

Under the assumptions that the sources $S$ and receivers $R$ are well-separated from the support of $\eta(x)$ and that $c_{0}(x)$ varies smoothly, the matrix $A$ satisfies a low-rank property similar to Theorem 1. To see this, we again partition $X$ into Cartesian squares $X_{0},\ldots,X_{P_{X}-1}$ of side-length equal to $1/\sqrt{\omega}$. Since $R=S$ is now the restriction of $X$ on the surface level, this partition also induces a partitioning for $R\times S$. When $c_{0}(x)$ varies smoothly, it is shown (see for example ) that the restriction of the matrix $[G_{0}(r,x)]_{r\in R,x\in X}$ (or $[G_{0}(s,x)]_{s\in S,x\in X}$) to each piece of the partitioning is numerically low-rank. Since the matrix $A$ is obtained by taking the Khatri-Rao product of $[G_{0}(r,x)]_{r\in R,x\in X}$, $[G_{0}(s,x)]_{s\in S,x\in X}$, the low-rank property is preserved with the guarantee that the rank at most squares in the worst case.

By following the same argument in Section 3.3, one can show that the matrix $A$ has a low-complexity matrix factorization $A\approx U\Sigma V^{*}$ of exactly the same structure as (29). The corresponding factorization for $A^{*}$ is $A^{*}\approx V\Sigma^{*}U^{*}$.

3. Neural networks

Based on the low-rank property in Section 4.2, we propose here SwitchNet for seismic imaging.

When the linear approximation is valid (i.e., (52) holds) $\eta$ can be obtained from $d$ via a filtered projection approach (or called migration in the seismic community)

where $\epsilon I$ is a regularizing term. Since $A^{*}$ has a low-complexity factorization $A^{*}\approx V\Sigma^{*}U^{*}$, the application $A^{*}$ to a vector can be represented by a Switch layer.

Concerning the $(A^{*}A+\epsilon I)^{-1}$ term, note that

which, unlike (41), is no longer a translation-invariant kernel as the data gathering setup is not so. For example, even when the background velocity $c_{0}(x)=1$, the different terms of the Green’s function $G_{0}(\cdot)$ in (55) scale like

which fail to give a translation-invariant kernel of form $K(x-y)$. As a direct consequence, the operator $(A^{*}A+\epsilon I)^{-1}$ is not translation-invariant either.

In order to capture the loss of translation-invariance, we include an extra pointwise multiplication layer PM (defined in Section 3.4.2) when dealing with the inverse map. The pseudo-code of the NN for the inverse map is given in Algorithm 3.

3.2. NN for the forward map η→d→𝜂𝑑\eta\rightarrow d

As in Section 3.4.3, for the forward map from $\eta\rightarrow d$, we simply reverse the architecture of the NN for the inverse map in Algorithm 3. For completeness we detail its structure in Algorithm 4. The main difference between Algorithm 2 and Algorithm 4 is again the inclusion of an extra pointwise multiplication layer.

4. Numerical results

In the numerical experiments, we set $\Omega=[-0.5,0.5]^{2}$ and discretize it by a $64\times 64$ Cartesian grid. As mentioned before, the sources $S$ and the receivers $R$ are located on a line near the top surface of $\Omega$, similar to the setting in Fig. 6. This line is discretized uniformly with $M=80$ points. Therefore, the size of $\eta$ and $d$ are $64\times 64$ and $80\times 80$, respectively. We assume a Gaussian mixture model for $\eta$ as in (49), where $\beta=0.2,\sigma=0.015$. Unlike before, the centers $\{c_{i}\}_{i=1}^{n_{s}}$ are kept away from the top surface of $\Omega$ in order to ensure that they are well-separated from the sources and receivers.

In Algorithm 3 and Algorithm 4, the parameters are set to be $t=3$, $P_{X}=8^{2}$, $P_{D}=4^{2}$, $N=64$, $M=80$, $w=8$, $\alpha=18$, and $L=3$, resulting NNs with 2900K parameters. The procedure of training the NNs is the same as the one used in Section 3.5. Table 2 presents the test errors for this model problem. The predicted and the ground truth $d,\eta$ are visually compared in Figure 7 for one typical test sample.

Discussion

In this paper, we introduce a neural network, SwitchNet, for approximating forward and inverse maps arising from the time-harmonic wave equation. For these maps, local information at the input has a global impact at the output, therefore they generally require the use of a fully connected NN in order to parameterize them. Based on certain low-rank property that arises in the linearized operators, we are able to replace a fully connected NN with the sparse SwitchNet, thus reducing complexity dramatically. Furthermore, unlike convolutional NNs with local filters, the proposed SwitchNet connects the input layer with the output layer globally. This enables us to represent highly oscillatory wave field resulted from scattering problems, and to solve for the associated inverse problems.

Acknowledgments

The work of Y.K. and L.Y. is partially supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) program and the National Science Foundation under award DMS-1818449. Y.K. thanks Prof. Emmanuel Candès for the partial support from a Math+X postdoctoral fellowship. This work is also supported by the GCP Research Credits Program from Google.

References