reply
**I appreciate the authors' response and explanation.**
**I looked at the proof of Theorem 3.1 and I am unable to fill in the gaps where diffusion terms should be. It is possible that the authors chose an unfortunate simplifying notation and that this is correct, but I cannot verify the result right now. I would also be surprised at true monotonicity for a stochastic quantity. In Corollary 3.3, the authors seem to use the result in an ODE fashion to conclude that |u-v|^2 in fact decays to zero.**
Thanks for this question. We clarify the details below. We chose a simplified notation in the original proof to avoid confusing the readers with too much technical details. More technically, the proof should start with Eq. (1), and the standard SDE formulation. The following is how one would arrive at the result in this route.
Also, a core aspect of the result of Theorem 3.1 is exactly that Eq. (2) is an ODE! This is because the diffusion term cancel each other for the time evolution of $\|u\|^2-\|w\|^2$. The mechanism for this cancellation is essential to the proof: the gradient noice covariance is low-rank due to symmetry.
Now, we outline the proof of Theorem 3.1 in a more technical and standard manner. By defining $\theta=(u^T,w^T)^T$, we can rewrite the quantity $\|u\|^2-\|w\|^2$ as $\theta^TB\theta$, where $B:=\quad
\begin{pmatrix} I_u & O \\\\ O & -I_w \end{pmatrix}
\quad$. The dynamics of $\theta$ is described by Eq. (1): $d\theta/dt=-\nabla_{\theta}L+\sqrt{TC(\theta)}dW/dt$, where $C(\theta):=\mathbb{E}[\nabla_{\theta}\ell\nabla_{\theta}^T\ell]-\mathbb{E}[\nabla_{\theta}\ell]\mathbb{E}[\nabla_{\theta}\ell]$. To obtain the dynamics of $\theta^TB\theta$, we use Ito's lemma: If the vector $X_t$ satisfies the stochastic process $dX_t=\mu_tdt+D_tdW_t$, then the dynamics of a function $f(X_t)$ of $X_t$ can be written as $d f(X_t) = \left(\nabla_X^T f \mu_t + \frac{1}{2}\mathrm{Tr}[D_t^T \nabla_X^2 f(X_t) D_t] \right)dt + \nabla_X^Tf(X_t) D_t dW_t$. Hence, the dynamics of $G(\theta)=\theta^TB\theta$ is given by $dG(\theta)/dt=-\theta^TB\nabla_{\theta}L+\theta^TB\sqrt{TC(\theta)}dW/dt+T\mathrm{Tr}[C(\theta)B]$. To simplify this, we use the infinitesimal form of the rescaling symmetry, which can be expressed as $\ell(\theta,x)=\ell(A\theta,x)$ with the matrix $A:=\quad
\begin{pmatrix} (1+\epsilon) I_u & O \\\\ O & (1-\epsilon)I_w \end{pmatrix}
\quad$. We expand the equation to first order in $\epsilon$ and obtain $\theta^TB\partial\ell/\partial\theta=0$. By taking expectation to both sides, we have $\theta^TB\partial L/\partial\theta=0$. In addition, $\theta^TBC(\theta)=\mathbb{E}[\theta^TB\nabla_{\theta}\ell\nabla_{\theta}^T\ell]-\mathbb{E}[\theta^TB\nabla_{\theta}\ell]\mathbb{E}[\nabla_{\theta}\ell]=0$. Therefore, we can see $\theta^TB\sqrt{C(\theta)}=0$ since $C(\theta)$ and $\sqrt{C(\theta)}$ share the same null space. By substituting $\theta^TB\partial L/\partial\theta=0$ and $\theta^TB\sqrt{C(\theta)}=0$ into the evolution equation of $G(\theta)$, we have $dG(\theta)/dt=T\mathrm{Tr}[C(\theta)B]$, which is nothing but Eq. (30). Here we can see the diffusion terms related to $dW$ vanish due to the rescaling symmetry.
**It is still not entirely clear to me how to show that the loss only depends on the rank 1 matrix and how the 'lifted' function would be defined on the space of matrices (not just the set of rank 1 matrices, which does not form a linear subspace). An explicit construction or formula would be appreciated.**
Due to the rescaling symmetry, the loss function only depends on this rank 1 matrix. To show this, we redefine a new function $\tilde{\ell}(u_iw_j,w_j,x):=\ell(u_i,w_j,x)$ to reorganize the parameters. Then the derivatives of $\ell$ can be rewritten as $\frac{\partial \ell}{\partial w_j}=\sum_{i}u_i\frac{\partial \tilde{\ell}}{\partial(u_iw_j)}$ and $\frac{\partial\ell}{\partial u_i}=\sum_{j}w_j\frac{\partial \tilde{\ell}}{\partial(u_iw_j)}+\frac{\partial \tilde{\ell}}{\partial(w_j)}$. However, due to the rescaling symmetry $\tilde{\ell}(u_iw_j,w_j,x)=\tilde{\ell}(u_iw_j,\lambda w_j,x)$, we have $\sum_jw_j\partial\tilde{\ell}/\partial(w_j)=0$. Since this equality holds for an arbitrary $w_j$, we always have $\partial\tilde{\ell}/\partial(w_j)=0$ for all $j$. Hence, the loss function only relies on the rank 1 matrix $u_iw_j$. We will clarify this in the revision.
Please ask for additional clarification for any point that is not yet clear.