[论文阅读]基于分层矩阵的多尺度神经网络

论文:A multiscale neural network based on hierarchical matrices
作者：Yuwei Fan∗ , Lin Liny , Lexing Yingz , Leonardo Zepeda-N´u~nezx
文章的主要工作是将$\mathcal{H}$矩阵的思想引入神经网络，从而求解非线性偏微分方程。文章结构如下：先提出$\mathcal{H}$矩阵的概念，然后将$\mathcal{H}$矩阵于神经网络对应起来，介绍了一维和二维的情况（重点是一维），最后提出多尺度神经网络。

$\mathcal{H}$矩阵

考虑方程
$$u(x)=\int_{\Omega}{g(x,y)v(y)dy},\Omega=[0,1)$$

$u,v$是周期函数，$g(x,y)$是smooth and numerically low-rank away from the diagonal，对区间进行$N=2^Lm$等分，将其离散化，我们可以得到线性方程组

$$u=Av$$

其中$A\in\mathbb{R}^{n\times n},u,v\in\mathbb{R}^n$
我们引入$L+1$层的hierarchical dyadic decomposition，首先引入第0层的分解，对应的网格点的集合为

$$\mathcal{I}^{(0)}=\{k/N:k=0,\cdots,N-1\}$$

在第$\ell$层，我们把网格分成$2^\ell$个不相交的segment
每个segment是由$\mathcal{I}^{(\ell)}_i=\mathcal{I}^{(0)}\cap[(i-1)/2^\ell,i/2^\ell),i=1,\cdots,2^l$确定的，用$\mathcal{I}^{(\ell)}(\text{或}\mathcal{J}^{(\ell)})$代表第$l$级的一个segment，如果$\mathcal{I}^{(\ell-1)}\subset\mathcal{J}^{(\ell-1)}(k>1)$，我们称$\mathcal{J}^{(\ell)}$是$\mathcal{I}^{(\ell-1)}$的parent，对应的，称$\mathcal{I}^{(\ell-1)}$是$\mathcal{J}^{(\ell)}$的child，显然，除了第$L$层，每一层的segment都有两个child segment

此外，对于在$\ell>2$层的一个segment$\mathcal{I}$，作如下定义：
NL($\mathcal{I}$) neighbor list of $\mathcal{I}$，在第$\ell$级上与$\mathcal{I}$相邻的segment（含自身）
IL($\mathcal{I}$) interaction list of $\mathcal{I}$，如果$\ell=2$，那么IL($\mathcal{I}$)是level2上除去NL($\mathcal{I}$)的segment，如果$\ell>2$，IL($\mathcal{I}$)包含了NL($\mathcal{P}$)的所有的child segment，这里的$\mathcal{P}$是$\mathcal{I}$除去NL($\mathcal{I}$)的parent segment

对于一个向量$v\in\mathbb{R}^n$，$v_\mathcal{I}$表示$v$中用$\mathcal{I}$标识的那部分元素，对于一个矩阵$A\in\mathbb{R}^{N\times N}$，$A_{\mathcal{I}, \mathcal{J}}$表示的是$\mathcal{I}\times\mathcal{J}$所框住的子矩阵，利用这些符号，我们得到矩阵$A$的一个分解


$\mathcal{H}$矩阵的核心是用低秩矩阵去近似非零块$A^{(\ell)}$，在本文，为了简单，我们假设每个非零块的秩最多是$r$，所以我们有

$$A_{\mathcal{I}, \mathcal{J}}^{(\ell)} \approx U_{\mathcal{I}}^{(\ell)} M_{\mathcal{I}, \mathcal{J}}^{(\ell)}\left(V_{\mathcal{J}}^{(\ell)}\right)^{T}, \quad U_{\mathcal{I}}^{(\ell)}, V_{\mathcal{J}}^{(\ell)} \in \mathbb{R}^{N / 2^{\ell} \times r}, \quad M_{\mathcal{I}, \mathcal{J}}^{(\ell)} \in \mathbb{R}^{r \times r}$$

所以矩阵$A$可以写为
$$A=\sum_{\ell=2}^{L} A^{(\ell)}+A^{(\mathrm{ad})} \approx \sum_{\ell=2}^{L} U^{(\ell)} M^{(\ell)}\left(V^{(\ell)}\right)^{T}+A^{(\mathrm{ad})}$$

之前的矩阵乘法可以写为

$$u \approx \sum_{\ell=2}^{L} u^{(\ell)}+u^{(\mathrm{ad})}=\sum_{\ell=2}^{L} U^{(\ell)} M^{(\ell)}\left(V^{(\ell)}\right)^{T} v+A^{(\mathrm{ad})} v$$

性质1：
$\bullet$ $U^{l}$和$V^{l}$是准对角矩阵，对角块的大小为$N/2^\ell\times r$
$\bullet$$A^{(ad)}$是块循环三对角阵
$\bullet M^{(\ell)},l=1,2,\cdots,L$是块循环带状矩阵，其中$block$大小是$r\times r$， band大小是$n_b^{(\ell)}$，当$\ell=2$时，band为2，$\ell\ge3$时band为3

局部全连接网络（Locally connected network）

我们希望是利用神经网络来解释$\mathcal{H}$矩阵的结构。对于一维的情况，我们将$\xi$视为一个二维的张量，大小是$\alpha\times N_x$, 其中的$\alpha$是通道数，$N_x$是空间的维数,通过LCN，我们可以得到另一个二维的张量$\zeta$，大小为$\alpha^\prime\times N_x^\prime$，而LCN的表达式如下
$$\zeta_{c^\prime, i}=\phi\Bigg(\sum_{j=(i-1)s+1}^{(i-1)s+w}{\sum_{c=1}^{\alpha}{W_{c^\prime,c;i,j}\xi_{c,j}}}+b_{c^{\prime},i}\Bigg),i=1,\cdots,N_x^\prime,c^\prime=1,\cdots,\alpha^{\prime}$$

其中$\alpha,\alpha^\prime$代表通道维数，$N_x,N_x^\prime$代表空间的维度，$w$叫核窗口大小，$s$叫步长
对于权重$W_{c^\prime,c;i,j}$，可以看成是四维的张量，四个维度分别代表了输出的通道数，输入的通道数，输出的空间维度，输入的空间维度。
在网络中，$N_x, \alpha, N_x^\prime,\alpha^\prime,w,s$都是需要指定的参数，定义三种不同LC网络

LCR Restriction network：设$s=w=\frac{N_x}{N_x^\prime},\alpha=1$，记为$\text{LCR}[\phi;N_x,N_x^\prime,\alpha^\prime]$

$$\zeta_{c',i}=\phi\Bigg(\sum_{j=(i-1)\frac{N_x}{N_x^\prime}+1}^{i*\frac{N_x}{N_x^\prime}}{W_{c',1;i,j}\xi_{1,j}}+b_{c',i}\Bigg)$$

这可以看成是准对角矩阵和向量的乘法。

LCK Kernal network: 设步长$s=1,N_x=N_x^\prime$，记为$\text{LCK}[\phi;N_x,\alpha,\alpha^\prime,w]$

$$\zeta_{c^\prime, i}=\phi\Bigg(\sum_{j=i}^{i+w-1}{\sum_{c=1}^{\alpha}{W_{c^\prime,c;i,j}\xi_{c,j}}}+b_{c^{\prime},i}\Bigg),i=1,\cdots,N_x^\prime,c^\prime=1,\cdots,\alpha^{\prime}$$

这个网络可以看成是一个循环分块带状矩阵和一个向量的乘积，其中分块矩阵的块的大小是$\alpha\times\alpha^\prime$，band size是$\frac{w-1}{2}$，向量长度为$N_x\alpha$

LCI Interpolation network，设$s=w=1,N_x=N_x^\prime$，记为$\text{LCI}[\phi;N_x,\alpha,\alpha^\prime]$

$$ \zeta_{c',i}=\phi\Bigg(\sum_{c=1}^{\alpha}{W_{c',c;i,i}\xi_{c,i}+b_{c',i}}\Bigg) $$

这也可以看成是一种准对角矩阵和向量的乘积

神经网络表示
为了将公式
$$u \approx \sum_{\ell=2}^{L} u^{(\ell)}+u^{(\mathrm{ad})}=\sum_{\ell=2}^{L} U^{(\ell)} M^{(\ell)}\left(V^{(\ell)}\right)^{T} v+A^{(\mathrm{ad})} v$$

用神经网络表示出来，我们用以下四个运算来表示

$$ \begin{align*} \xi^{(l)}&=(V^{(l)})^Tv,\\ \zeta^{(l)}&=M^{(l)}\xi^{(l)},\\ u^{(l)}&=U^{(l)}\zeta^{(l)},\\ u^{(\text{ad})}&=A^{(\text{ad})}v \end{align*} $$

根据性质1的第一条，我们可以得到

$$\xi^{(l)}=\text{LCR}[\text{linear};N,2^\ell,r](v),u^{(\ell)}=\text{LCI}[\text{linear};2^\ell,r,\frac{N}{2^{\ell}}](\zeta^{(\ell)})$$

根据第三条，我们有
$$\zeta^{(\ell)}=\text{LCK}[\text{linear};2^\ell,r,r,2n_b^{(\ell)}+1](\xi^{(\ell)})$$

我们用$\text{Reshape}[n_1,n_2]$表示将一个$n_1n_2$长度的向量转成一个$n_1\times n_2$大小的张量，Flatten表示它的逆运算，利用性质1的第二条，我们可以把第四个式子写成

$$ \begin{align*} \tilde{v}&=\text{Reshape}[m,2^L](v),\\ \tilde{u}^{(ad)}&=\text{LCK}\Big[\text{linear};2^L,m,m,2n_b^{(ad)}+1\Big](\tilde{v}),\\u^{(ad)} &= \text{Flatten}(\tilde{u}^{(ad)}) \end{align*} $$

参数量分析

可以直接看出三种不同的LCN的参数量为
$$N_p^{\text{LCR}}=N_x\alpha^\prime,\quad N_p^{\text{LCK}}=N_x\alpha\alpha^\prime w,\quad N_p^{\text{LCI}}=N_x\alpha\alpha^\prime$$

所以算法1的参数量为

$$ \begin{aligned} N_{p}^{\mathcal{H}} &=\sum_{\ell=2}^{L}\left(N r+2^{\ell} r^{2}\left(2 n_{b}^{(\ell)}+1\right)+2^{\ell} r \frac{N}{2^{\ell}}\right)+2^{L} m^{2}\left(2 n_{b}^{(\mathrm{ad})}+1\right) \\ & \leq 2 L N r+2^{L+1} r^{2}\left(2 \max _{\ell=2}^{L} n_{b}^{(\ell)}+1\right)+N m\left(2 n_{b}^{(\mathrm{ad})}+1\right) \\ & \leq 2 N \log (N) r+3 N m\left(2 n_{b}+1\right) \sim O(N \log (N)) \end{aligned} $$

这里的$n_b=\max{(n_b^{\text{ad}}, n_b^{\ell},\ell=1,\cdots,L)}$ 且 $r\le m$.

多维情形

这里以二维的情况为例，考虑下面的方程
$$u(x)=\int_{\Omega}{g(x,y)v(y)dy},\quad\Omega=[0,1)\times[0,1)$$

仍然使用一致网格，网格数为$n=N^2(N=2^Lm)$,所有的网格点如下
$$\mathcal{I}^{(d,0)}=\{(k_1/N,k_2/N):k_1,k_2=0,\cdots,N-1\}$$

在第$\ell$层上，我们可以将网格分解成$4^\ell$个box$\mathcal{I}^{(d,\ell)}_i=\mathcal{I}_{i_1}^{(\ell)}\bigotimes\mathcal{I}_{i_2}^{(\ell)},i_1,i_2=1,2,\cdots,2^\ell$
类似可以定义$\text{NL}$和$\text{IL}$，同样对离散的矩阵也有前面类似的分解，对于非对角线上的块状矩阵，同样有下面的近似
$$A_{\mathcal{I}, \mathcal{J}}^{(\ell)} \approx U_{\mathcal{I}}^{(\ell)} M_{\mathcal{I}, \mathcal{J}}^{(\ell)}\left(V_{\mathcal{J}}^{(\ell)}\right)^{T}, \quad \mathcal{I}, \mathcal{J} \in \mathcal{I}^{(\ell)}, \quad U_{\mathcal{I}}^{(\ell)}, V_{\mathcal{J}}^{(\ell)} \in \mathbb{R}^{\left(N / 2^{\ell}\right)^{2} \times r}, M_{\mathcal{I}, \mathcal{J}}^{(\ell)} \in \mathbb{R}^{r \times r}$$

二维的LC层是将$\alpha\times N_{x,1}\times N_{x,2}$变成$\alpha^\prime\times N^\prime_{x,1}\times N^\prime_{x,2}$，具体形式为
$$\zeta_{c^{\prime}, i}=\phi\left(\sum_{j=(i-1) s+1}^{(i-1) s+w} \sum_{c=1}^{\alpha} W_{c^{\prime}, c ; i, j} \xi_{c, j}+b_{c^{\prime}, i}\right), \quad i_{1}=1, \ldots, N_{x, 1}^{\prime}, i_{2}=1, \ldots, N_{x, 2}^{\prime}, c^{\prime}=1, \ldots, \alpha^{\prime}$$

我们可以类似地定义之前提到的三种不同类型的LCN，只不过这里的$N_x=\{N_{x,1},N_{x,2}\},s=\{s_1,s_2\},w=\{w_1,w_2\}$，一维情况下的$s=w=\frac{N_x}{N_x^\prime}$相当于$s_i=w_i=\frac{N_{x,j}}{N_{x,j}^\prime}$
二维情况下的算法如下

其中的$\text{Reshape}[r^2,n_1,n_2]$是指将大小为$rn_1\times rn_2$的二维张量转化为$r^2\times n_1\times n_2$的三维张量

多尺度神经网络

多尺度神经网络是对前面的扩展，使用前面的形式来表示一个nonlinear generalization of pseudo-differential operators of the form
$$u=\mathcal{M}(v),u,v\in\mathcal{R}^{N^d}$$

这里只考虑一维的情况，

和前面的算法1比较，这里是引入了非线性的激活函数，且非线性激活函数主要是用在LCK层，网络的参数量估计如下

$$ \begin{aligned} N_{p}^{\mathrm{MNN}} &=\sum_{\ell=2}^{L}\left(N r+K 2^{\ell} r^{2}\left(2 n_{b}^{(\ell)}+1\right)+2^{\ell} r \frac{N}{2^{\ell}}\right)+K 2^{L} m^{2}\left(2 n_{b}^{(\mathrm{ad})}+1\right) \\ & \leq 2 L N r+K 2^{L+1} r^{2}\left(2 \max _{\ell=2}^{L} n_{b}^{(\ell)}+1\right)+N K m\left(2 n_{b}^{(\mathrm{ad})}+1\right) \\ & \leq 2 N \log (N) r+3 N K m\left(2 n_{b}+1\right) \end{aligned} $$

平移不变的情况

对于方程
$$u(x)=\int_{\Omega}{g(x,y)v(y)dy},\Omega=[0,1)$$

如果$g(x,y)$是平移不变的，如$g(x,y)=g(x-y)$，这时候矩阵$A$是一个Toeplitz矩阵（在同一条对角线上的元素相等），这时候$M^{(\ell)}$和$A^{\text{ad}}$也是Toeplitz矩阵，准对角阵$U^{(\ell)}$（或$V^{(\ell)}$）对角线上的分块矩阵是相等的，这时候LCN相当于一个卷积网络

$$ \zeta_{c^{\prime}, i}=\phi\left(\sum_{j=(i-1) s+1}^{(i-1) s+w} \sum_{c=1}^{\alpha} W_{c^{\prime}, c ; j} \xi_{c, j}+b_{c^{\prime}}\right), \quad i=1, \ldots, N_{x}^{\prime}, c^{\prime}=1, \ldots, \alpha^{\prime}$$ 这时候$W$和$b$与$i$无关，类似可以定义出CR，CK和CI等网络层，把算法1中的LCN层换成卷积层，我们就可以得到核函数$g(x,y)$是平移不变的情况下的算法。 这种情况可以扩展到当核函数是 equivariant to translation，即对于任意的translation $\mathcal{T}$ $$\mathcal{T}\mathcal{M}(v)=\mathcal{M}(\mathcal{T}v)$$ 这是算法3中的LCN层可以替换成卷积层，CR，CK和CI等网络层的参数量为 $$N_p^{\text{CR}}=\frac{N_x}{N_x^\prime}\alpha^\prime,\quad N_p^{\text{CK}}=\alpha\alpha^\prime w,\quad N_p^{CI}=\alpha\alpha^\prime$$ 这时候算法3的参数估计为 $$\begin{aligned} N_{p,\text{CNN}}^{\text{MNN}}&=\sum_{\ell=1}^{L}{\Big(r\frac{N}{2^\ell}+Kr^2(2n_b^{(\ell)}+1)+r\frac{N}{2^\ell}\Big)}+Km^2(2n_b^{\text{ad}}+1)\\ &\le rN+(r^2\log{N}+m^2)(2n_b+1)K. \end{aligned}$$