I'm going through Why Momentum Really Works and am unable to understand the following line in the article.

"By writing the contributions of each eigenspace’s error to the loss $$f(w^{k})-f(w^{\star})=\sum(1-\alpha\lambda_{i})^{2k}\lambda_{i}[x_{i}^{0}]^2$$ ​​ we can visualize the contributions of each error component to the loss."

I am not able to figure out how the above equation was derived despite fully understanding how:

$$w^k - w^\star = Qx^k = \sum\limits_{i}^n x_i^0(1-\alpha\lambda_i)^kq_i$$

was derived and also understanding that :

$$f(w) = \frac12w^TAw-b^Tw, \quad w \in \mathbb{R}^n, A \in \mathbb{R}^{n,n}$$ has gradient: $$\nabla f(w) = Aw-b$$ where $A$ is symmetric and invertable and thus $f$ has optimal solution $w^\star = A^{-1}b$. ​​

I've already seen below related posts but these posts do not explain the top equation which I am trying to understand the derivation of.

Gradient descent of $f(w)=\frac12w^TAw-b^Tw$ viewed in the space of Eigenvectors of $A$

Gradient descent derivation in Eigenspace


This is just the result of a bit of tedious algebra. From the other definitions in the article, we have the following identities: $$w^* = A^{-1}b$$ $$w^k = A^{-1}b + Qx^k$$ One identity they assume you know is that $Q^TQ = I$, which implies: $$ Q^T A Q = Q^T Q \text{diag}[\lambda_1,\dots,\lambda_n]Q^T Q = \text{diag}[\lambda_1,\dots,\lambda_n] $$

The next steps involve writing out $f(w^k)-f(w^*)$ using the identities above and cancelling out terms. You should be able to cancel all but one term, leaving you with the final result: $$f(w^k) - f(w^*) = \frac{1}{2}\sum (x_i^k)^2\lambda_i = \frac{1}{2} \sum \lambda_i (1-\alpha \lambda_i)^{2k} [x_{i}^0]^2$$ I wound up off by a factor of 1/2 (I'm pretty sure the article left this factor out, since it's not that important).

[Edit] Tedious details are below for others who may get stuck.

We start by re-writing the loss function in terms of the quadratic and linear terms: $$\begin{aligned} f(w^k) - f(w^*) &= \frac{1}{2}(w^k)^TAw^k - b^T w^k \\ &\ - \frac{1}{2}(w^*)^T A w^* + b^T w^*\\ & = \frac{1}{2}\left( (w^k)^T A w^k - (w^*)^T A w^*\right) - b^T (w^k - w^*)\\ & = \frac{1}{2}\left( (w^k)^T A w^k - (w^*)^T A w^*\right) - b^TQ x^k \end{aligned}$$ Where the last line uses the definition of $x^k$. Now let's work out the quadratic terms separately to try and keep things managable. The quadratic term for $w^*$ is straightforward: $$ (w^*)^T A w^* = b^T A^{-1} A A^{-1} b = b^T A^{-1} b$$ and for $w^k$ we get: $$\begin{aligned} (w^k)^T A w^k &=(A^{-1} b + Q x^k)^T A (A^{-1}b + Q x^k)\\ &= (A^{-1}b + Q x^k)^T(b + A Q x^k)\\ &= (b^T A^{-1} + (x^k)^T Q^T) (b + A Q x^k)\\ &= b^TA^{-1}b + b^TQ x^k + (x^k)^TQ^Tb + (x^k)^T Q^T A Q x^k\\ &= b^TA^{-1}b + 2 b^TQ x^k + (x^k)^T \text{diag}[\lambda_1,\dots,\lambda_n] x^k\\ &= (w^*)^T A w^* + 2 b^TQ x^k + (x^k)^T \text{diag}[\lambda_1,\dots,\lambda_n] x^k\\ \end{aligned}$$ and so we get $$ \frac{1}{2} \left((w^k)^T A w^k - (w^*)^T A w^* \right) = b^T Q x^k + \frac{1}{2} (x^k)^T \text{diag}[\lambda_1, \dots, \lambda_n]x^k$$

Now putting this back into our expression for the loss function we have: $$\begin{aligned} f(w^k) - f(w^*) & = \frac{1}{2}\left( (w^k)^T A w^k - (w^*)^T A w^*\right) - b^TQ x^k\\ &= b^T Q x^k + \frac{1}{2} (x^k)^T \text{diag}[\lambda_1, \dots, \lambda_n]x^k - b^T Q x^k\\ &= \frac{1}{2} (x^k)^T \text{diag}[\lambda_1, \dots, \lambda_n]x^k\\ &= \frac{1}{2}\sum (x_i^k)^2\lambda_i \end{aligned}$$


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.