on U-Net regression

Note

This is a short note about the training dynamics of a two-layer convolution network. Just for some understanding.

Let the network ($t$ for time variable of the dynamics) be

$$f(X,t)=\sum _{i=1}^n{a}_i(t)\sigma (K_i(\cdot ,t)\ast X),$$

where $K_i$ are convolution kernels with a finite bandwidth. Suppose the target function is $h(X)$ is a linear affine mapping, then the usual loss function is:

$$L(t)=\frac{1}{2}{\int }_M\Vert f(X,t)-h(X){\Vert }^{2}d\mu (X)$$

Gradient Flow

The gradient flow is

$$\begin{array}{rl}\frac{d}{dt}{a}_i(t)& =-{\int }_M(f(X,t)-h(X))\cdot \sigma (K_i(\cdot ,t)\ast X)d\mu (X)\\ \frac{d}{dt}{K}_i(w,t)& =-a_i(t){\int }_M(f(X,t)-h(X))\cdot \left[{\sigma }^{\prime }(K_i(\cdot ,t)\ast X)\odot X(\cdot -w)\right]d\mu (X)\end{array}$$

One thing we can observe is that $\sigma (x)=x{\sigma }^{\prime }(x)$. Therefore,

$$\frac{d}{dt}{a}_i(t)={\int }_M(f(X,t)-h(X))\cdot \left[{\sigma }^{\prime }(K_i(\cdot ,t)\ast X)\odot K_i(\cdot ,t)\ast X\right]d\mu (X)$$

We can write the convolution kernel $K_i(x,t)\ast X=\int X(x-w)K_i(w,t)dw$, thus

$$\int K_i(w,t)\frac{d}{dt}{K}_i(w,t)dw=a_i(t)\frac{d}{dt}{a}_i(t)$$

Showing that the $|a_i(t){|}^2-\Vert K_i(\cdot ,t){\Vert }^2$ is invariant for each $i\in [n]$. Therefore, the training is confined on a tensor product of Lorentzian manifolds of type $(N,1)$, where $N$ is the dimension of convolution kernel. However, this does not reveal any training dynamics. For images, asymptotically we can assume the support of kernel is finite.

First, we know the dynamics will converge to a stationary point.

$$\begin{array}{rl}\frac{d}{dt}L(t)& ={\int }_M(f(X,t)-h(X))\cdot \frac{df(X,t)}{dt}d\mu (X)\\ & =\sum _i{a}_i^{\prime }(t){\int }_M(f(X,t)-h(X))\sigma (K_i(\cdot ,t)\ast X)d\mu (X)-\sum _i{a}_i(t){\int }_M(f(X,t)-h(X))\cdot {\sigma }^{\prime }(K_i(\cdot ,t)\ast X)\odot \left(\frac{d}{dt}{K}_i(\cdot ,t)\ast X\right)d\mu (X)\\ & =-\sum _i|a_i^{\prime }(t){|}^2-\sum _i\Vert K_i^{\prime }(\cdot ,t){\Vert }^2.\end{array}$$

Regarding the coefficients $a_i(t)$, the dynamics will satisfy

$$\frac{d}{dt}\text{a}=-G(t)\text{a}+{\int }_{M}h(X)\sigma (\text{K}(\cdot ,t)\ast X)d\mu (X)$$

Then it can be solved by

$$\text{a}=\exp (-{\int }_0^{t}G(s)ds)\text{a}(0)+{\int }_0^t\exp \left(-{\int }_u^{t}G(s)ds\right){\int }_{M}h(X)\sigma (\text{K}(\cdot ,u)\ast X)d\mu (X)du$$

It means, once the Gram matrix $G(t)={\int }_M\left[\sigma (K_i(\cdot ,t)\ast X)\right]\cdot \left[\sigma (K_j(\cdot ,t))\ast X\right]d\mu (X)$ satisfies some positive definiteness, and $h(X)$ can be always captured well (in terms of eigenvalue) by the basis $\sigma (\text{K}(\cdot ,u)\ast X)$.

In our case, $X$ is an image of truncated frequencies and $h(X)$ is a continuation of $X$ (assuming we have some analyticity in frequency, which is true), and $f(X,t)$ needs to fit the extended image, essentially is to retrieve high frequencies. Our idea is, the lower frequencies are all easy to capture, but higher frequencies are difficult to retrieve in reasonable time. It should be significantly longer time (not just polynomial relation). The plan is the following:

• For lower frequencies, show the error can be under a threshold.
• For higher frequencies, show it will take unreasonable time to arrive. The transition from lower frequency band $B$ to $B(1+ϵ)$ should take exponentially long in terms of $ϵB$ gap.

Fourier Space

Since it is about convolution, it might be suitable to consider the dual space (Fourier), then Plancherel identity shows that $\Vert K_i(\cdot ,t)\Vert =\Vert \widehat{K_i}(\cdot ,t)\Vert$. The initialization of each $K_i$ can be viewed as a random Gaussian field and decomposed into

$$K_i(x,0)=\sum _{|\zeta |<B}{C}_{i,\zeta }(\omega )e^{i2\pi \zeta \cdot x}.$$

The norm $\Vert {\hat{K}}_i{\Vert }^2={\sum }_{|\zeta |<B}|C_{i,\zeta }{|}^2$ (assume support is unit square), the distribution is governed by the law of large number when $B$ is large. We only assume $B$ is large enough to cover the band of $X$ on $M$. For simplicity, we assume that the affine mapping $h$ in the following diagonal form under Fourier representation.

$$h(X)=\sum _{|\zeta |<B}{h}_{\zeta }\hat{X}(\zeta )e^{i2\pi \zeta \cdot x}$$

ReLU Activation

The ReLU activation function imposes some challenges for its nonlinearity. But formally, we still can write

$$\frac{d}{dt}{a}_i(t)=-{\int }_M\sum _{|\zeta |<B}\left(\hat{f}(\zeta ,t)-h(\zeta )\hat{X}(\zeta )\right)\overline{P_i(\zeta ,t)}dX.$$

where $\sigma (K_i(\cdot ,t)\ast X)={\sum }_{\zeta }{P}_i(\zeta ,t)e^{2\pi i\zeta \cdot x}$.

Notes

Links