on U-Net regression

Overview

This post explores the training dynamics of a two-layer convolutional neural network (CNN) under gradient flow. We analyze the invariance properties of the kernel-coefficient pairs and discuss the challenges of spectral bias, particularly the exponential time required to recover high-frequency information from band-limited data.

🏷️ Problem Formulation

We consider a two-layer convolutional network $f(X,t)$ governed by time-dependent coefficients $a_i(t)$ and kernels $K_i(\cdot ,t)$:

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

where $\ast$ denotes convolution and $\sigma$ is a non-linear activation function. Given a target linear affine mapping $h(X)$, we minimize the $L^2$ loss functional over a data manifold $M$:

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

🏷️ Gradient Flow and Invariance

The parameters evolve according to the gradient descent dynamics:

$$\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}$$

For homogeneous activations such as the ReLU ($\sigma (x)=x{\sigma }^{\prime }(x)$), we observe a fundamental conservation law. The difference between the coefficient energy and the kernel energy remains invariant:

$$\int K_i(w,t)\frac{d}{dt}{K}_i(w,t)dw=a_i(t)\frac{d}{dt}{a}_i(t)⟹\frac{d}{dt}\left(|a_i(t){|}^2-\Vert K_i(\cdot ,t){\Vert }^2\right)=0$$

This indicates that the training trajectory is confined to a tensor product of Lorentzian manifolds of type $(N,1)$.

🏷️ Stationary Points and Convergence

The loss function is monotonically non-increasing under gradient flow:

$$\frac{d}{dt}L(t)=-\sum _i|a_i^{\prime }(t){|}^2-\sum _i\Vert K_i^{\prime }(\cdot ,t){\Vert }^2$$

The coefficient vector $\mathbf{a}$ satisfies a non-homogeneous linear evolution:

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

where $G(t)$ is the Gram matrix (Neural Tangent Kernel) of the features. Convergence depends on the spectral properties of $G(t)$ and the ability of the basis $\sigma (\mathbf{K}\ast X)$ to resolve the target $h(X)$.

🏷️ Spectral Bias in Frequency Retrieval

In the context of image regression, we model $X$ as a frequency-truncated signal. The task of $f(X,t)$ is effectively to perform analytic continuation in the Fourier domain to retrieve high-frequency components.

Our analysis suggests a sharp transition in the training timescales:

• Low Frequencies: Captured rapidly within a threshold $ϵ$.
• High Frequencies: The transition from a frequency band $B$ to $B(1+ϵ)$ requires a training time that scales exponentially with the gap $ϵB$.

🏷️ Fourier Representation

Utilizing the Plancherel identity, we analyze the kernel initialization in the dual space. The initial kernels $K_i(x,0)$ are viewed as random Gaussian fields:

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

The distribution of the coefficients $C_{i,\zeta }$ is governed by the Law of Large Numbers for high bandwidth $B$, ensuring that the initial spectrum covers the necessary data bandwidth on $M$.

🏷️ Notes

• The Lorentzian structure implies that small initial weights can lead to “exploding” or “vanishing” gradients if not properly balanced.
• The spectral bias observed here is a specific instance of the “F-Principle” where neural networks fit low frequencies before high frequencies.

🔗 See Also

• on two-layer ReLU networks
• on convergence of graph Laplacian to manifold’s Laplacian

📚 References