1
00:00:01,000 --> 00:00:55,195
[Hal Turing] Alrighty! Thanks for tuning in! Hello AI world! I am your host, Hal Turing, and my co-host is Dr. Ada Shannon.
And today we're digging into a paper called "Scalable Second Order Optimization for Deep Learning," by Rohan Anil et al. — that's five authors total: Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer, out of Google Research, Google Inc, Tel Aviv University, and Princeton University. It's dated March 8th, 2021 right on the paper itself. And Ada, the number that actually stopped me mid-scroll reading this one — they trained a Transformer on the WMT'14 English-to-French translation task in half as many steps as a well-tuned Adam baseline. Half. Not a few percentage points shaved off. Half.

2
00:00:55,195 --> 00:01:30,629
[Dr. Ada Shannon] Half the steps is a real number, not a rounding-error flex. And it's not one cherry-picked task either — they've got results on BERT, on a click-through-rate model at Criteo scale, on ResNet-50. What makes this worth our listeners' time is that second-order optimization has been the field's permanent "yeah but it doesn't scale" idea for literally decades. Everyone agrees the underlying math is more powerful than what we're actually running in production. Almost nobody uses it at scale. This paper is Google saying: we got it to actually work at a scale you'd care about, not on some toy MNIST classifier tucked into an appendix.

3
00:01:30,629 --> 00:01:48,647
[Hal Turing] Okay so before we get into what they actually built, let's back up, because I bet a chunk of our audience has trained plenty of models with Adam and never really stopped to unpack what "first-order" even means as opposed to "second-order" optimization. Ada, break it down for me — what's the actual mechanical difference between the two, in plain terms?

4
00:01:48,647 --> 00:02:41,542
[Dr. Ada Shannon] Sure. First-order just means the update only looks at the gradient — the slope of the loss with respect to each parameter. SGD does that directly; Adam dresses it up with a running per-parameter second-moment estimate, but it's still only ever tracking how each individual weight's own gradient has behaved over time, never how weights relate to each other. Second-order methods bring in curvature — not just which way is downhill, but how that downhill direction bends. Classically that's the Hessian, the full matrix of pairwise second derivatives. Newton's method uses the exact inverse Hessian and can converge in dramatically fewer steps than gradient descent. The catch is cost: for N parameters, the Hessian is N-by-N. That's O(N squared) just to store it, O(N cubed) to invert it. Totally fine for a few thousand parameters. Absurd for a billion-parameter model.

5
00:02:41,542 --> 00:02:59,979
[Hal Turing] Wait, hold on — where does "preconditioning" actually fit into all this? I hear that word thrown around constantly whenever people talk about optimizers, and I don't think I've ever gotten a straight, plain-English definition of what it actually means, or why anyone bothers with it instead of just using the raw gradient.

6
00:02:59,979 --> 00:04:00,165
[Dr. Ada Shannon] Preconditioning is just: before you apply the gradient step, you transform the gradient by multiplying it with some matrix — the preconditioner — that reshapes the update direction. A diagonal preconditioner, which is effectively what Adam uses, can only stretch or shrink each coordinate independently. A full-matrix preconditioner can also rotate — it can say "these two parameters tend to move together, adjust accordingly," which a diagonal method is structurally blind to. Full-matrix AdaGrad, the theoretically ideal version of this, is exactly as expensive as the Hessian problem I just described. That's where Shampoo comes in — that's Gupta, Koren, and Singer, 2018, and Koren and Singer are actually two of the five authors on today's paper too. Shampoo's trick is Kronecker factorization: instead of one giant preconditioner over all parameters jointly, it keeps one small matrix per dimension of each weight tensor — call them L and R for a 2D weight matrix — and approximates the full preconditioner as their Kronecker product.

7
00:04:00,165 --> 00:04:15,537
[Hal Turing] Oh wait wait wait — hold on, so instead of one gigantic N-by-N monster matrix, you end up with two much smaller matrices, like one sized for the rows and one for the columns, and you multiply those together to approximate the whole thing? That's a genuinely clever shortcut.

8
00:04:15,537 --> 00:04:51,713
[Dr. Ada Shannon] Exactly, and that's the whole reason it's tractable — the cost drops from cubic in the total parameter count to roughly cubic in the dimension sizes of each individual tensor. And Shampoo isn't the only one playing this game. K-FAC, from Martens and Grosse in 2015, does something structurally similar but approximates the Fisher information matrix instead of raw gradient outer-product statistics. K-BFGS, Goldfarb and coauthors in 2020, applies a comparable factored approximation directly to the Hessian. Those two are the natural rivals this paper keeps circling back to for comparison.

9
00:04:51,713 --> 00:05:37,967
[Hal Turing] So that's the lineage — Shampoo generalizing full-matrix AdaGrad, with K-FAC and K-BFGS as the other factored contenders chasing the same problem. And the headline claim in the abstract here is that once you actually get Shampoo running at scale, with the systems work this team did, you get real convergence gains and real wall-clock gains — not just theoretical ones — across Transformer machine translation, BERT, DLRM on Criteo, and ResNet-50. Four genuinely different production-scale workloads, not one lucky benchmark. Next up, Ada's going to walk us through what it actually took under the hood to make that happen — because, spoiler, the naive version of this idea falls over almost immediately.

10
00:05:37,967 --> 00:06:14,237
[Dr. Ada Shannon] Right, real convergence gains that actually show up on the wall clock, not just in a loss curve. But let's be honest about where the cost goes, because Shampoo isn't free just because you've factored the matrix. There are three extra chunks of work per step compared to Adam or AdaGrad. First, updating the preconditioner statistics themselves, accumulating those gradient outer products into L and R. Second, computing the inverse fourth root of L and R so you can actually apply them. Third, doing the preconditioned gradient multiply. The first and third are annoying but manageable. The second one is the killer.

11
00:06:14,237 --> 00:06:17,116
[Hal Turing] Wait, how much of a killer are we talking? Give me a number.

12
00:06:17,116 --> 00:06:49,810
[Dr. Ada Shannon] Up to 100 times the cost of a normal training step, just for that one root computation. And it's not optional — you can't approximate your way out of it with single precision either, because the L and R matrices in real workloads turn out to be badly ill-conditioned. Run the coupled Newton iteration or an SVD in single precision on those and it either diverges or gives you garbage. So they're stuck doing this in double precision, which is exactly the arithmetic mode modern accelerators are worst at, because TPUs and GPUs are built and priced around low-precision throughput.

13
00:06:49,810 --> 00:07:00,770
[Hal Turing] So you've got this thing that's simultaneously mandatory and 100x more expensive and runs badly on the exact hardware you're using. That sounds like a paper-killer, not a paper.

14
00:07:00,770 --> 00:07:37,411
[Dr. Ada Shannon] It would be, except they made a structural observation: the loss landscape doesn't shift that dramatically step to step, so you don't need a fresh inverse root every single step. They show empirically you can recompute L to the minus a quarter and R to the minus a quarter only once every few hundred steps with no meaningful accuracy hit. And once it's infrequent, you can shove it somewhere else entirely — onto the CPUs sitting next to the TPU or GPU that are otherwise mostly idle during training, computing asynchronously while the accelerator just keeps training on the last preconditioner it has.

15
00:07:37,411 --> 00:07:46,374
[Hal Turing] Oh wait, hold on — so the CPU is literally just crunching double-precision matrix roots in the background the whole time, and the accelerator never has to wait on it?

16
00:07:46,374 --> 00:08:31,653
[Dr. Ada Shannon] Exactly, that's the whole trick. Pipeline it. The accelerator computes the cheap statistics every step and applies whatever the current preconditioner is; the CPU cluster is off computing the next one in parallel, and by the time it's ready, it just gets swapped in. The most expensive operation in the whole algorithm ends up adding almost nothing to wall-clock time. That's paired with two more extensions for layers where even the factored version is too big — for a giant embedding or softmax layer, they only use one of L or R, not both, and that's not a hack, it's justified formally as Lemma 1 in the paper. For oversized fully-connected layers, they block the tensor into smaller chunks and precondition each block separately, backed by Lemma 2.

17
00:08:31,653 --> 00:08:37,272
[Hal Turing] Okay, so that's the engineering story. What did all of this actually buy them when they ran it on real models?

18
00:08:37,272 --> 00:09:24,083
[Dr. Ada Shannon] Solid gains, though not uniformly. On WMT'14 English-to-French with a Transformer, they hit the same perplexity in about half the steps compared to well-tuned Adam, and that translated into a 45% wall-time reduction on the base model, 37% on Transformer-Big. DLRM on the Criteo click-through dataset also converged in about half the steps, 37.5% wall-time reduction, and it set a new state-of-the-art AUC of 80.56%, which in that task is a meaningfully large jump. ResNet-50 on ImageNet needed 31.7% fewer steps than the prior state-of-the-art and came in 13% faster in wall time, though that one leaned on emulated lower precision to make the accelerator math tractable.

19
00:09:24,083 --> 00:09:29,981
[Hal Turing] And BERT-Large? You said 16% fewer steps earlier, but that doesn't sound like it's keeping pace with the others.

20
00:09:29,981 --> 00:09:55,616
[Dr. Ada Shannon] It isn't — only a 4% wall-time reduction, and the paper is upfront that BERT wasn't yet tuned for performance in their system. So the step-count win is real, but it didn't convert into wall-clock savings the way the other tasks did. That's a pretty telling asterisk to plant early: whatever's driving the wall-clock number isn't just 'run Shampoo,' it's also how much systems effort went into that particular task.

21
00:09:55,616 --> 00:10:04,765
[Hal Turing] Last thing before we move on — did they actually put Shampoo up against K-FAC and K-BFGS directly, or just against Adam and friends the whole way through?

22
00:10:04,765 --> 00:10:35,694
[Dr. Ada Shannon] Only at toy scale. MNIST, FACES, and CURVES autoencoders — small models, tens of thousands of parameters. On those, Shampoo, K-FAC, and K-BFGS all land in roughly the same place, and all clearly beat RMSprop and Adam. But scaling K-FAC or K-BFGS up to Transformer or BERT-Large scale is explicitly punted to future work in this paper. Every one of the big production-scale numbers we just walked through is Shampoo versus first-order baselines only.

23
00:10:35,694 --> 00:11:16,236
[Hal Turing] Right, toy scale only, and that's exactly the thing I want to press on, Ada. The abstract's headline line is that this is 'comparable to second-order methods such as K-FAC and K-BFGS.' But you just told me that comparison only happened on MNIST, FACES, and CURVES autoencoders — tens of thousands of parameters. At Transformer, BERT-Large, DLRM, ResNet-50 scale, every single baseline is first-order: Adam, AdaGrad, LAMB, Nesterov, LARS. So does that 'comparable to SOTA second-order methods' framing actually hold where it matters, or is it borrowed credibility from a toy result?

24
00:11:16,236 --> 00:11:58,450
[Dr. Ada Shannon] Honestly, borrowed credibility. The toy-scale result is real and it's consistent across three different autoencoder benchmarks, so it's not nothing — Shampoo genuinely tracks K-FAC and K-BFGS when the problem is small enough that all three are tractable. But the paper never tests whether that equivalence survives at the scale where the interesting claims actually live. Extrapolating a tens-of-thousands-of-parameters result to justify a headline about Transformer and BERT-Large-scale training is exactly the move a careful reader should flag. The honest version of the abstract would say Shampoo matches second-order methods at small scale and matches or beats first-order methods at large scale — those are two separate claims, and only one of them is tested where it counts.

25
00:11:58,450 --> 00:12:33,790
[Hal Turing] That's a fair place to land it. So the real headline here isn't 'Shampoo equals K-FAC at scale' — it's that a preconditioned method can go toe-to-toe with heavily tuned Adam and LAMB on production-sized workloads like WMT'14, DLRM, and ResNet-50, with the BERT-Large case as an honest asterisk showing the systems engineering still matters as much as the algorithm. Thanks for walking through all of that, Ada — great breakdown as always. And thank you all for listening to AI Post Transformers. Until next time.