Machine Learning

Gauss-Newton Optimization for Deep Neural Networks and LLM Training This document provides a brief technical overview of modern optimization techniques for large-scale deep learning. It starts with classical optimization foundations, moves through advanced second-order methods, and ends with practical curvature approximations that make these techniques feasible for billion-parameter models. What you’ll learn: How classical optimization methods scale (or don’t) to modern deep learning The mathematical foundation of second-order optimization Practical approximations like K-FAC, Shampoo, and Rank-1 methods When and why to use each optimization approach 📋 Table of Contents Classical Optimization Foundations Problem Setup Gradient Descent Newton’s Method Gauss–Newton and Generalized Gauss–Newton Setup: Network Outputs and Loss Hessian Structure Gauss–Newton (GN) Approximation Generalized Gauss–Newton (GGN) Why Use GGN? Practical Challenge: Computing J and G Approximating GGN in Practice Gradient Descent (First-Order Baseline) Adam K-FAC: Kronecker-Factored Approximate Curvature Shampoo Optimizer Rank-1 Curvature Approximation Summary of GGN Approximations Intuitive Analogy: Mountain Navigation Practical Guidance for Optimizer Selection Key Takeaways 1. Classical Optimization Foundations 1.1 Problem Setup The foundation of all optimization in deep learning is the empirical risk minimization problem. We consider the general problem of minimizing an objective function: $$ L(\theta) = \mathbb{E}_{(x,y)\sim \mathcal{D}}[\ell(f(x; \theta), y)] $$ ...

LSH Attention: Multiple Hyperplanes and Bucketing Great follow-up 👌 — let’s go carefully. When you use multiple random hyperplanes (say $L$ of them), you’re no longer just splitting into two buckets; you get a bit-string signature for each token. 1. Hashing with multiple hyperplanes Suppose you pick $L$ random hyperplanes: $r_1,\dots,r_L \in \mathbb{R}^d$. For each token $x$, compute: $$ h_{r_\ell}(x) = \text{sign}(r_\ell^\top x) \in {+1,-1}. $$ Collect them into a vector (bitstring): ...

Self-Attention Approximation Methods Approximate $S = QK^\top$ (the $n\times n$ score matrix) to reduce the $O(n^2 d)$ computational cost. 1) Kernelization / “Linear” Attention (e.g., Performer FAVOR+) Softmax attention can be seen as a kernel $k(q,k)=\exp\!\big(q^\top k/\sqrt{d}\big)$. If we find a feature map $\phi:\mathbb{R}^d\to\mathbb{R}^m$ so that $$ \exp\! \Big(\frac{q^\top k}{\sqrt{d}}\Big)\;\approx\; \phi(q)^\top \phi(k), $$ then $$ \underbrace{\mathrm{softmax}\!\big(QK^\top/\sqrt{d}\big)}_{\text{row-wise}} V \;\approx\; \frac{ \big(\phi(Q),[\phi(K)^\top V]\big) } { \big(\phi(Q),[\phi(K)^\top \mathbf{1}]\big) }, $$ where division is row-wise. We never form $n\times n$ matrices. ...