1
00:00:01,000 --> 00:00:07,269
[Hal Turing] Alrighty! Thanks for tuning in! Hello AI world! I am your host, Hal Turing, and my co-host is Dr. Ada Shannon.

2
00:00:07,269 --> 00:00:32,161
[Hal Turing] Today's paper is Large Memory Layers with Product Keys, by Guillaume Lample and colleagues at Facebook AI Research and Sorbonne Université, from NeurIPS 2019. The claim is that you can bolt a memory with over a million trainable slots into a transformer, and a 12-layer model with it beats a 24-layer model. It also runs almost twice as fast.

3
00:00:32,161 --> 00:00:42,377
[Dr. Ada Shannon] Which sounds like free capacity. The question is whether it's free, or whether the comparison hides the bill. Over three parts we'll test that. First, the machinery.

4
00:00:42,377 --> 00:00:50,876
[Hal Turing] Start with the setting. These are autoregressive language models, so they predict the next word from the previous ones. The score is perplexity.

5
00:00:50,876 --> 00:01:00,210
[Dr. Ada Shannon] Lower is better. Roughly, it's how many equally likely next words the model is torn between. A perplexity of 20 means it's about as unsure as a fair twenty-sided die.

6
00:01:00,210 --> 00:01:02,997
[Hal Turing] And the motivation is that models underfit big corpora?

7
00:01:02,997 --> 00:01:27,238
[Dr. Ada Shannon] Yes. Capacity normally costs compute, because every parameter runs for every token. The authors also say the One Billion Word corpus is too small, since models overfit it. On a 16-layer, 1024-wide model, test perplexity reaches 25.3 just as validation perplexity starts rising. So they train on CC-News, about 28 billion words, where more capacity can still help.

8
00:01:27,238 --> 00:01:31,790
[Hal Turing] Where does the memory plug in? I want the transformer block pinned down first.

9
00:01:31,790 --> 00:01:41,031
[Dr. Ada Shannon] A block is self-attention, then a feed-forward network, the FFN. That's a two-layer network applied to each token separately. The memory layer replaces that FFN.

10
00:01:41,031 --> 00:01:42,796
[Hal Turing] Spell out the memory layer itself.

11
00:01:42,796 --> 00:02:03,322
[Dr. Ada Shannon] A query network turns the token's hidden state into a query vector. You score the query against a table of learned keys and keep the top k. A softmax weights those k, and the output is the weighted sum of the value vectors attached to those keys. The values hold most of the parameters. Only k slots are read, and only those get gradient. That's the sparse part.

12
00:02:03,322 --> 00:02:05,458
[Hal Turing] So why not do that with a million keys today?

13
00:02:05,458 --> 00:02:19,530
[Dr. Ada Shannon] Because finding the top k by scoring every key is linear in memory size. Score a million keys per token and you've spent the compute you were trying to save. That's why memory networks, Neural Turing Machines and neural caches all stayed small.

14
00:02:19,530 --> 00:02:21,620
[Hal Turing] Someone must have tried to dodge that.

15
00:02:21,620 --> 00:02:43,679
[Dr. Ada Shannon] Rae and colleagues, 2016, in Sparse Access Memory. They used an approximate nearest-neighbor index to find the top slots quickly. The Lample paper says that index has to be re-learned regularly as the keys change in training, or it drifts catastrophically. Their pitch is exact search instead. I'll hold judgment on that, because as far as the source shows, they don't test it head to head against Rae.

16
00:02:43,679 --> 00:02:46,511
[Hal Turing] And the other route to cheap capacity is mixture of experts.

17
00:02:46,511 --> 00:03:04,344
[Dr. Ada Shannon] Shazeer and colleagues, 2017, the sparsely-gated mixture-of-experts layer. A gate sends each input to a few big expert sub-networks. Capacity grows, per-token compute stays small. Keep it in mind, because it's the obvious competitor, and we'll come back to whether it shows up as a baseline.

18
00:03:04,344 --> 00:03:08,710
[Hal Turing] Okay, product keys. How do you search a million keys without scoring a million keys?

19
00:03:08,710 --> 00:03:28,865
[Dr. Ada Shannon] You never build the million keys. You keep two small codebooks of sub-keys, each half the query's dimension. Every full key is one sub-key from the first codebook concatenated with one from the second. Two codebooks of 1024 sub-keys give 1024 times 1024 keys, about a million. But you only store 2048 half-size vectors.

20
00:03:28,865 --> 00:03:30,165
[Hal Turing] Give me a toy version.

21
00:03:30,165 --> 00:03:53,756
[Dr. Ada Shannon] Three sub-keys per codebook, so nine implied keys. Split the query into a left half and a right half. The left half scores the three left sub-keys, the right half scores the three right ones. Keep the best two on each side. Pair them up and you get four candidate full keys. Each candidate's score is just the left score plus the right score. Pick your top k from those four.

22
00:03:53,756 --> 00:03:57,611
[Hal Turing] So at full scale the search touches two thousand-ish sub-keys, not a million.

23
00:03:57,611 --> 00:04:06,992
[Dr. Ada Shannon] Right, roughly a thousandfold fewer comparisons. The next part covers why the answer is still exact, and what exact means here, because it comes with a caveat.

24
00:04:06,992 --> 00:04:12,797
[Hal Turing] So next is how it's trained and wired in, then the benchmarks, and finally what the comparison leaves out.

25
00:04:12,797 --> 00:04:16,419
[Dr. Ada Shannon] Yes. Does the headline hold, and what exactly is it matched against?

26
00:04:16,419 --> 00:04:19,623
[Hal Turing] Let's trace one token through the memory layer, start to finish.

27
00:04:19,623 --> 00:04:49,670
[Dr. Ada Shannon] The hidden state goes into a query network, a linear layer or small MLP, that outputs a 512-dimensional query, followed by batch normalization. The query is split into two 256-dimensional halves. Each half is scored against its own codebook of about a thousand sub-keys, and each keeps its top k. That leaves k squared candidate pairs. The final top k comes from those pairs, a softmax weights them, and the output is the weighted sum of the value vectors in those k slots.

28
00:04:49,670 --> 00:04:52,503
[Hal Turing] And the savings come from never touching the full table.

29
00:04:52,503 --> 00:05:05,970
[Dr. Ada Shannon] Right. Cost goes from the number of slots times the query dimension to roughly the square root of the slots plus k squared, times the query dimension. At a thousand-by-thousand grid that is on the order of a thousand times fewer operations.

30
00:05:05,970 --> 00:05:12,147
[Hal Turing] Why can you trust that shortcut? Taking the top k of each half and then combining sounds like it could miss the true winner.

31
00:05:12,147 --> 00:05:42,519
[Dr. Ada Shannon] It can't, and the argument is short. A full key's score is the sum of its two half scores. Say the true best pair used a sub-key outside its half's top k. Then k better sub-keys in that half could each be paired with the same partner and beat it. So the true top k must sit inside the candidate set. But the exactness has a scope. It is exact over product-structured keys. Those keys are a Cartesian product, not a million freely chosen vectors, so the structure limits what the keys can express.

32
00:05:42,519 --> 00:05:46,559
[Hal Turing] Training is sparse too, right? Only the selected slots get gradient.

33
00:05:46,559 --> 00:06:19,578
[Dr. Ada Shannon] Only k slots per head. Because each value gets updated rarely, the authors train the values with Adam at a learning rate of 1e-3, against 2.5e-4 for the rest of the network. Keys and the query network use the base rate. They also use multiple heads: several independent query networks and sub-key sets that share one value table, with outputs summed. The main setup is four heads, k of 32, and 512 squared, so 262,144 slots. Different heads select largely different keys.

34
00:06:19,578 --> 00:06:24,082
[Hal Turing] And it replaces the FFN in chosen layers, as x plus memory of x.

35
00:06:24,082 --> 00:06:38,618
[Dr. Ada Shannon] Yes, in a handful of layers, spaced out. In a 16-layer model that means something like layers six and twelve. Position matters. In their 6-layer test, memory at layer four or five was best and layer one was worst.

36
00:06:38,618 --> 00:06:40,568
[Hal Turing] So the headline. What do we actually see?

37
00:06:40,568 --> 00:07:06,110
[Dr. Ada Shannon] At width 1024, a 12-layer model with one memory layer reaches perplexity 15.6. The 24-layer dense model gets 16.0, and the 12-layer dense model gets 17.7. The memory model also runs in about half the time per word. At width 1600 the margin narrows, 13.7 against 14.0. These are single runs on the news corpus, with no variance reported.

38
00:07:06,110 --> 00:07:10,522
[Hal Turing] Half the layers, better score. Does the paper say how many parameters that memory adds?

39
00:07:10,522 --> 00:07:21,111
[Dr. Ada Shannon] No. Parameter counts, memory footprint and inference hardware are all missing, and so is any table of training speed, even though the abstract mentions it. Hold that for the critique.

40
00:07:21,111 --> 00:07:26,126
[Hal Turing] Before we do, how do product keys compare to flat keys? Same memory, exhaustive search.

41
00:07:26,126 --> 00:07:55,337
[Dr. Ada Shannon] I need two terms first. Usage is the fraction of slots ever accessed on the test set. KL divergence to uniform measures how lopsided those accesses are. Low usage means paid-for capacity sits idle. At 147 thousand slots, flat keys touch about one slot in ten and run under 14 thousand words per second. Product keys touch about 99.6 percent of the slots and run above 36 thousand, with perplexity 20.7 against 22.1.

42
00:07:55,337 --> 00:07:58,077
[Hal Turing] So structure helped optimization, not just speed.

43
00:07:58,077 --> 00:08:07,458
[Dr. Ada Shannon] That's the suggestive part. But the flat model couldn't be trained much beyond that size, and product keys also have far fewer key parameters, so the comparison isn't matched.

44
00:08:07,458 --> 00:08:09,408
[Hal Turing] Does a bigger memory keep paying off?

45
00:08:09,408 --> 00:08:40,058
[Dr. Ada Shannon] Only if the slots are used. In the 6-layer model, going from 16 thousand to a million slots takes perplexity from 23.0 to 18.0, with speed essentially flat. And 23.0 is exactly the memoryless baseline, so the small memory bought nothing. Without batch normalization on the query, a million slots reach only about 26 percent usage. With it, about 80 percent. So scaling a sparse memory is partly an optimization problem.

46
00:08:40,058 --> 00:08:41,219
[Hal Turing] And the heads and k?

47
00:08:41,219 --> 00:08:50,925
[Dr. Ada Shannon] Raising heads times k improves perplexity and usage. Pairs with the same product behave alike. Speed depends on heads and k, not on memory size.

48
00:08:50,925 --> 00:08:58,913
[Hal Turing] Parameter counts are the gap I keep snagging on. If half the layers plus a memory beats the deep model, what did we pay in parameters?

49
00:08:58,913 --> 00:09:24,966
[Dr. Ada Shannon] Probably a lot, and the paper never tallies it. Take 262,144 slots times a value dimension. If that dimension is near 1024, which the text doesn't confirm, that's roughly 268 million added parameters. So the headline shows that a model that is fast per token, with a big lookup table, beats a deeper model per unit of latency. It does not show that sparse lookup beats simply having more parameters.

50
00:09:24,966 --> 00:09:26,638
[Hal Turing] What control would settle that?

51
00:09:26,638 --> 00:09:46,514
[Dr. Ada Shannon] Two. A 12-layer model with a wider FFN, or the same model with the memory swapped for a dense FFN holding the same parameter count. Neither is trained. There's a second confound too. The memory replaces an FFN in those layers, so the comparison mixes removing dense capacity with adding a memory.

52
00:09:46,514 --> 00:09:49,579
[Hal Turing] And mixture of experts? It's cited as the alternative.

53
00:09:49,579 --> 00:10:12,752
[Dr. Ada Shannon] Cited and never trained. So the paper can't separate any big sparse layer from product-key addressing specifically. The only addressing comparison is product versus flat keys. That one is confounded, since product keys have far fewer key parameters. Flat keys were also run only up to 147 thousand slots, because training took an unreasonable amount of time, so a slow flat baseline may be exaggerating the gap.

54
00:10:12,752 --> 00:10:16,607
[Hal Turing] What about the exactness pitch against Rae and colleagues' drifting index?

55
00:10:16,607 --> 00:10:33,232
[Dr. Ada Shannon] Argued, not tested. No approximate-index baseline is run. And exactness only covers the product-structured key set. Keys have to be a Cartesian product of two codebooks, so what the memory can express is restricted, and that cost is never measured.

56
00:10:33,232 --> 00:10:37,691
[Hal Turing] Then there's the speed claim. Does almost twice as fast survive outside their setup?

57
00:10:37,691 --> 00:10:58,264
[Dr. Ada Shannon] Unknown. It's inference throughput in words per second, and the inference hardware isn't given. Each token also reads a few hundred rows from a huge table, 128 in the main setup, and a table that big costs real RAM even though FLOPs stay low. The paper doesn't measure footprint or what happens if the table has to be sharded or offloaded.

58
00:10:58,264 --> 00:11:01,375
[Hal Turing] The query batch normalization also bothers me. It's doing a lot of work.

59
00:11:01,375 --> 00:11:25,524
[Dr. Ada Shannon] It is. It moves a million-slot memory from about a quarter of slots used to about four in five. But usage only says a slot was touched at least once. It doesn't say how much each slot contributes, and the paper reports no per-slot breakdown. And that evidence comes from a 6-layer ablation, far smaller than the headline models. With single runs and no variance, the small gaps can't be called reliable.

60
00:11:25,524 --> 00:11:28,171
[Hal Turing] Everything is perplexity on one news corpus too.

61
00:11:28,171 --> 00:11:46,793
[Dr. Ada Shannon] Yes. No downstream task, so we can't tell which tokens improved or whether the gain carries beyond next-word prediction. Also, the baseline is underfitting at this data scale, so a bigger memory may be filling an undertrained regime. The paper doesn't say whether the gain holds once the dense model is trained further.

62
00:11:46,793 --> 00:11:48,697
[Hal Turing] Has anyone run the check that was missing?

63
00:11:48,697 --> 00:12:20,137
[Dr. Ada Shannon] The best one is Memory Layers at Scale, by Berges and colleagues at Meta, 2024. They scale product-key-style layers to 128 billion memory parameters on 8-billion-parameter base models, with custom kernels, against dense and mixture-of-experts baselines at matched compute and parameters. From the description I have, the gains show up mainly in factual question answering. So the follow-up tests exactly what 2019 left open, and the knowledge-lookup framing looks right. I won't quote numbers I haven't verified.

64
00:12:20,137 --> 00:12:21,484
[Hal Turing] Practically, who should care?

65
00:12:21,484 --> 00:12:36,623
[Dr. Ada Shannon] People who need cheap per-token capacity for knowledge-heavy lookup, and who can afford the footprint. The costs are real. The table lives in memory, you need custom kernels for the gather, and utilization is an optimization problem, as that 26 percent usage showed.

66
00:12:36,623 --> 00:12:38,667
[Hal Turing] What experiments would close this out?

67
00:12:38,667 --> 00:12:50,880
[Dr. Ada Shannon] A parameter-matched dense control. A mixture of experts at matched parameters and FLOPs. An approximate-index baseline in Rae's style. A downstream evaluation. And multiple seeds.

68
00:12:50,880 --> 00:13:06,577
[Hal Turing] So in this setup, product keys make a huge sparse memory trainable and fast, and it beats extra depth at equal width. What the paper doesn't show is that lookup beats more parameters, or that it beats other sparse layers. Thanks for listening, everyone.

69
00:13:06,577 --> 00:13:07,970
[Dr. Ada Shannon] Thanks, all. Goodbye.