18.S995: Topics in Deep Learning Theory

Linear Algebra

Let $H \in \mathbb{R}^{d \times d}$ be symmetric and positive definite, let $w \in \mathbb{R}^d$, and define $L(w) = \tfrac{1}{2} w^\top H w$.

1. Compute the gradient $\nabla L(w)$.
2. Let $w$ follow gradient descent: $w_{t+1} = w_t - \eta \nabla L(w_t)$. Solve for $w_t$ in closed form as a function of $t$.
3. For what $\eta$ does $w_t \to 0$ as $t \to \infty$? What changes if $H$ is only positive semidefinite?

Analysis

Let $f:\mathbb{R}^d\to\mathbb{R}$ be differentiable and suppose that $\nabla f$ is $L$-Lipschitz with respect to a norm $\|\cdot\|$. Let $r_0,r_1,\ldots\in\mathbb{R}^d$ satisfy $\lVert r_t\rVert\leq\varepsilon$, and let $\eta > 0$. Starting from $x_0=y_0$, define $$x_{t+1}=x_t-\eta\nabla f(x_t),\qquad y_{t+1}=y_t-\eta\qty{\nabla f(y_t)+r_t}.$$

1. Show that $\lVert x_{t+1}-y_{t+1}\rVert\leq\qty{1+\eta L}\lVert x_t-y_t\rVert+\eta\varepsilon$.
2. Using $1+x\leq e^x$, show that $\lVert x_t-y_t\rVert\leq\tfrac{\varepsilon}{L}\qty{e^{\eta Lt}-1}$.

Probability

A mean-zero random variable $X$ is $\sigma$-sub-Gaussian if, for every $\lambda \in \mathbb{R}$, $$\mathbb{E}\left[\exp\qty{\lambda X}\right] \leq \exp\qty{\frac{\lambda^2\sigma^2}{2}}.$$

1. Let $X \sim \mathcal{N}(0,\sigma^2)$. Show that $X$ is $\sigma$-sub-Gaussian.
2. Let $X_1, \ldots, X_n$ be independent mean-zero random variables such that $X_i$ is $\sigma_i$-sub-Gaussian. Show that $\sum_{i=1}^n X_i$ is $\sqrt{\sum_{i=1}^n \sigma_i^2}$-sub-Gaussian.
3. Let $X$ be $\sigma$-sub-Gaussian. Use Markov's inequality on $\exp\qty{\lambda X}$ and optimize over $\lambda$ to show that, for every $t \geq 0$, $$\mathbb{P}\qty{X \geq t} \leq \exp\qty{-\frac{t^2}{2\sigma^2}}.$$
4. Let $W \in \mathbb{R}^{m \times d}$ have i.i.d. $\mathcal{N}(0,1)$ entries. Using Lemma 1 and a union bound, prove that, for some absolute constant $C$ and any $\delta \in (0,1]$, with probability at least $1-\delta$, $$\lVert W \rVert_{\mathrm{op}} \leq C\qty{\sqrt{m}+\sqrt{d}+\sqrt{2\log\qty{1/\delta}}}.$$ Hint: Note that for fixed $\hat u, \hat v$, $\hat u^\top W\hat v \sim \mathcal{N}(0,1)$, and is therefore sub-Gaussian, so you can apply part 3.
   Note: This bound actually holds with $C=1$, but proving this requires more advanced techniques.

   Lemma 1. Let $\mathcal{N}_m \subset \mathbb{S}^{m-1}$ and $\mathcal{N}_d \subset \mathbb{S}^{d-1}$ be $1/4$-nets with $|\mathcal{N}_m| \leq 9^m$ and $|\mathcal{N}_d| \leq 9^d$. Then $$\lVert W \rVert_{\mathrm{op}} \leq 2\max_{(\hat u,\hat v) \in \mathcal{N}_m \times \mathcal{N}_d}\hat u^\top W\hat v.$$

   Proof. Let $u,v$ be the top left and right singular vectors of $W$, so that $u^\top Wv = \lVert W\rVert_{\mathrm{op}}$, and pick $\hat u \in \mathcal{N}_m$ and $\hat v \in \mathcal{N}_d$ with $\lVert u - \hat u\rVert \leq 1/4$ and $\lVert v - \hat v\rVert \leq 1/4$. Then $$\lVert W\rVert_{\mathrm{op}} = u^\top Wv = \hat u^\top W\hat v + \qty{u-\hat u}^\top Wv + \hat u^\top W\qty{v-\hat v} \leq \hat u^\top W\hat v + \tfrac{1}{2}\lVert W\rVert_{\mathrm{op}},$$ and rearranging completes the proof.

Neural Networks

Let $x \in \mathbb{R}^d$ with $\lVert x \rVert^2 = d$. Let $W \in \mathbb{R}^{m \times d}$, $b \in \mathbb{R}^m$, $A \in \mathbb{R}^{k \times m}$, and let $\sigma(z) = \max\qty{0,z}$ be the ReLU activation, applied coordinatewise. Define the logits $$f_\theta(x) = A\sigma\qty{Wx+b} \in \mathbb{R}^k, \qquad \theta = \qty{W,b,A}.$$

1. Suppose $W$ has i.i.d. $\mathcal{N}(0,1/d)$ entries, $A$ has i.i.d. $\mathcal{N}(0,1/m)$ entries, and $b=0$.
   1. Show that each coordinate of $Wx$ is $\mathcal{N}(0,1)$, and that $\mathbb{E}_{z \sim \mathcal{N}(0,1)}\left[\sigma(z)^2\right] = 1/2$.
   2. Compute $\operatorname{Cov}\left[f_\theta(x)\right] \in \mathbb{R}^{k \times k}$.
   3. What is the limiting distribution of $f_\theta(x)$ as $m \to \infty$? Hint: CLT.
2. For a label $y \in \{1, \ldots, k\}$, let $L = -\log \mathcal{S}\qty{f_\theta(x)}_y$, where $\mathcal{S}$ is the softmax function. Using the chain rule / backpropagation, compute $\nabla_A L$, $\nabla_W L$, and $\nabla_b L$. What are the ranks of $\nabla_A L$ and $\nabla_W L$?
   Hint: the gradient of $f \mapsto -\log \mathcal{S}\qty{f}_y$ is $\mathcal{S}\qty{f} - e_y$, where $e_y$ is the one-hot vector for $y$.
3. Complete the following notebook, which involves checking your answer to part 1 and using the gradients you computed in part 2 to train the model on MNIST using gradient descent. Use only NumPy, and do not add imports beyond those provided (e.g. no JAX/PyTorch).