% Compile with pdflatex. Replace the demonstration content before delivery.
\documentclass[aspectratio=169,11pt]{beamer}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{amsmath,amssymb,graphicx,booktabs}

\definecolor{AcademicBlue}{RGB}{26,57,86}
\definecolor{AcademicTeal}{RGB}{0,105,112}
\setbeamercolor{structure}{fg=AcademicBlue}
\setbeamercolor{frametitle}{fg=AcademicBlue,bg=white}
\setbeamercolor{title}{fg=AcademicBlue}
\setbeamercolor{alerted text}{fg=AcademicTeal}
\setbeamercolor{block title}{fg=white,bg=AcademicBlue}
\setbeamercolor{block body}{fg=black,bg=AcademicBlue!5}
\setbeamertemplate{navigation symbols}{}
\setbeamertemplate{itemize items}[circle]
\setbeamertemplate{footline}{%
  \hfill\usebeamerfont{page number in head/foot}%
  \insertframenumber\,/\,\inserttotalframenumber\hspace{4mm}\vspace{3mm}}
\setbeamersize{text margin left=8mm,text margin right=8mm}
\setbeamerfont{frametitle}{size=\Large,series=\bfseries}
\setbeamerfont{title}{size=\LARGE,series=\bfseries}
\newcommand{\papersource}[1]{\par\smallskip{\scriptsize\color{gray}#1\par}}

\title[FwPKM]{Fast-weight Product Key Memory}
\subtitle{Sparse test-time learning for episodic memory}
\author{Paper: Tianyu Zhao and Llion Jones, Sakana AI}
\date{arXiv:2601.00671v2 \quad | \quad Paper discussion, 16 September 2026}
\begin{document}
\begin{frame}\titlepage\end{frame}
\begin{frame}{Sparse memory can learn from the current context}
\begin{itemize}
\item Problem: store many new associations without touching a dense memory on every token.
\item FwPKM combines product-key retrieval with chunk-level fast-weight updates.
\item Best-supported benefit: episodic retrieval alongside recurrent or local-attention backbones.
\item Evidence boundaries: rereading costs, slower measured training, and weak long-term retention.
\end{itemize}
\papersource{Secs. 1, 4, 6; PDF pp. 2--3, 7--11, 15.}
\end{frame}
\begin{frame}{Capacity and access cost are different constraints}
\begin{tabular}{lll}\toprule
Memory & Stored history & Main constraint\\\midrule
Full attention & Explicit token KV & Sequence-wide quadratic work\\
GDN / recurrent state & Fixed-size state & Information capacity\\
Dense TTT & Adaptive dense network & Read/write cost grows with size\\
FwPKM & Adaptive sparse table & Storage, collisions, sparse kernels\\\bottomrule\end{tabular}
\bigskip
Fixed memory size gives history-length-independent FwPKM access work; it does not make memory-size scaling free.
\papersource{Secs. 1--2; PDF pp. 2--4. Complexity qualification: reading analysis.}
\end{frame}
\begin{frame}{FwPKM sits before the backbone token mixer}
\centering\includegraphics[width=.85\linewidth,height=.78\textheight,keepaspectratio]{figures/architecture.pdf}
\papersource{Fig. 1, PDF p. 4; architecture crop. Blue: slow; green/red: fast keys/values.}
\end{frame}
\begin{frame}{Two small key tables address 262,144 value slots}
\begin{itemize}
\item $q\in\mathbb R^{512}$ splits into two 256-dimensional subqueries.
\item Two $512\times256$ codebooks: score 1,024 subkeys.
\item Top-8 per codebook $\rightarrow$ 64 candidate pairs $\rightarrow$ final Top-8.
\item $V\in\mathbb R^{262144\times512}$; retrieve $\hat v_t=\sum_{i\in\mathcal I_t}s^{\prime}_{t,i}V_i$.
\end{itemize}\[s_i=-\log(10^{-3}+\|q-K_i\|_2^2)\]
\papersource{Secs. 2--3; Appendix A.2/C; PDF pp. 3--6, 22, 25. IDW scores apply to subkeys.}
\end{frame}
\begin{frame}{A learned gate selects the memory contribution}
\[o_t=g_t\hat v_t+(1-g_t)v_t,\quad g_t\in(0,1)\]\begin{itemize}
\item Slow projections construct query, value, and scalar gate from the hidden state.
\item The mixed value is normalized and projected back to the backbone.
\item The same gate weights the local writing loss.
\item Read with existing weights; update only after finishing the chunk.
\end{itemize}
\papersource{Sec. 3.1--3.3; PDF pp. 4--5, Eqs. (8)--(13).}
\end{frame}
\begin{frame}{Lookahead stores what follows a similar query}
\[\mathcal L_{\mathrm{mem}}^{\mathrm{LA}}=\sum_{t=1}^{C-1}\tfrac12g_t\|v_{t+1}-\hat v_t\|_2^2\]\begin{itemize}
\item Default target: next-token value representation, z-score normalized over features.
\item Targets are available before the chunk-end write; same-chunk predictions cannot use that write.
\item Training drops one unpaired position per 512-token chunk.
\item Inference caches available query--target pairs; update frequency depends on the evaluation protocol.
\end{itemize}
\papersource{Appendix A.1--A.5; PDF pp. 22--23, Eq. (19).}
\end{frame}
\begin{frame}{Slot updates aggregate competing residuals}
\[V_i\leftarrow V_i-\frac{1}{N_i^{\mathrm{read}}}\nabla_{V_i}\mathcal L_{\mathrm{mem}}^{\mathrm{LA}}\]\begin{itemize}
\item $N_i^{\mathrm{read}}$: accesses to slot $i$ in the chunk; only accessed rows update.
\item Sum the reconstruction loss, then normalize gradients per accessed row.
\item A single slot with unit gate and unit retrieval weight can be rewritten in one step.
\item With soft weights and collisions, one pass need not reconstruct perfectly. This motivates rereading.
\end{itemize}
\papersource{Sec. 3.3; Appendix A.4; PDF pp. 5, 22. Single-slot example: teaching derivation.}
\end{frame}
\begin{frame}{Address balancing has a sign inconsistency in the paper}
\[\mathcal L_{\mathrm{addr}}=-H(\bar p^{(1)})-H(\bar p^{(2)})\]\begin{itemize}
\item Intended objective: maximize entropy of chunk-averaged subkey usage.
\item Eq. (17) expands this with a minus sign before each sum of p log p, which is positive entropy.
\item The text and Algorithm 1 instead imply minimizing the sum of p log p.
\item We flag the inconsistency; the author implementation was not executed or audited.
\end{itemize}
\papersource{Sec. 3.4, Eq. (17), PDF p. 6; Algorithm 1, p. 24. Sign diagnosis: reading analysis.}
\end{frame}
\begin{frame}{The study tests small hybrid language models}
\begin{itemize}
\item 12 layers, hidden size 768; baselines about 112M parameters.
\item 10B tokens: 5B LongContext64 + 5B Fineweb-Edu; training length 4K.
\item Backbones: GDN, GDN+SWA, GDN+FA, FA; hybrids use a 3:1 ratio.
\item PKM replaces MLPs; FwPKM precedes token mixers at selected layers 2, 6, 10.
\item FwMLP and LaCT provide dense fast-weight comparisons; their update rules and state sharing differ.
\end{itemize}
\papersource{Sec. 4.1; Appendix C; PDF pp. 6--7, 25--26.}
\end{frame}
\begin{frame}{Evaluation changes the memory update protocol}
\begin{tabular}{p{.17\linewidth}p{.48\linewidth}p{.23\linewidth}}\toprule
Test & State/update protocol & Metric\\\midrule
Corpus PPL & Ordered 4K segments, no fast-state reset between adjacent segments; chunk 512 & All-token PPL\\
NIAH & One update after each full-context read; 1--4 reads & Retrieval accuracy\\
LongBench & One context read/update; no answer-time updates & Answer-token PPL\\
Adaptation & Six domains, slow weights frozen & Change in NLL\\\bottomrule\end{tabular}
\papersource{Secs. 4.2--4.5; PDF pp. 7--11. NIAH also freezes fast weights during answer generation.}
\end{frame}
\begin{frame}{Sparse fast memory improves GDN corpus perplexity}
\centering\begin{tabular}{lrrr}\toprule
GDN variant & Fineweb-Edu & LC64 & LAMBADA\\\midrule
Baseline & 19.73 & 10.12 & 32.91\\
PKM at 2,6,10 & 18.27 & 9.27 & 29.97\\
FwPKM at 2,6,10 & 18.09 & 8.89 & 26.73\\
PKM at 6 + FwPKM at 2,10 & 17.39 & 8.54 & 25.91\\\bottomrule\end{tabular}
\bigskip
\begin{itemize}\item All entries are PPL (lower is better).
\item The mixed variant reduces LC64 PPL by 15.6\% relative to GDN.
\item These are ordered-stream results with fast memory carried across segments.\end{itemize}
\papersource{Fig. 2, PDF p. 8; selected rows transcribed. Relative reduction calculated from reported values.}
\end{frame}
\begin{frame}{Full attention can suppress use of the added memory}
\begin{itemize}
\item With FA or GDN+FA, gates often cluster near zero and PPL gains are limited.
\item Training-time pSWA limits full attention to a 512-token window with probability 0.9.
\item This encourages FwPKM use; stronger benefits emerge on retrieval tasks.
\item Interpretation: memory utility depends on competition between available information paths.
\end{itemize}
\papersource{Sec. 4.2; Figs. 2--3, PDF pp. 8--9.}
\end{frame}
\begin{frame}{128K retrieval is demonstrated with a rereading budget}
\begin{tabular}{lrr}\toprule
Model, four reads & 4K accuracy & 128K accuracy\\\midrule
GDN + FwPKM at 2,6,10 & 88.6\% & 73.8\%\\
GDN+SWA + PKM/FwPKM mix & 80.8\% & 74.8\%\\
GDN+FA + FwPKM (pSWA) & 85.0\% & 84.0\%\\\bottomrule\end{tabular}
\bigskip
\begin{itemize}\item Five needles per sample; four-character keys and six-digit values.
\item Each read uses the entire haystack as one update chunk.
\item Training is 4K; longer evaluations also apply YaRN to RoPE.
\item Four-read results are not one-pass or equal-latency results.\end{itemize}
\papersource{Fig. 4, PDF p. 10; protocol p. 9; YaRN in Appendix C, p. 25.}
\end{frame}
\begin{frame}{LongBench measures answer probability, not official QA scores}
\begin{itemize}
\item 950 examples across five QA tasks; average 17,421 tokens, maximum 81,616.
\item Only answer tokens are scored; the first reference is used when several exist.
\item GDN to PKM/FwPKM mix: 2WikiMultihopQA PPL 19.50 to 12.69; HotpotQA 36.25 to 21.00.
\item LaCT is omitted from the plot because PPL exceeds 100.
\item This supports better answer likelihood, not an established gain in F1 or exact match.
\end{itemize}
\papersource{Sec. 4.4, Fig. 5; PDF pp. 9--11.}
\end{frame}
\begin{frame}{Fast adaptation does not ensure long-term retention}
\begin{itemize}
\item Six Pile domains; each has separate 4M-token adaptation and evaluation sets.
\item After adapting to one domain, evaluate all six domains against initial NLL.
\item Current-domain NLL improves, but later stages often erase the gain.
\item The paper motivates consolidation and retention mechanisms; it does not provide a solution.
\end{itemize}
\papersource{Sec. 4.5, Fig. 6; PDF pp. 10--12.}
\end{frame}
\begin{frame}{Lookahead is crucial for the retrieval behavior}
\centering\begin{tabular}{lr}\toprule
GDN+SWA mixed variant & 128K NIAH, four reads\\\midrule
Default & 74.8\%\\
Without addressing loss & 37.2\%\\
Without gate-weighted writing loss & 19.4\%\\
Without lookahead targets & 1.8\%\\\bottomrule\end{tabular}
\bigskip
\begin{itemize}\item Removing addressing loss uses MSE to update both keys and values.
\item Small PPL changes can hide a large loss of retrieval ability.
\item Larger Top-k / more heads do not improve every task consistently.\end{itemize}
\papersource{Appendix D.1, Figs. 9--12; PDF pp. 27--28.}
\end{frame}
\begin{frame}{Address balancing helps, but usage remains nonuniform}
\centering\begin{tabular}{lrrr}\toprule
LC64 configuration & Collision & Coverage & KLD\\\midrule
Default & 85.2\% & 28.4\% & 2.1\\
Without addressing loss & 96.4\% & 8.0\% & 5.4\\
One head, Top-32 & 95.4\% & 55.0\% & 1.9\\\bottomrule\end{tabular}
\bigskip
\begin{itemize}\item Statistics are computed every 32K tokens, then averaged.
\item Lower collision/KLD and higher coverage are preferred within comparable settings.
\item Top-k changes the access budget; uniform subkey marginals do not guarantee uniform joint slots.\end{itemize}
\papersource{Appendix D.3, Table 4; PDF p. 31. Joint-distribution qualification: reading analysis.}
\end{frame}
\begin{frame}{Freezing memory separates dynamic use from added capacity}
\begin{tabular}{lrr}\toprule
GDN+SWA + PKM/FwPKM mix & Updating & Frozen\\\midrule
Fineweb-Edu PPL & 16.76 & 16.97\\
LC64 PPL & 8.27 & 8.77\\
LAMBADA PPL & 26.43 & 28.07\\\bottomrule\end{tabular}
\bigskip
\begin{itemize}\item Long-context corpora are more affected by disabling updates.
\item Cross-segment PPL visualizations offer supporting evidence.
\item Curve shapes alone do not establish the proposed causal mechanism.\end{itemize}
\papersource{Appendix D.2, Fig. 14, p. 30; Appendix E.1, Fig. 15, p. 32.}
\end{frame}
\begin{frame}{Memory access is traceable, with limits to the interpretation}
\begin{itemize}
\item A 4K/four-read case traces retrieved slots while generating a six-digit passkey.
\item Layer-2 and layer-10 slots often include the correct next token, with some noisy slots.
\item A separate text example shows selective high-layer gates on rare entities.
\item That visualization uses 32-token updates, not the default 512.
\item These are informative cases, not population-level retrieval or novelty-detection guarantees.
\end{itemize}
\papersource{Sec. 5, Figs. 7--8; PDF pp. 12--14; extra example in Fig. 16, p. 33.}
\end{frame}
\begin{frame}{Low active compute has not yet delivered high throughput}
\centering\begin{tabular}{lrrrr}\toprule
GDN variant & Total M & Active M & FLOPs T & Samples/s\\\midrule
Baseline & 112.11 & 112.11 & 22.04 & 54.44\\
PKM & 509.08 & 106.62 & 20.92 & 42.90\\
FwPKM & 519.12 & 116.48 & 22.74 & 19.04\\
PKM/FwPKM mix & 515.77 & 113.19 & 22.14 & 21.27\\\bottomrule\end{tabular}
\bigskip
\begin{itemize}\item FwPKM: 4.63$\times$ total parameters, 1.03$\times$ FLOPs, 35\% of baseline throughput.
\item Single H100, batch 8, 100 training steps; not a decoding-latency benchmark.
\item Better sparse kernels are a proposed next step, not a measured result.\end{itemize}
\papersource{Sec. 6, Table 1; PDF pp. 14--15. Ratios calculated from reported values.}
\end{frame}
\begin{frame}{Take-home: a promising memory component with clear boundaries}
\begin{itemize}
\item Contribution: a large sparse memory that can rewrite context-specific associations online.
\item Mechanism: lookahead reconstruction, selective writing, and anti-collapse addressing.
\item Evidence: streaming PPL, budgeted long-range retrieval, and answer-token likelihood.
\item Open issues: entropy-sign specification, state isolation, retention, and actual serving efficiency.
\item Next experiment: compare under matched memory, rereading compute, and end-to-end latency.
\end{itemize}
\papersource{Paper conclusions: Sec. 8, p. 17. Proposed next experiment: reading analysis.}
\end{frame}
\begin{frame}{Reference and reading scope}
\textbf{Tianyu Zhao and Llion Jones.}\\
\emph{Fast-weight Product Key Memory}. Sakana AI.\\
arXiv:2601.00671v2, revised 22 February 2026.\\[1em]
\url{https://arxiv.org/abs/2601.00671v2}\\[1em]
Reading: complete main text and Appendices A--E, author TeX and 33-page PDF.\\
No training replication or author-code execution.\\[1em]
Presentation prepared as a paper discussion, not by the paper authors.
\end{frame}
\end{document}