1
00:00:01,000 --> 00:00:09,300
[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:09,300 --> 00:00:30,675
[Hal Turing] Today we're looking at "Memory Layers at Scale," from Vincent-Pierre Berges, Barlas Oğuz, and colleagues at Meta FAIR, posted to arXiv in December 2024. The question is whether you can add a huge number of parameters to a language model, especially parameters that store facts, without adding any arithmetic per token.

3
00:00:30,675 --> 00:00:43,850
[Dr. Ada Shannon] The claim is that you can, and that at the same compute it beats both dense models and mixture-of-experts. That claim has an asterisk, and the asterisk is the whole episode: what does "same compute" leave out?

4
00:00:43,850 --> 00:00:50,925
[Hal Turing] Start with why facts are the target. A birthday or a capital city is a lookup. It's not reasoning.

5
00:00:50,925 --> 00:01:10,475
[Dr. Ada Shannon] Right. In a standard transformer, that lookup lives in the feed-forward network, the FFN. That's the dense block in every layer that processes each token. Every weight in it touches every token. So if you want to store more facts, you widen the FFN, and the arithmetic per token grows with it. Storage and compute are bolted together.

6
00:01:10,475 --> 00:01:13,700
[Hal Turing] And that's what the scaling laws describe.

7
00:01:13,700 --> 00:01:32,150
[Dr. Ada Shannon] Yes. Kaplan and colleagues at OpenAI, 2020, showed that loss falls smoothly as you grow parameters, data and compute together. But that's loss averaged over everything. How a specific capability like factual recall scales is much less studied, and this paper goes after exactly that.

8
00:01:32,150 --> 00:01:38,500
[Hal Turing] So how do you decouple parameters from compute? I know mixture-of-experts is one answer.

9
00:01:38,500 --> 00:02:00,974
[Dr. Ada Shannon] It's the established one. A mixture-of-experts layer has several FFN experts and a router that sends each token to only a few. You hold many more parameters than you use per token. Fedus and colleagues at Google popularized it with the Switch Transformer in 2021. One variant matters later: expert-choice routing, from Zhou and colleagues at Google in 2022. There, each expert picks its favorite tokens, which balances the load during training.

10
00:02:00,974 --> 00:02:04,724
[Hal Turing] Fine, so what's the memory layer doing differently?

11
00:02:04,724 --> 00:02:27,275
[Dr. Ada Shannon] Think of attention. A query is compared against keys, and you take a weighted sum of values. In attention, those keys and values are activations from the context. In a memory layer, the keys and values are learned weights, just a big trainable table. A token produces a query, finds the top-k most similar keys, and returns a softmax-weighted sum of those k value vectors. Only those k rows are read, and only those get gradient updates.

12
00:02:27,275 --> 00:02:29,700
[Hal Turing] Give me a concrete picture.

13
00:02:29,700 --> 00:02:49,475
[Dr. Ada Shannon] Say the table has about a million entries. The token's hidden state at "the capital of France is" becomes a query. It scores against the keys, and a handful of rows come back, with k a small number. Those rows are blended and added to the residual stream. The other roughly million minus k entries cost nothing for this token.

14
00:02:49,475 --> 00:02:56,275
[Hal Turing] But scoring a query against a million keys sounds like exactly the arithmetic you're avoiding.

15
00:02:56,275 --> 00:03:25,025
[Dr. Ada Shannon] It would be, done naively. That's where product keys come in, from Lample and colleagues at Facebook AI Research, 2019, in "Large Memory Layers with Product Keys." You split the query in half. Each half is scored against a small set of sub-keys, around the square root of N. Combine the best halves and you've effectively searched all N keys without building the full table of keys. With 1,024 sub-keys per half you get 1,024 squared, so 2^20 values, about a million.

16
00:03:25,025 --> 00:03:28,300
[Hal Turing] So this isn't new. 2019.

17
00:03:28,300 --> 00:03:48,825
[Dr. Ada Shannon] Correct, and it's important to say so. Lample's result was that a 12-layer model with one memory layer beat a 24-layer model, at roughly 100 million parameter scale. This paper is a scale-up. The contribution is making it work at billions of parameters with custom GPU kernels, plus some design changes we'll get to. It's not a new idea.

18
00:03:48,825 --> 00:03:52,525
[Hal Turing] There's also a closer relative, I think. PEER?

19
00:03:52,525 --> 00:04:10,950
[Dr. Ada Shannon] "Mixture of a Million Experts," Xu Owen He, 2024, out of Google DeepMind. It also uses product keys, but each of its experts is a single neuron, built from two retrieved embeddings. It's the nearest neighbor to a memory layer, and it needs twice the parameters per key. It shows up as a baseline.

20
00:04:10,950 --> 00:04:16,350
[Hal Turing] Now the "matched FLOPs" part. FLOPs are just the arithmetic per token, right?

21
00:04:16,350 --> 00:04:38,600
[Dr. Ada Shannon] Yes, and it's only one cost. A model can store far more parameters than it touches per token. Then there's memory traffic: dense matrix multiplies are limited by arithmetic, so they're FLOP-bound. A memory lookup does almost no arithmetic. It fetches scattered rows from GPU memory, so it's memory-bandwidth-bound. FLOP counts say nothing about its speed. And once the table is sharded across GPUs, you also pay communication.

22
00:04:38,600 --> 00:04:43,175
[Hal Turing] So three costs: compute, stored parameters, and traffic.

23
00:04:43,175 --> 00:04:45,025
[Dr. Ada Shannon] Right. Keep those separate.

24
00:04:45,025 --> 00:04:46,851
[Hal Turing] And the yardsticks?

25
00:04:46,851 --> 00:05:09,151
[Dr. Ada Shannon] Closed-book question answering: NaturalQuestions and TriviaQA. The model answers real questions from its weights alone, with no retrieval. NaturalQuestions is scored by exact match, TriviaQA by F1 overlap with the gold answer. Higher is better on both. For ablations, they use negative log-likelihood, the loss on the correct answer, where lower is better. Accuracy on small models is too noisy to compare.

26
00:05:09,151 --> 00:05:16,351
[Hal Turing] So the tension going in is strong factual gains at matched compute, versus what "matched" hides.

27
00:05:16,351 --> 00:05:20,176
[Dr. Ada Shannon] That's it. And whether the competitors were given a fair shot.

28
00:05:20,176 --> 00:05:32,576
[Hal Turing] Let's get into what they actually built. The vanilla memory layer is a single lookup that replaces one FFN. The paper's improved version is called Memory+. What's different?

29
00:05:32,576 --> 00:05:54,601
[Dr. Ada Shannon] Three changes, plus one stability fix. First, three memory layers instead of one, and all three share a single pool of values. Second, the lookup output goes through an input-dependent gate with a swilu non-linearity, then an output projection. Third, the key dimension is set to half the model dimension. The stability fix is qk-norm, which they use when training blows up on large memories, especially in small base models.

30
00:05:54,601 --> 00:06:01,901
[Hal Turing] Why would sharing help? Three layers reading one table sounds like it should just be one layer with extra steps.

31
00:06:01,901 --> 00:06:20,626
[Dr. Ada Shannon] The plausible reason is that different depths work with different levels of abstraction. A middle layer and a later layer ask the table different questions, so the same stored entries get used in more ways. And sharing adds no extra memory parameters, so you get the extra lookups for free in parameter terms.

32
00:06:20,626 --> 00:06:24,026
[Hal Turing] So why not replace every FFN?

33
00:06:24,026 --> 00:06:44,626
[Dr. Ada Shannon] Because each swap removes dense parameters. The ablation shows it. With three shared layers, the answer loss on the training-set questions improves. At six, it gets worse again. Their reading is that sparse and dense layers are complementary. One caveat: these ablations measure negative log-likelihood on the training set, and the gaps are small. Suggestive, not decisive.

34
00:06:44,626 --> 00:06:51,526
[Hal Turing] Now the systems side. They wrote custom CUDA kernels, and that's where the claim of a speedup lives.

35
00:06:51,526 --> 00:07:24,951
[Dr. Ada Shannon] The lookup is an EmbeddingBag: gather k rows, take a weighted sum. Almost no arithmetic, so speed is set by how fast you can fetch scattered rows from GPU memory. PyTorch's version gets under 400 gigabytes per second on an H100. Their forward kernel reaches about 3 terabytes per second, against a 3.35 spec. The backward pass is harder, because many tokens write gradients into the same row. They tried atomics, row locks, and a reverse-index approach that avoids atomics. End to end, the custom kernels are about 6 times faster than PyTorch's EmbeddingBag.

36
00:07:24,951 --> 00:07:28,051
[Hal Turing] Faster than the dense FFN?

37
00:07:28,051 --> 00:07:38,026
[Dr. Ada Shannon] No. That 6x is against PyTorch's own operator. It says nothing about the memory layer versus a dense FFN. Keep that straight.

38
00:07:38,026 --> 00:07:41,776
[Hal Turing] And the table is too big for one GPU.

39
00:07:41,776 --> 00:08:02,576
[Dr. Ada Shannon] Right. The values are sharded across GPUs along the embedding dimension. Indices are shared across the group, each GPU sums its slice of the embeddings, and the partial results are sent back. It runs in its own process group, independent of tensor, pipeline or context parallelism. That communication is a real cost, and we'll come back to it.

40
00:08:02,576 --> 00:08:06,676
[Hal Turing] Okay, the setup. Who is Memory+ up against?

41
00:08:06,676 --> 00:08:29,351
[Dr. Ada Shannon] Llama-style dense models at 134 million, 373 million, 720 million and 1.3 billion parameters, each trained on 1 trillion tokens. MoE baselines with 16, 8, 6 and 4 experts, picked as the smallest count that exceeds the memory model's parameters. PEER with 768 half-keys. And memory with 1024 half-keys, so 2 to the 20, about a million values.

42
00:08:29,351 --> 00:08:33,551
[Hal Turing] Give me the headline at the 1.3 billion base.

43
00:08:33,551 --> 00:08:53,251
[Dr. Ada Shannon] On NaturalQuestions, the dense model answers about 8 per hundred correctly. Memory+ answers about 14, roughly 1.76 times as many. The 4-expert MoE, with a similar parameter count, also gets about 8. So extra parameters alone didn't help. The architecture had to use them well.

44
00:08:53,251 --> 00:08:57,251
[Hal Turing] But PEER has even more parameters. Where does it land?

45
00:08:57,251 --> 00:09:15,576
[Dr. Ada Shannon] Close. On TriviaQA, PEER is at 42.5 and Memory+ at 42.9. On NQ at this size, PEER beats vanilla Memory, 12.3 to 9.8. So Memory+ over PEER is a small edge, not a blowout. And these are single runs with no error bars.

46
00:09:15,576 --> 00:09:19,726
[Hal Turing] What about the claim that it beats larger dense models?

47
00:09:19,726 --> 00:09:42,951
[Dr. Ada Shannon] The 720 million Memory+ against the 1.3 billion dense model. On NQ, 9.4 versus 7.8. On TriviaQA, 36.7 versus 32.6. So yes on factual QA. On PIQA, a physical commonsense test, it's 72.4 versus 72.7, and on OBQA, a science question set, 24.0 versus 23.4. That's mixed, nothing to hang a claim on.

48
00:09:42,951 --> 00:09:46,851
[Hal Turing] Now the one everyone quotes. Scaling the table.

49
00:09:46,851 --> 00:10:08,576
[Dr. Ada Shannon] They hold the 1.3 billion base fixed and grow the memory. Factual QA rises steadily. At 64 million values, about 139 billion total parameters, TriviaQA reaches 62.1 against 64.0 for Llama2 7B. That model used roughly ten times the compute per token and twice the training tokens, so it's a different data mix and budget, not a controlled comparison.

50
00:10:08,576 --> 00:10:11,976
[Hal Turing] So it matches Llama2 7B?

51
00:10:11,976 --> 00:10:33,051
[Dr. Ada Shannon] On TriviaQA, nearly. On NaturalQuestions it's still about 4 points behind, 20.8 against 25.1. And returns diminish at the top. Going from 16 million to 64 million values quadruples the table for about 0.6 more NQ points. Also, the paper shows plots, not fitted scaling exponents. So "scaling law" is generous.

52
00:10:33,051 --> 00:10:38,801
[Hal Turing] And that 139 billion parameter table has to live somewhere.

53
00:10:38,801 --> 00:10:41,201
[Dr. Ada Shannon] Yes. The paper never costs that.

54
00:10:41,201 --> 00:10:48,251
[Hal Turing] Let's take the 8B result first. It tests whether the headline survives at a realistic scale.

55
00:10:48,251 --> 00:11:08,401
[Dr. Ada Shannon] It survives, but smaller. The abstract's "over 100 percent" is true for the small base with the biggest table. At 8B with a 16-million-entry table, NaturalQuestions goes from about 25 right per hundred to about 27. That's under two points, roughly a 7 percent relative gain. TriviaQA and MMLU gain several points each at a trillion tokens.

56
00:11:08,401 --> 00:11:10,626
[Hal Turing] And the other benchmarks?

57
00:11:10,626 --> 00:11:28,926
[Dr. Ada Shannon] Mixed. MBPP, a coding benchmark, gets worse. PIQA gets slightly worse. OpenBookQA is basically a tie. HellaSwag and HumanEval gain a point or two. These are single runs with no seeds or error bars, so I wouldn't trust a one- or two-point gap in either direction.

58
00:11:28,926 --> 00:11:32,676
[Hal Turing] You said gains were bigger earlier in training?

59
00:11:32,676 --> 00:11:53,601
[Dr. Ada Shannon] Yes. At 200 billion tokens the MMLU gain is much larger than at a trillion. One reading is that memory helps the model pick up facts faster rather than raising its final ceiling. The paper doesn't test that. Same caution on the comparison to Llama 3.1 8B, trained on 15 trillion tokens. It has a different data mix, so it says little about token efficiency.

60
00:11:53,601 --> 00:11:58,776
[Hal Turing] Now the framing. Matched FLOPs is the selling point. What does it leave out?

61
00:11:58,776 --> 00:12:21,726
[Dr. Ada Shannon] Matched FLOPs isn't matched cost. The 1.3B model carries 3.4 billion parameters. The lookups are bandwidth-bound, and sharding the table adds communication. The paper reports no tokens per second, no MFU, no step time against dense or MoE. The authors concede that matching FFN speed takes non-trivial effort. The six-times kernel speedup is against PyTorch's own operator, not against a dense FFN.

62
00:12:21,726 --> 00:12:23,851
[Hal Turing] So what would settle it?

63
00:12:23,851 --> 00:12:40,651
[Dr. Ada Shannon] Tokens per second and MFU at equal quality, for training and inference, at several memory-group sizes. Then an iso-parameter or iso-wall-clock comparison. I'd also want the optimizer state for a table that size costed, and the paper doesn't do that.

64
00:12:40,651 --> 00:12:44,301
[Hal Turing] The MoE baseline looks thin to me.

65
00:12:44,301 --> 00:13:14,226
[Dr. Ada Shannon] It is. Four to sixteen coarse experts, Llama hyperparameters reused, expert-choice routing in training but top-1 at evaluation. That's a possible train/test mismatch. At 1.3B it barely beats dense and is worse on TriviaQA. There's no fine-grained design with shared experts like DeepSeekMoE, by Damai Dai and colleagues at DeepSeek, 2024. That design would test whether many small experts close the gap on facts. So "MoE underperforms" is established only against this MoE, not MoE in general.

66
00:13:14,226 --> 00:13:16,926
[Hal Turing] And PEER? It already looked strong.

67
00:13:16,926 --> 00:13:36,976
[Dr. Ada Shannon] Right. It roughly ties vanilla Memory, and at 1.3B it beats it on NaturalQuestions while using more parameters. Memory+ adds multi-layer sharing, the gating and a tuned key dimension, plus the kernels. So its edge over PEER may come from those choices and from tuning, not from the lookup idea itself. The paper doesn't separate them.

68
00:13:36,976 --> 00:13:39,976
[Hal Turing] What about the original product-key work?

69
00:13:39,976 --> 00:14:01,826
[Dr. Ada Shannon] The 2019 work changed little in the core idea, so what's different now is scale, tuning and kernels. Lample and colleagues also analyzed key usage and collapse. This paper doesn't, even at 64 million values. We don't know how many slots are ever hit, or whether the extra table holds distinct facts or dead entries. The scaling claim has the same limit: it's plots, not fitted exponents, and NaturalQuestions flattens at the top.

70
00:14:01,826 --> 00:14:04,126
[Hal Turing] The facts claim itself?

71
00:14:04,126 --> 00:14:17,051
[Dr. Ada Shannon] It rests on short-form closed-book QA. Long-form hallucination is explicitly deferred. And there's no comparison to retrieval-augmented generation, the other known way to fix facts.

72
00:14:17,051 --> 00:14:19,477
[Hal Turing] So who should care right now?

73
00:14:19,477 --> 00:14:38,052
[Dr. Ada Shannon] Teams building knowledge-heavy models that have plenty of memory but limited compute. The cost is a huge sharded table in GPU memory, its optimizer state, and a serving footprint. The paper measures none of the serving side. Whether that trade is worth it depends on hardware numbers it omits.

74
00:14:38,052 --> 00:14:42,777
[Hal Turing] The authors also hope for continual learning and less forgetting.

75
00:14:42,777 --> 00:14:52,577
[Dr. Ada Shannon] Those are hopes, not results. Same for faster kernels, sparse updates, and fewer hallucinations. This paper doesn't test editing or forgetting.

76
00:14:52,577 --> 00:15:05,702
[Hal Turing] Then here's where I land. At matched FLOPs, a scaled memory layer clearly stores more facts than the dense and coarse-MoE baselines tested. That's solid for small bases and modest at 8B.

77
00:15:05,702 --> 00:15:18,202
[Dr. Ada Shannon] And the boundary is what it doesn't show. It isn't shown to be cheaper in time or money, or better than a strong MoE. And the benefit is much smaller at 8B, with no evidence yet beyond that.

78
00:15:18,202 --> 00:15:22,927
[Hal Turing] That's our look at memory layers. Thanks for listening, everyone.