% 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{MLP Memory}
\author{Research paper discussion}
\date{2026-11-11}
\begin{document}
\begin{frame}{MLP Memory: learn retrieval before deployment}
\textbf{A Retriever-Pretrained Memory for Large Language Models}
\medskip

Rubin Wei, Jiaqi Cao, Jiarui Wang, Jushi Kai, Qipeng Guo, Bowen Zhou, Zhouhan Lin
\medskip

\href{https://arxiv.org/abs/2508.01832v4}{arXiv:2508.01832v4} $\cdot$ Research group discussion $\cdot$ 20--30 min
\begin{block}{Core idea}
Distill token-level retrieval distributions into a separate MLP, then mix its probabilities with a frozen language model.
\end{block}
Published at \href{https://proceedings.iclr.cc/paper_files/paper/2026/hash/d72ae75abaa70a3b19c5d4f436c680d1-Abstract-Conference.html}{ICLR 2026}. \href{https://github.com/LUMIA-Group/MLPMemory}{Official code and checkpoints} available.
\papersource{Status checked 2026-11-11. This presentation analyzes arXiv v4.}
\end{frame}
\begin{frame}{Knowledge access has deployment trade-offs}
\begin{center}\begin{tabular}{lll}\toprule
Approach & Knowledge access & Main cost\\\midrule
RAG & Retrieved documents & Search + longer context\\
$k$NN-LM & Retrieved token labels & Datastore + token search\\
CPT / LoRA & Adapted model weights & Potential task degradation\\
MLP Memory & Separate parametric branch & Offline distillation + extra MLP\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item Move retrieval work into training; retain a direct forward pass at inference.
\item Long-term memory here means corpus knowledge stored in parameters.
\end{itemize}
\papersource{Sec. 1--2, PDF pp. 2--4.}
\end{frame}
\begin{frame}{The teacher is a token-level retriever}
\[
(\mathcal K,\mathcal V)=\{(f(c_t),w_t):(c_t,w_t)\in\mathcal D\}
\]
\[
p_{kNN}(y\mid c)\propto\sum_{(k_i,v_i)\in\mathcal N}
\mathbb I[y=v_i]\exp[-d(k_i,f(c))]
\]\begin{itemize}
\item $c_t$: prefix; $w_t$: next token; $f(c_t)\in\mathbb R^d$: context vector.
\item Retrieve nearby context vectors, then aggregate their next-token labels.
\item Sparse vocabulary distribution supplies alternatives beyond one observed token.
\end{itemize}
\papersource{Sec. 2, Eqs. (2)--(3), PDF p. 4. Normalization is implicit in the paper.}
\end{frame}
\begin{frame}{A separate MLP replaces search at inference}
\begin{block}{Frozen backbone: encode the prefix}
\[c_t\ \longrightarrow\ \text{Transformer}\ \longrightarrow\ (h_t,\ p_{LM})\]
\end{block}
\begin{block}{Memory branch: predict retrieval-like probabilities}
\[h_t\ \longrightarrow\ \text{stacked MLPs}\ \longrightarrow\ p_{MLP}\]
\end{block}
\[
p_{final}=\lambda p_{MLP}+(1-\lambda)p_{LM}
\]
\papersource{Teaching schematic based on Fig. 4, PDF p. 5. We follow Eq. (7)'s memory-weight convention; Fig. 4 reverses it.}
\end{frame}
\begin{frame}{No token mixing is needed inside the memory}
\[\mathcal M:\mathbb R^d\longrightarrow\mathbb R^{|V|}\]\begin{itemize}
\item Each query vector already summarizes the prefix through the Transformer.
\item Stack feed-forward, residual and normalization operations; output a vocabulary distribution.
\item Llama2 / Mistral: 8 layers; hidden size 4096; intermediate size 11008 / 14336.
\item Initialize from the corresponding backbone MLP modules. Roughly 1B parameters, excluding embeddings.
\end{itemize}
\papersource{Sec. 3.1 and Appendix I, PDF pp. 5--6, 26--27; Table 11.}
\end{frame}
\begin{frame}{Build the retrieval targets once, offline}
\begin{enumerate}
\item Run the backbone over the training corpus; store context vectors and next tokens.
\item Retrieve neighbors for each training context, excluding the query itself.
\item Cache vector--distribution pairs for memory training.
\item Optimize the memory branch; deploy without the datastore.
\end{enumerate}\begin{itemize}
\item Self-exclusion prevents trivial lookup of the training example.
\item Deployment savings do not remove the offline indexing and search costs.
\end{itemize}
\papersource{Sec. 3.2, PDF p. 6.}
\end{frame}
\begin{frame}{Hybrid supervision combines retrieval and truth}
\begin{align*}
\mathcal L_{KL}&=\mathrm{KL}(p_{kNN}\parallel p_{MLP}),\\
\mathcal L_{CE}&=-\log p_{MLP}(w_t\mid c_t),\\
\mathcal L&=\alpha\mathcal L_{KL}+(1-\alpha)\mathcal L_{CE},\quad\alpha=0.4.
\end{align*}\begin{itemize}
\item KL transfers the similarity-weighted distribution of plausible continuations.
\item CE anchors the prediction to the observed next token.
\item For a fixed teacher, this equals cross-entropy against $\alpha p_{kNN}+(1-\alpha)\delta_{w_t}$, up to a constant.
\end{itemize}
\papersource{Sec. 3.2, Eqs. (4)--(6), PDF p. 6. Final bullet: teaching derivation.}
\end{frame}
\begin{frame}{Mix probabilities after softmax}
\[
p_{final}(w_t\mid c_t)=\lambda p_{MLP}(w_t\mid c_t)
+(1-\lambda)p_{LM}(w_t\mid c_t)
\]\begin{itemize}
\item $\lambda$ weights memory at inference; $\alpha$ weights KL during training.
\item Tune $\lambda$ on each task's validation split. Probability mixing differs from logit averaging.
\item Run the backbone plus the memory MLP; no document access or nearest-neighbor search.
\item Differentiability permits joint optimization; main experiments train the components separately.
\end{itemize}
\papersource{Sec. 3.3, Eq. (7), PDF p. 7; Fig. 4 reverses the weight convention.}
\end{frame}
\begin{frame}{Evaluation separates five questions}
\begin{center}\begin{tabular}{ll}\toprule
Question & Setting / metric\\\midrule
Scaling & GPT-2 + WikiText / Web; PPL $\downarrow$\\
Factual QA & Two 7B backbones; five benchmark scores $\uparrow$\\
General NLP & Mistral; nine classification / entailment tasks\\
Hallucination detection & HaluEval accuracy $\uparrow$\\
Design choices & Loss weight, input layer, teacher and capacity\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item QA uses Wikipedia-2021; general NLP uses a separate mixed-domain corpus.
\item EM for NQ, WebQA, TriviaQA, HotpotQA; MC1/MC2/MC3 average for TruthfulQA.
\end{itemize}
\papersource{Sec. 4 and Appendix A, PDF pp. 7, 18.}
\end{frame}
\begin{frame}{Budgets and baselines shape the comparison}
\begin{itemize}
\item 32 A800 80GB GPUs; AdamW; memory learning rate $4\times10^{-4}$; $\alpha=0.4$.
\item Paper states a training budget equivalent to one epoch of a 7B model. Offline retrieval costs are not itemized.
\item RAG: BGE, top-5 documents. $k$NN-LM: $\lambda=0.1$, temperature 10.
\item LoRA: Q/K/V and MLP layers, rank chosen to match memory parameters. CPT updates all weights.
\item Memory interpolation is tuned per task; the $k$NN baseline uses a fixed interpolation weight.
\end{itemize}
\papersource{Sec. 4, PDF p. 7; Appendix A, Table 4, PDF p. 18.}
\end{frame}
\begin{frame}{Scaling gains refer to fitted exponents}
\centering\includegraphics[width=0.96\linewidth,height=0.47\textheight,keepaspectratio]{figures/scaling.pdf}
\begin{center}\begin{tabular}{lrrr}\toprule
Dataset & $|\gamma_d|$ & $|\gamma_m|$ & Relative increase\\\midrule
WikiText-103 & 0.143 & 0.168 & 17.5\%\\
Web & 0.216 & 0.268 & 24.1\%\\
\bottomrule\end{tabular}\end{center}\centering $PPL=(\beta N)^\gamma$; $N$ includes backbone and memory parameters.
\papersource{Sec. 5.1, Fig. 5, Eqs. (8)--(9), PDF p. 8. These are not PPL reductions.}
\end{frame}
\begin{frame}{More compute continues to help the memory model}
\begin{itemize}
\item WikiText-103: about 100M tokens. Web mixture: about 600M tokens.
\item GPT-2 baselines span 124M--1.5B parameters; augmented models count both components.
\item At the GPT2-xl scale, the Web compute curve keeps declining for the augmented model.
\item Continued decoder training flattens and then rises in the observed range.
\item Interpretation: retrieval imitation supplies a different learning signal; extrapolation beyond the measured scale remains open.
\end{itemize}
\papersource{Sec. 5.1, Fig. 5(c), PDF pp. 8--9.}
\end{frame}
\begin{frame}{Average QA scores improve on both backbones}
\begin{center}\begin{tabular}{lrr}\toprule
Method & Llama2-7B & Mistral-7B\\\midrule
Base & 32.81 & 32.12\\
RAG & 32.94 & 33.38\\
$k$NN-LM & 33.16 & 32.57\\
CPT & 29.66 & 30.53\\
LoRA & 31.34 & 32.28\\
MLP Memory & \textbf{35.38} & \textbf{36.06}\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item Relative increase of the average: 7.8\% / 12.3\%.
\item Mistral gains 3.94 absolute points; this is not a pooled EM score.
\end{itemize}
\papersource{Sec. 5.2, Table 1, PDF p. 9.}
\end{frame}
\begin{frame}{Explicit retrieval still wins on some QA tasks}
\begin{center}\begin{tabular}{lrrr}\toprule
Mistral task & Base & RAG & Memory\\\midrule
NQ & 20.63 & 22.56 & 25.20\\
WebQA & 29.28 & 24.90 & 37.45\\
TriviaQA & 57.65 & 54.21 & 60.99\\
TruthfulQA & 32.09 & 35.47 & 32.54\\
HotpotQA & 20.96 & 29.77 & 24.14\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item Largest memory gains over Base: WebQA +8.17; NQ +4.57.
\item RAG leads Memory on TruthfulQA and HotpotQA.
\item CPT / LoRA do improve some tasks, contrary to the body's blanket degradation wording.
\end{itemize}
\papersource{Table 1, PDF p. 9. Scores follow each benchmark's evaluation metric.}
\end{frame}
\begin{frame}{General NLP improves without changing the backbone}
\begin{center}\begin{tabular}{lr}\toprule
Method & Nine-task average\\\midrule
Base & 67.86\\
RAG & 70.79\\
$k$NN-LM & 68.50\\
CPT & 70.58\\
LoRA & 67.88\\
MLP Memory & 73.07\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item All nine tasks improve over Base; average gain is 5.21 points.
\item Examples: CB 69.64 $\to$ 76.79; RTE 59.57 $\to$ 64.62.
\item A separately trained mixed-domain memory is used for this experiment.
\end{itemize}
\papersource{Sec. 5.3, Table 2, PDF p. 10. Domain-conditional PMI scoring.}
\end{frame}
\begin{frame}{HaluEval improves, but v4 has a column conflict}
\begin{center}\begin{tabular}{lrrr}\toprule
Main Table 3 & Dialogue & QA & Summary\\\midrule
Base & 57.18 & 53.99 & 50.27\\
RAG & 59.06 & 65.09 & --\\
Memory & 66.86 & 64.07 & 52.41\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item These are hallucination detection accuracies, not free-generation hallucination rates.
\item Appendix Table 5 instead gives best Dialogue = 64.07 and QA = 66.86.
\item The best values are swapped across tasks; task-specific gains cannot be reconciled from v4.
\item RAG exceeds Memory on QA under the main-table labeling.
\end{itemize}
\papersource{Sec. 5.4, Table 3, PDF p. 10; Appendix B, Table 5, PDF p. 19.}
\end{frame}
\begin{frame}{Balanced loss and deeper inputs work best}
\centering\includegraphics[width=0.96\linewidth,height=0.45\textheight,keepaspectratio]{figures/ablation.pdf}
\begin{itemize}
\item Best loss balance is near $\alpha=0.4$ across GPT-2 sizes.
\item Input-layer minima lie roughly in the 70--90\% depth range; authors summarize this as about 70\%.
\item Normalized PPL is per curve: values do not compare absolute quality across models.
\end{itemize}
\papersource{Sec. 5.5, Fig. 6, PDF p. 11. High alpha means more KL, despite contrary prose.}
\end{frame}
\begin{frame}{TTFT and throughput show different speedups}
\begin{center}\begin{tabular}{lrr}\toprule
Method & TTFT (ms) $\downarrow$ & TPS $\uparrow$\\\midrule
Base 7B & 22 & 42\\
RAG top-5 & 67 & 24\\
$k$NN-LM & 151 & 6\\
Memory + 7B & 27 & 36\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item Memory vs RAG: 2.48$\times$ TTFT ratio; 1.5$\times$ throughput.
\item Memory vs $k$NN-LM: 5.59$\times$ TTFT ratio; 6$\times$ throughput.
\item Relative to Base, Memory adds about 23\% TTFT and loses about 14\% TPS.
\item Same Wikipedia corpus; $k$NN uses 4096 $\to$ 256 dimensionality reduction.
\end{itemize}
\papersource{Fig. 1, PDF p. 1; Sec. 3.3, PDF p. 7. Ratios computed from plotted values.}
\end{frame}
\begin{frame}{Distribution evidence suggests a sharper head}
\begin{center}\begin{tabular}{lrrr}\toprule
Statistic (20K examples) & LM & $k$NN & Memory\\\midrule
Tokens with nonzero mass & 50257 & 251 & 50257\\
Tokens for 80\% mass & 23 & 22 & 13\\
Tokens for 99\% mass & 617 & 126 & 308\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item Memory retains a dense softmax tail but concentrates mass more than the LM.
\item At 80\% mass, Memory is more concentrated than either comparator.
\item Token case studies support rare-entity gains; they do not establish the mechanism across all contexts.
\end{itemize}
\papersource{Appendix D/H, Tables 7--8/10, PDF pp. 20--21/26.}
\end{frame}
\begin{frame}{Small memory captures much of the measured gain}
\begin{center}\begin{tabular}{lr}\toprule
GPT2-large configuration & WikiText-103 PPL $\downarrow$\\\midrule
Base & 15.81\\
+ 8 layers, 221M & 11.41\\
+ 15 layers, 359M & 11.35\\
+ 36 layers, 772M & 11.25\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item Larger teacher distributions modestly help the small student: 12.91 $\to$ 12.84.
\item Teacher $k$NN PPL improves 9.43 $\to$ 9.42 from $k=1024$ to 2048.
\item The latter is a teacher sensitivity test, not retrained memory performance.
\end{itemize}
\papersource{Appendices E/F/J; Fig. 8 and Tables 9/13, PDF pp. 22--23/27.}
\end{frame}
\begin{frame}{Deploying a stable corpus is the strongest use case}
\begin{itemize}
\item Useful when offline indexing / distillation can be amortized over many low-latency requests.
\item Fixed-size memory inference does not scale with the original datastore size.
\item Updating or deleting knowledge requires parameter changes; no online write mechanism is demonstrated.
\item The branch is tied to a backbone representation space and vocabulary.
\item Explicit retrieval retains source provenance and immediate document updates.
\end{itemize}
\papersource{Analysis based on Sec. 3 and reported experiments.}
\end{frame}
\begin{frame}{Take-home: retrieval can be a training signal}
\begin{enumerate}
\item \textbf{Mechanism:} compile retrieval distributions into a separate MLP, with mixed soft and hard supervision.
\item \textbf{Evidence:} stronger average QA and general NLP scores, plus lower retrieval-related inference latency.
\item \textbf{Trade-off:} offline retrieval and training replace online search; corpus updates become a learning problem.
\item \textbf{Reading caution:} preserve metric definitions and the v4 HaluEval / interpolation inconsistencies.
\end{enumerate}
\end{frame}
\begin{frame}{References and discussion prompts}
\begin{itemize}
\item Wei et al. \href{https://arxiv.org/abs/2508.01832v4}{MLP Memory: A Retriever-Pretrained Memory for Large Language Models}. arXiv:2508.01832v4.
\item \href{https://proceedings.iclr.cc/paper_files/paper/2026/hash/d72ae75abaa70a3b19c5d4f436c680d1-Abstract-Conference.html}{ICLR 2026 proceedings}; \href{https://github.com/LUMIA-Group/MLPMemory}{official implementation}.
\end{itemize}
\medskip
\textbf{Discussion}
\begin{itemize}
\item When do lifetime inference savings offset offline retrieval costs?
\item Can a gate choose between parametric memory and source-backed retrieval?
\item How should updates and provenance work for a distilled memory?
\end{itemize}
\end{frame}
\begin{frame}{Backup: interpolation strength matters}
\begin{center}\begin{tabular}{lrr}\toprule
Appendix Table 5 & $\lambda$ & Best accuracy\\\midrule
Dialogue & 0.45 & 64.07\\
QA & 0.55 & 66.86\\
Summarization & 0.35 & 52.41\\
\bottomrule\end{tabular}\end{center}\begin{itemize}
\item At $\lambda=0.9$, Dialogue = 49.71 and Summary = 46.67, below their baselines.
\item At $\lambda=0.6$, Summary = 49.79, below Base = 50.27.
\item A broad claim of uniform robustness over 0.3--0.6 overstates these results.
\item These appendix task labels conflict with main Table 3.
\end{itemize}
\papersource{Appendix B, Table 5, PDF p. 19.}
\end{frame}
\begin{frame}{Backup: FLOPs are not measured end-to-end speed}
\[
\frac{F_t}{F_m}\approx
\frac{4Ld(2d_a+d_f)+2Ln_{ctx}d_a}{6Ldd_f}
=1+\frac{n_{ctx}}{12d}
\]
\[d_a=d,\qquad d_f=4d.\]\begin{itemize}
\item $L$: layers; $d$: hidden width; $d_f$: feed-forward width; $d_a$: attention width.
\item Compares matched Transformer and MLP non-embedding computations.
\item Does not imply that adding a memory branch makes the original backbone faster.
\end{itemize}
\papersource{Appendix C, Eq. (10), Table 6, PDF pp. 19--20.}
\end{frame}
\end{document}