• Refs, mostly TODO

    • MIT 6.7220 https://www.mit.edu/~gfarina/notes/
      • https://www.mit.edu/~gfarina/2024/67220s24_L12_newton/
    • Deep Learning book chapter 8 optimization
    • https://www.cs.toronto.edu/~rgrosse/courses/csc2541_2022/

  • Gradient descent

    • Invariant to rigid transformations (not scale)

  • Line search: Rather than a fixed learning rate, keep moving along the gradient direction until you minimize the objective

    • Since we are usually in a stochastic setting, this is actually not used

  • Second-order optimization algorithms

    • Newton’s method aka Newton-Raphson: iteratively find tangent line (or quadratic, etc.) and find its root (or min for quadratic). This is for finding zeros, but apply it to the gradient to find the minimum, ie where the first derivative is 0. This requires knowledge of the second derivatives ie Hessians
      • Often uses a quadratic instead of a line, ie the second order Taylor expansion.

    1000000443.png

    • BFGS/L-BFGS (source): quasi Newton's method. TODO

      • L is limited memory
      • Preconditions the gradient with curvature information
      • Covered in Deep Learning book

    • Line search (source)

    • Shampoo TODO

      • video
      • without gradient accumulation, they both agree and also agree with shampoo without accumulation

      image.png

    • Muon: record holder for nanoGPT (source) TODO

      • Relation to shampoo: (source)

        A connection between Muon and Shampoo (https://arxiv.org/abs/1802.09568) is that in the momentum=0 case, running Shampoo with the preconditioner updated every step and with no accumulation also yields the nearest orthogonal matrix as its update, although this has a slower runtime.

        The nearest orthogonal matrix is also equivalent to:

        1. The semi-unitary factor of the update's polar decomposition.
        2. UV^T where USV^T is the update's SVD.

        Using a rescaled version of (2) was proposed in prior work (

        https://ieeexplore.ieee.org/document/7347351…).

    • Duality

      • Given gradient G of a weights M, update M by UV', where USV' is the SVD of G. This generalizes g / ||g||
      • Muon similar but accumulates the gradient. Without this, they agree and agree with shampoo.

  • First-order optimization algorithms

    • Coordinate descent: optimize only along individual dims at a time

      Untitled

      • Used in variational inference
      • When conditionals are easier to optimize
      • No step size. Choose dims round robin or randomly.

    • GD

      • should strictly make loss go down every step (every step is a full batch epoch)

    • SGD

      image.png

      • vs GD is just full vs mini batch. GD is slow per step. SGD is just slow throughput.
      • The same final performance should be attainable using any batch size (See Shallue et al. 2018 and Why shouldn't the batch size be tuned to directly improve validation set performance?.

    • SGD with momentum: first moment (velocity)

      image.png

      • Nesterov momentum: update velocity first

      image.png

    • AdaGrad

      • Update each parameter differently, based on how much it has changed so far in training.
      • Intuition: if a component changed a lot (measured by the sum of its squared gradient), then it must have made a lot of progress toward the target, so slow down its effective learning rate with a larger denominator.

      image.png

      image.png

      • This lets you perform well if you’re in something less like a bowl and more like a valley, since you prioritize moving along the valley toward the goal instead of dropping straight to the bottom of the valley first and then moving toward the goal hi

    • RMSprop (Hinton)

      • Problem with AdaGrad: it only decreases the LR
      • EWMA of square of gradient lets us increase or decrease LR, lets it forget the past

      image.png

    • ADAM: first and second moments (Kingma, Ba)

      • Adam ~ RMSProp with momentum, but in the orig paper

        There are a few important differences between RMSProp with momentum and Adam: RMSProp with momentum generates its parameter updates using a momentum on the rescaled gradient, whereas Adam updates are directly estimated using a running average of first and second moment of the gradient.

        image.png

      • From paper

        image.png

      • Intuition

        • Momentum: updates w by v, not directly by g. Updates v by g.

          v += lr * g
w += v

        • RMSprop: updates w with LR that goes up/down by g^2.

          v = ewma of g^2
w += (lr = 1/sqrt(v)) * g

        • ADAM

          m = EWMA of g = like "v" of momentum
v = EWMA of g^2 = like "v" of RMSprop
w += (lr = 1/sqrt(v)) * m

      • $m$ is first moment (mean)

      • $v$ is second moment (uncentered variance)

      image.png

      • Code

        class Adam(SGD):
    def __init__(self, params, lr, wd=0., beta1=0.9, beta2=0.99, eps=1e-5):
        super().__init__(params, lr=lr, wd=wd)
        self.beta1,self.beta2,self.eps = beta1,beta2,eps

    def opt_step(self, p):
        if not hasattr(p, 'avg'): p.avg = torch.zeros_like(p.grad.data)
        if not hasattr(p, 'sqr_avg'): p.sqr_avg = torch.zeros_like(p.grad.data)
        p.avg = self.beta1*p.avg + (1-self.beta1)*p.grad
        unbias_avg = p.avg / (1 - (self.beta1**(self.i+1)))
        p.sqr_avg = self.beta2*p.sqr_avg + (1-self.beta2)*(p.grad**2)
        unbias_sqr_avg = p.sqr_avg / (1 - (self.beta2**(self.i+1)))
        p -= self.lr * unbias_avg / (unbias_sqr_avg + self.eps).sqrt()

      • Adam is "invariant to diagonal rescaling of the gradients”, i.e. Adam is invariant to multiplying the gradient by a diagonal matrix with only positive factors

        • Explainer

        • Another paper says:

          Second, while the magnitudes of Adam parameter updates are invariant to rescaling of the gradient, the effect of the updates on the same overall network function still varies with the magnitudes of parameters.

        • So presumably grad norm clipping has no effect? (Unless grad norm clipping is somehow achieving something on numerics)

      • With LR=0, we are still updating the moments (and grad norm clipping), just not updating weights.

    • Nice animation comparison of algorithms

  • Weight decay

    • This is the same as L2 norm: loss = ... + 0.5 * wd * weight**2 or $L_{new}\left(w\right) = L_{original}\left(w\right) + \lambda{w^{T}w}$
    • Or in gradients:
    new_weight = weight - lr*weight.grad - lr*wd*weight
# or
weight.grad += wd*weight

  • Grad norm clipping

    • Given grads g=[g1, g2, g3], the L2 norm is sum(g_i^2)^.5.
    • Clipping means: downscale g_i uniformly until L2 norm is 1 (or whatever is the clip).
      • (Not downscaling by their squares, or downscaling the large ones first, or something else)
    • This is done across all parameters, or can be done per-layer

    image.png

  • In practice

    • Learning rates: say, 1e-4 for larger networks

  • Learning rate schedules

    • Can “restart” or “continue” training

    https://twitter.com/__kolesnikov__/status/1687911223096971264

  • Hyperparameter optimization

    • bayesian optimization

  • TODO An Empirical Model of Large-Batch Training — Sam McCandlish

    • blog
    • explainer video
    • in-depth explainer video (haven’t watched)