Global Convergence of Sobolev Training for Overparameterized Neural Networks

Abstract

Sobolev loss is used when training a network to approximate the values and derivatives of a target function at a prescribed set of input points. Recent works have demonstrated its successful applications in various tasks such as distillation or synthetic gradient prediction. In this work we prove that an overparameterized two-layer relu neural network trained on the Sobolev loss with gradient flow from random initialization can fit any given function values and any given directional derivatives, under a separation condition on the input data.

Appendices

A Supplementary proofs for Sect. 3.1

In this section we provide the remaining proofs of the results in Sect. 3.1. We begin recalling the following matrix Chernoff inequality (see for example [15, Theorem 5.1.1]).

Theorem 3

(Matrix Chernoff). Consider a finite sequence $X_k$ of $p \times p$ independent, random, Hermitian matrices with $0 \preceq X_k \preceq L I$. Let $X = \textstyle {\sum }_k X_k$, then for all $\epsilon \in [0,1)$

$$\begin{aligned} \mathbb {P}\Big [\lambda _{\text {min}}(X) \le \epsilon \lambda _{\text {min}}\big (\mathbb {E}[X]\big ) \Big ] \le p e^{-(1-\epsilon )^2 \lambda _{\text {min}}(\mathbb {E}[X]) /2L} \end{aligned}$$

(16)

In order to lower bound the smallest eigenvalue of H(0) we use Lemma 1 together with the previous concentration result.

Proof

(Lemma 2). We first note that $\mathbb {E}[H(0)] = \mathbb {E}[\sum _r H_r(0)] = H^{\infty }$, and moreover $H_r(0)$ is symmetric positive semidefinite with $\lambda _{\text {max}}(H_r) \le n (k+1)/m$ by Lemma 1. Applying then the concentration bound (16) with the assumption $m \ge \frac{32}{{{\lambda }_*}}\, n(k+1) \ln ( n(k +1)/\delta )$ gives the thesis.

We next upper bound the errors at initialization.

Proof

(Lemma 3). Note that for any $x_i$, due the the assumption on the independence of the weights at initialization and the normalization of the data:

$$ \mathbb {E}[(f(W,x_i))^2] = \sum _{r=1}^m \frac{1}{m} \mathbb {E}[\sigma (w_r^T x_i)^2] \le 1 $$

and similarly for the directional derivatives

$$ \mathbb {E}[ \Vert \bar{F}(W,x_i)\Vert _2^2 ] = \mathbb {E}_{g \sim \mathcal {N}(0,I)} [\Vert \sigma '(g^T x_i ) V_i^T g\Vert _2^2 ] \le \sum _{j=1}^k \mathbb {E}[ (v_{i,j}^T g)^2 ] \le k. $$

We conclude the proof by using Jensen’s and Markov’s inequalities.

B Proof of Proposition 1

Consider the $d\times (k+1)$ matrices $\mathbf {X}_i = [x_i, V_i]$, and for define

$$ \hat{\psi }_w(x_i) = \sigma '(w^T x_i) \mathbf {X}_i. $$

and the $d \times (k+1)n$ matrix:

which corresponds to a column permutation of $\varOmega (w)$. Next observe that the matrix $\widehat{H}^\infty = \mathbb {E}_{w \sim \mathcal {N}(0,I_d)}[\widehat{\varOmega }(w)^T \widehat{\varOmega }(w)]$ is similar to $H^\infty$ and therefore has the same eigenvalues. In this section we lower bound ${\lambda }_*$ by analyzing $\widehat{H}^\infty$.

We begin recalling some facts about the spectral properties of the products of matrices.

Definition 2

([8]). Let $\mathbf {A} = [A_{\alpha \beta }]_{\alpha = 1, \dots , n}^{\beta = 1, \dots , n}$ and $\mathbf {B} = [B_{\alpha \beta }]_{\alpha = 1, \dots , n}^{\beta = 1, \dots , n}$ be $n p \times np$ matrices in which each block is in ${p \times p}$. Then we define the block Hadamard product of $\mathbf {A} \square \mathbf {B}$ as the $n p \times np$ matrix with:

$$ \mathbf {A} \square \mathbf {B} := [A_{\alpha \beta } B_{\alpha \beta } ]_{\alpha = 1, \dots , n}^{\beta = 1, \dots , n} $$

where $A_{\alpha \beta } B_{\alpha \beta }$ denotes the usual matrix product between $A_{\alpha \beta }$ and $B_{\alpha \beta }$.

Generalizing Schur’s Lemma one has the following regarding the eigenvalues of the block Hadamard product of two block matrices.

Proposition 2

([8]). Let $\mathbf {A} = [A_{\alpha \beta }]_{\alpha = 1, \dots , n}^{\beta = 1, \dots , n}$ and $\mathbf {B} = [B_{\alpha \beta }]_{\alpha = 1, \dots , n}^{\beta = 1, \dots , n}$ be $n p \times np$ positive semidefinite matrices. Assume that every $p\times p$ block of $\mathbf {A}$ commutes with every $p\times p$ block of $\mathbf {B}$, then:

$$ \lambda _{\text {min}}(\mathbf {B} \square \mathbf {A})= \lambda _{\text {min}}(\mathbf {A} \square \mathbf {B}) \ge \lambda _{\text {min}}(A) \cdot \min _{\alpha }\lambda _{\text {min}}(B_{\alpha \alpha }) $$

We finally recall the following on the eigenvalues of Kronecker product of matrices.

Proposition 3

([11]). Let with eigenvalues $\{\lambda _i\}$ and with eigenvalues $\{\mu _i\}$, then Kronecker product $A \otimes B$ between A and B has eigenvalues $\{\lambda _i \mu _j \}$.

We next define the following random kernel matrix.

Definition 3

Let $w \sim \mathcal {N}(0,I)$ then define the random matrix with entries $[\mathcal {M}(w)]_{ij}= \sigma '(w^T x_i) \sigma '(w^T x_j)$.

The next result from [12] establishes positive definiteness of this matrix in expectation, under the separation condition (7).

Lemma 5

([12]). Let $x_1, \dots , x_d$ in with unit Euclidean norm and assume that (7) is satisfied for all $i = 1, \dots d$. Then the following holds:

$$ \mathbb {E}_{w \sim \mathcal {N}(0,I)} [\mathcal {M}(w)] \succeq \frac{\delta _1 }{100 n^2} $$

Finally let block matrix with $d\times (k+1)$ blocks $\mathbf {X}_i$. Thanks to the assumption (8) the following result on the Gram matrices $\mathbf {X}_i^T\mathbf {X}_i$ holds.

Lemma 6

Assume that the condition (8) is satisfied, then for any $i = 1, \dots , n$ we have $\lambda _{\text {min}} (\mathbf {X}_i^T\mathbf {X}_i) \ge 1 - k \delta _2 > 0$.

Proof

The claim follows by observing that by Gershgorin’s Disk Theorem:

$$ | \lambda _{\text {min}} (\mathbf {X}_i^T\mathbf {X}_i) - 1| \le \sum _{ 1 \le j \le k} |x_i^T v_{i,j}|\le k \delta _2. $$

Finally observe that we can write:

$$ \widehat{H}^\infty = \mathbb {E}_{w \sim \mathcal {N}(0,I)}\big [ (\mathbf {X}^T \mathbf {X})\square (\mathcal {M}(w) \otimes I ) \big ]. $$

so that Proposition 2, Proposition 3, Lemma 5 and Lemma 6 allow to derive the thesis of Proposition 1.