1
00:00:01,000 --> 00:00:42,192
[Hal Turing] Alrighty! Thanks for tuning in! Hello AI world! I am your host, Hal Turing, and my co-host is Dr. Ada Shannon.
Today we're digging into Shampoo: Preconditioned Stochastic Tensor Optimization. First author Vineet Gupta, with two co-authors, Tomer Koren and Yoram Singer, out of Google Brain and Princeton University, posted to arXiv in March 2018. Ada, here's the number that stopped me cold: a full preconditioner for an m by n weight matrix needs an mn by mn matrix. For a modest thousand-by-thousand layer, that's a trillion entries. So the question this paper asks is: can you get some of that curvature-awareness without paying anywhere close to that bill?

2
00:00:42,192 --> 00:01:05,272
[Dr. Ada Shannon] Here's the part that matters before we even open the algorithm box: this thing runs at close to SGD speed per step while using second-order-flavored information that should, on paper, cost orders of magnitude more. Normally in optimization you pick a lane — cheap and dumb, or expensive and smart. Shampoo is refusing that tradeoff, and the fact that it mostly holds up is worth sitting with.

3
00:01:05,272 --> 00:01:13,399
[Hal Turing] Before the mechanics, let's ground this. What does "preconditioning" even mean, for people who last touched this in a numerical methods class?

4
00:01:13,399 --> 00:01:53,013
[Dr. Ada Shannon] A preconditioner is a matrix you multiply the gradient by before you take your step, so a badly-shaped loss surface — a long narrow valley — gets rescaled into something closer to a round bowl. Plain SGD steps learning-rate times gradient, treating every direction the same. Newton's method is the gold standard: it uses the actual Hessian, the true curvature, as the preconditioner. Problem is, for a network with millions of parameters the Hessian is an N by N matrix — you can't store it, let alone invert it. That's why quasi-Newton methods exist: they approximate curvature cheaply from gradient history instead of computing it exactly.

5
00:01:53,013 --> 00:01:59,561
[Hal Turing] And that's where AdaGrad comes in. What's the diagonal-versus-full-matrix distinction this paper leans on so hard?

6
00:01:59,561 --> 00:02:39,035
[Dr. Ada Shannon] AdaGrad, from Duchi, Hazan, and Singer out of Berkeley and Google in 2011, comes in two flavors. The diagonal version — ancestor of Adam, what basically everyone uses — tracks one running number per parameter: the accumulated sum of squared gradients. Cheap, scales fine. The paper also proposed a full-matrix version tracking the entire outer-product covariance of accumulated gradients, capturing correlations diagonal methods miss. It's theoretically stronger and almost never used, because for n parameters that's an n by n matrix. Shampoo's pitch is recovering a usable chunk of that full-matrix version without the quadratic cost.

7
00:02:39,035 --> 00:02:44,422
[Hal Turing] So instead of flattening everything into one vector, Shampoo keeps the parameters in their natural shape?

8
00:02:44,422 --> 00:03:21,666
[Dr. Ada Shannon] Exactly, and that's the structural move. A fully-connected layer is an m by n matrix, not secretly a vector. A conv filter bank is a four-dimensional tensor. TensorFlow and Torch already represent weights this way. So instead of one impossible mn by mn preconditioner, Shampoo keeps one moderately-sized preconditioner per dimension — an m by m matrix, an n by n matrix — and combines them through a Kronecker product, a recipe for building one big block matrix out of two small ones without ever materializing it. Those two small matrices implicitly represent —

9
00:03:21,666 --> 00:03:28,447
[Hal Turing] Oh wait wait wait — hold on, implicitly represent, meaning you never actually form the giant matrix, you just act like you did?

10
00:03:28,447 --> 00:03:57,657
[Dr. Ada Shannon] Right. You apply the two small preconditioners to the gradient tensor directly, and it's mathematically equivalent to having used the Kronecker-product approximation of the full thing. There's a prior method built on similar instinct — K-FAC, Kronecker-Factored Approximate Curvature, from Martens and Grosse in 2015 — which approximates a network's Fisher information matrix the same way, per layer. We'll come back to how Shampoo and K-FAC relate later; for now, file it as the other major structure-aware preconditioner out there.

11
00:03:57,657 --> 00:04:06,063
[Hal Turing] And the convergence proofs use something called online convex optimization. I see that term constantly and only sort of understand it.

12
00:04:06,063 --> 00:04:35,877
[Dr. Ada Shannon] OCO analyzes an algorithm making a sequence of decisions and getting scored after each one, without assuming a fixed data distribution — the loss can even be adversarial. Success is measured as regret: how much worse your cumulative loss is versus the single best fixed decision in hindsight. It's the framework Zinkevich set up in 2003 for online gradient descent, and it's what AdaGrad's own regret bounds live in. Shampoo's convergence proof lives entirely in this stochastic convex world too.

13
00:04:35,877 --> 00:04:51,202
[Hal Turing] Okay but honestly, doesn't that bother you a little? Neural network landscapes are wildly non-convex — saddle points, weird geometry everywhere. If the proof only holds in the convex setting, how much are we supposed to trust it back in transformer-land?

14
00:04:51,202 --> 00:05:11,868
[Dr. Ada Shannon] I think that's the wrong bar. Nobody's claiming the regret bound literally holds during ResNet training. The convex analysis is a design tool — it tells you the algorithm has sound worst-case behavior in a controlled setting, a sanity check most heuristic optimizers never get. The real justification for using it on non-convex nets is empirical, not the theorem.

15
00:05:11,868 --> 00:05:21,063
[Hal Turing] Sure, but then why lean on the convex framing so heavily if it's not doing the load-bearing work? That feels like borrowing credibility from math that doesn't transfer.

16
00:05:21,063 --> 00:05:34,670
[Dr. Ada Shannon] Because it's still the only rigorous language the field has for vetting an update rule before throwing it at a billion-parameter model. I'm not saying trust the bound blindly — I'm saying it earns the algorithm a seat at the table, and the experiments have to close the deal.

17
00:05:34,670 --> 00:05:45,909
[Hal Turing] Fair, I'll hold the experiments to it. One last thing — the name. The footnote literally says it's called Shampoo because, quote, "it has to do with preconditioning." That's the whole joke.

18
00:05:45,909 --> 00:05:48,184
[Dr. Ada Shannon] That's it. That's the entire etymology.

19
00:05:48,184 --> 00:05:54,872
[Hal Turing] Okay, let's actually open the algorithm box. On every step, what are L_t and R_t doing, and how do they turn into the update?

20
00:05:54,872 --> 00:06:26,358
[Dr. Ada Shannon] You keep two small matrices instead of one giant one. L_t accumulates the sum of G G-transpose, R_t accumulates G-transpose G, capturing second-moment statistics along the rows and columns separately. Then W at t plus one equals W_t minus eta times L_t to the negative one-quarter, times G_t, times R_t to the negative one-quarter. You're sandwiching the gradient between a left preconditioner rescaling rows and a right one rescaling columns. No flattening anywhere in that expression.

21
00:06:26,358 --> 00:06:32,581
[Hal Turing] That negative one-quarter power is doing quiet work. Why a quarter, not a half like a normal square-root preconditioner?

22
00:06:32,581 --> 00:07:03,603
[Dr. Ada Shannon] Because you apply it from both sides — a quarter on the left composes with a quarter on the right into an effective half power, the standard decay rate for root-t convergence. And if you flatten W and G into vectors, that two-sided sandwich is mathematically identical to premultiplying by the Kronecker product of L_t and R_t, each to the quarter. That Kronecker product is provably close to the full, unapproximated AdaGrad preconditioner. The quarter exponent is the seam where the matrix trick and the flattened full-matrix picture agree.

23
00:07:03,603 --> 00:07:06,761
[Hal Turing] So what's the actual payoff versus ever forming that full thing?

24
00:07:06,761 --> 00:07:28,727
[Dr. Ada Shannon] Memory drops from m-squared n-squared down to m-squared plus n-squared — for a thousand-by-thousand layer, two million numbers instead of a trillion. Compute for inverting or rooting L_t and R_t is order m-cubed plus n-cubed, versus m-cubed n-cubed for the full preconditioner. Both scale linearly with the larger dimension instead of quadratically.

25
00:07:28,727 --> 00:07:35,182
[Hal Turing] That's for a plain matrix layer. Conv filters are four-dimensional tensors — how does this generalize?

26
00:07:35,182 --> 00:08:04,207
[Dr. Ada Shannon] Algorithm 2 handles arbitrary order-k tensors. You keep one preconditioning matrix per dimension, each accumulating the contraction of the gradient against itself along every other dimension. Instead of left- and right-multiplying, you apply each one to the gradient through a tensor-matrix contraction along its own mode, using the power negative one over two-k instead of negative one-quarter. Do that for every dimension — order doesn't matter, it's commutative — and the whole tensor's preconditioned without ever flattening it.

27
00:08:04,207 --> 00:08:14,981
[Hal Turing] Wait, hold on — some of those dimensions are huge, right? A twenty-five-hundred-wide embedding layer means inverting a twenty-five-hundred by twenty-five-hundred matrix. That can't be free.

28
00:08:14,981 --> 00:08:44,749
[Dr. Ada Shannon] It's not, so there's a diagonal Shampoo variant — you replace L_t and R_t with just their diagonals, storing m plus n numbers and updating in linear time. It triggers automatically past a size threshold, around twelve hundred in their experiments, decided independently per dimension. They also lean on two heuristics to keep wall-clock time down: recomputing the matrix roots only every twenty to a hundred steps instead of every step, and folding in ordinary momentum at alpha equals point-nine before preconditioning.

29
00:08:44,749 --> 00:08:46,978
[Hal Turing] And theoretically, what did they actually prove?

30
00:08:46,978 --> 00:09:05,182
[Dr. Ada Shannon] Theorems 7 and 10 bound the regret by root of two r D times the product of those trace terms, scaling as O of root T under standard conditions — assuming the gradient matrices have rank at most r, which sits explicitly in the constant. Theorem 10 is the tensor version of the same rate.

31
00:09:05,182 --> 00:09:14,703
[Hal Turing] Hold on — O of root T is exactly what vanilla online gradient descent gets with no preconditioning at all. What did all that machinery buy, rate-wise?

32
00:09:14,703 --> 00:09:38,294
[Dr. Ada Shannon] I'd push back on that framing, Hal. Nobody beats root T here, that's the known floor. What the preconditioning buys you is in the constant — Lemma 8 shows the Kronecker product of L and R lower-bounds the true full preconditioner, scaled by rank r, so the small eigenvalues that curvature actually cares about don't get washed out. Vanilla gradient descent has no such guarantee; its constant is just the raw gradient norm.

33
00:09:38,294 --> 00:09:45,167
[Hal Turing] A better constant is a smaller claim than the abstract implies, though. I'm not fully sold the theory is where the real story is.

34
00:09:45,167 --> 00:09:54,966
[Dr. Ada Shannon] Fair — I don't think it's meant to carry the paper alone, it's a sanity check the construction isn't nonsense. The experiments are where you see whether it actually trains faster.

35
00:09:54,966 --> 00:09:56,313
[Hal Turing] So what's the test bed?

36
00:09:56,313 --> 00:10:18,186
[Dr. Ada Shannon] CIFAR-10 on a 32-layer ResNet and a small 20-layer Inception network, CIFAR-100 on a 55-layer ResNet that deliberately skips batchnorm, and LM1B language modeling with the Attention model straight out of Vaswani et al. Everything on a single Tesla K40, batch size 128, just swapping the optimizer into existing training code.

37
00:10:18,186 --> 00:10:21,112
[Hal Turing] And Table 1 — does it actually keep pace per step?

38
00:10:21,112 --> 00:10:36,623
[Dr. Ada Shannon] Close to it. Steps per second on ResNet-32, Inception, and LM1B sit within a small margin of SGD, AdaGrad, and Adam, and on the 55-layer ResNet it's actually faster than all three — the one result they flag as surprising.

39
00:10:36,623 --> 00:10:38,759
[Hal Turing] What about the training curves themselves?

40
00:10:38,759 --> 00:11:01,700
[Dr. Ada Shannon] Figures 2 through 4. On both CIFAR-10 architectures and the CIFAR-100 ResNet, Shampoo reaches lower training loss within the same epoch budget — a visible margin. LM1B is the standout: test log-perplexity separates from the pack early and stays lower through all five hundred thousand steps, a noticeably bigger gap than either image run produced.

41
00:11:01,700 --> 00:11:34,301
[Hal Turing] That gap on LM1B is the number I'd lead with. But it's also where I want to push, because there's a comparison missing from every figure in this paper. Related Work spends real space on K-FAC — Martens and Grosse, University of Toronto, 2015 — the closest existing method that's also structure-aware and Kronecker-based. Then it just disappears. Not in Table 1, not in Figures 2 through 4. Every win we've cited is against SGD, AdaGrad, and Adam — not the one method built on the same idea.

42
00:11:34,301 --> 00:12:10,617
[Dr. Ada Shannon] It's a real hole. Citing K-FAC at length in the intro doesn't excuse skipping it in the experiments. The implicit defense is architectural — K-FAC needs to sample from the model's predictive distribution and leans on feed-forward backprop structure, while Shampoo only needs tensor shapes and works in a general stochastic setting. Fair simplicity argument. But 'easier to implement' and 'beats the best structure-aware alternative' are different claims, and this paper only earns you the first. Anyone citing Shampoo as proof that per-tensor Kronecker preconditioning wins outright is reading more into Table 1 than it contains.

43
00:12:10,617 --> 00:12:45,493
[Hal Turing] And it gets worse once you look at what's tested. Every experiment tops out at 13.5 million parameters, on one 2018 Tesla K40. Preconditioner cost scales as m-cubed plus n-cubed — fine at 256 by 256. Less fine once dimensions hit the thousands, which is just an ordinary transformer feedforward width now. The paper even has its own built-in admission of this — the diagonal fallback past that ~1200 threshold, quietly turning Shampoo into plain diagonal AdaGrad above it.

44
00:12:45,493 --> 00:13:11,732
[Dr. Ada Shannon] Wait, sorry to cut in — that threshold isn't a footnote, it's Section 5 in black and white. Above roughly 1200, you don't get Kronecker preconditioning anymore, you get diagonal AdaGrad with extra bookkeeping. Modern transformer FFN widths run 4096, 8192, sometimes over 16,000. So the honest read isn't 'might not scale' — the paper's own trigger condition tells you the structured advantage disappears for most tensors in anything built years after this.

45
00:13:11,732 --> 00:13:29,472
[Hal Turing] I actually disagree with how far you're pushing that. That threshold is an engineering default, not a theoretical concession — nothing stops someone raising it. Judging a 2018 method by a constant nobody tuned for 16,000-wide layers feels like moving the goalposts.

46
00:13:29,472 --> 00:13:58,218
[Dr. Ada Shannon] No — the paper frames itself as motivated by 'modern machine learning practices' generally, arbitrary tensor order, general-purpose. That's their scope claim, not mine. Claim general-purpose, and the burden's on you to show the general case, or at least flag where the approximation kicks in — there's no ablation anywhere on how much non-diagonal preconditioning buys as dimensions climb. Honestly, we're both a little right: real engineering knob, real gap in the evidence for their own stated scope.

47
00:13:58,218 --> 00:14:36,020
[Hal Turing] Fair, call it a draw. Here's the one that bothers me more — it's baked into the theory, not the implementation. That rank-r assumption behind Theorems 7 and 10 traces to the authors' own prior paper — Gupta, Koren, and Singer again, A Unified Approach to Adaptive Regularization in Online and Stochastic Optimization, 2017 — supplying the key lemma the proof rests on. A 55-layer ResNet or an attention stack is about as non-convex as it gets. What's the actual status of Figures 2 through 4 relative to that theorem?

48
00:14:36,020 --> 00:15:30,355
[Dr. Ada Shannon] Suggestive, not proven. The bound says: if the problem were convex and gradients had bounded rank, Shampoo tracks the best fixed point at the optimal rate. That doesn't transfer to 'therefore CIFAR-100 training loss drops faster' — the curves are separate empirical evidence pointing the same direction, with nothing formally connecting them. Worth remembering too: that CIFAR-10 Inception network comes straight out of Zhang, Bengio, Hardt, and Recht's Understanding Deep Learning Requires Rethinking Generalization, 2017 — a paper whose whole point is that lower training loss doesn't guarantee better generalization, and Figures 2 through 4 only show the training side. There's also an asymmetry on LM1B: Shampoo ran at the default learning rate, eta equals one, while every other algorithm — including Adam, Kingma and Ba's 2014 method — got swept over a range.

49
00:15:30,355 --> 00:16:06,532
[Hal Turing] Which is a shame, because underneath the tuning asymmetry there's a genuinely useful idea. The architecture-agnostic part is real — Shampoo only needs tensor shapes, not the model around it — and that block-diagonal, per-tensor independence is what let later work scale this up. What it never measures is the cost of that independence: it throws away correlation between, say, a QKV projection and the residual stream feeding it — coupling that's everywhere in the Attention model from Vaswani and colleagues at Google Brain, 2017, the very architecture tested here.

50
00:16:06,532 --> 00:16:35,185
[Dr. Ada Shannon] So the open list is short: an actual head-to-head against K-FAC, not a paragraph in Related Work; evidence at the dimensions transformers actually use instead of a 2018 K40 and 14 million parameters; and a non-convex extension of the theory that doesn't lean on an assumption real gradients may not satisfy. None of that makes this a bad paper — it's a clever engineering idea whose empirical case is narrower than the introduction implies.

51
00:16:35,185 --> 00:17:03,792
[Hal Turing] So, back to where we started: can a per-dimension preconditioner buy curvature-awareness without the full mn-by-mn bill? Inside the regime this paper actually tests — small models, one GPU, convex-flavored theory — yes, convincingly. Whether that holds at the scale and architectures we train today is the question this paper hands off, not one it answers. Thanks for sticking with us through all three parts — we'll catch you next time.