MGiaD: Multigrid in all dimensions - Efficiency and robustness by weight sharing and coarsening in resolution and channel dimensions

Method Overview

The solution $u$ of large system of (non)-linear equations $Au=f$, with $A$ the system matrix and $f$ the right-hand side, is approximated iteratively via stationary iterative schemes $$ u^{(i+1)} = u^{(i)} + B(f- Au^{(i)}), \;\;\; i = 0,1, ...,$$ where $f - Au^{(i)}$ is called the residual.

Typically $B \approx A^{-1}$ and the (unknown) error is propagated in each iteration as $$ e^{(i)} = (I - BA) e^{(i)}.$$ The system matrix $A$ does not change with the iteration, the operator $B$ can be changed with the iteration.

_ResNet-block

The factor $(I-BA)$ structurally resembles a ResNet-block, where $A$ and $B$ are learnable convolutions. Nevertheless, the structural similar yields motivation for reusing the operators $A$ and $B$ ins ResNet-blocks. Though iterative methods are characterized by a fast decay of the error in the first iterations (smoothing), they suffer from a drastic decrease of the convergence rate. To speed up convergence, solutions from coarser resolutions are used to approximate the true solution. To that end the residual is restricted to the next coarser resolution level $\ell + 1$ as $r_{\ell + 1} = R_\ell^{\ell +1}e_{\ell}$, where the restriction operator $R_\ell^{\ell +1}$ is a pooling operation. The residual problem on the coarse grid is solved directly and with the coarse-grid solution the approximated solution on the fine grid is corrected accordingly. Recursively repeating this scheme yields a multi-grid structure.

Since we deal with non-linear problems in deep learning, an initial guess for the solution on the coarser resolution is required, given by $u^{\ell +1} = \Pi_\ell^{\ell +1}$.

Multigrid in CNNs

In deep learning we assume that data $f$ and features $u$ are related by (convolutional) mappings $A$ and $B$, which can be chosen as $$A: \mathbb{R}^{m \times n \times c} \rightarrow \mathbb{R}^{m \times n \times h}, \text{s.t. } \; A(u) =f$$ $$B: \mathbb{R}^{m \times n \times h} \rightarrow \mathbb{R}^{m \times n \times c}, \text{s.t. } \; u \approx B(f) \;$$

Where $m$ and $n$ are the input/output feature map dimensions and $c$ and $h$ the number of in- and output channels respectively.

_{Data-feature relations on resolution level $\ell$ followed by transfer to coarser resolution $\ell +1$}

Besides the structural analogy in MG the operations $A,B$ and $f$ are given and the solution $u$ is sought. On the other hand in deep learning $f$ and $u$ are given, and we want to learn $A$ and $B$. The transfer operators are pooling operators, where the $R_{\ell}^{\ell +1}$ operates on the data domain and $\Pi_{\ell}^{\ell+1}$ operators on the feature domain.

These mappings can be used to construct iterative schemes, i.e., residual blocks.

_{Weight sharing in $\text{ResNet}^4$ and $\text{MgNet}^5$; left: ResNet-blocks, no weight sharing; middle: MgNet-blocks shared $A_\ell$; right: MgNet-blocks shared layers $A_\ell$ and $B_\ell$.}

By reusing the operators $A$ and or $A$ and $B$ a significant reduction of weights with marginal loss in accuracy can be achieved, However the weight count scales quadratically with respect to the number of channels as $A$ and $B$ are convolutions with a $k^2$ of size $c \times h \times k^2$. Due to their local action, convolutions are sparse in the resolution dimension and dense in the channel dimension as each input channel couples with each output channel.

Smoothing in Channels ($\mathtt{Sic}$)

To achieve sparsity in the channel dimension we construct a hierarchical MG structure build from grouped convolutions. The number of channels is successively reduced until fully coupled convolutions are reasonable, which restore channel interaction and ensuring information exchange. On each in channel level: smoothing with shared weights ($\mathtt{Sic}$). Fully coupled convolutions on the coarsest in channel levels correspond to direct solving in MG. One smoothing in channel hierarchy replaces one MgNet-block.

\(\mathtt{SIC}\)-block on one resolution level with \(\kappa = 2\) channel levels. The convolutions \(\widehat{A}\) and \(\widehat{B}\) are in groups of size \(4\), indicated by different colors. On each level, \(\widehat{A}\) and \(\widehat{B}\) are shared between pre-smoothing (left leg) and post-smoothing (right leg), respectively. Operators on the second, coarser level are different from the first level, but also shared. The transfer operators \(\widehat{\Pi}\) and \(\widehat{R}\) halve the number of channels, and the prolongation mapping \(\widehat{P}\) doubles the number of channels.