% Self-authored reading deck, based on the arxiv-paper-reader Beamer template. % Build: python3 build_slides.py slides.en.tex \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}{% \hspace{8mm}{\scriptsize\color{gray}JRT \enspace|\enspace arXiv:2407.05483v1} \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}} \newcommand{\jrt}{JRT-RNN} \newcommand{\jp}{JRT-Prompt} \newcommand{\E}{\mathrm{E}} \newcommand{\D}{\mathrm{D}} \title[JRT]{Just read twice: closing the recall gap\\for recurrent language models} \subtitle{Reading order, finite memory, and prefix linear attention} \author{Paper by Simran Arora et al.\\[3mm] \small Independent research-group reading presentation} \date{arXiv:2407.05483v1 (7 July 2024)\\Read on 29 September 2026} \begin{document} \begin{frame} \titlepage \vspace{-5mm} \begin{center} \color{AcademicTeal} Let the task inform what a recurrent model stores. \end{center} \end{frame} \begin{frame}{Knowing the question changes what to remember} \begin{block}{A fixed state cannot retain every fact in a long document} The selection problem is harder when the question arrives \emph{after} the document. \end{block} \medskip \begin{tabular}{@{}p{0.40\linewidth}p{0.54\linewidth}@{}} \textbf{Input order} & \textbf{What selection can use}\\[2mm] Document $\to$ Question & Predict which facts will be useful.\\[3mm] Question $\to$ Document & Select facts for a known request.\\[3mm] \shortstack[l]{Document $\to$ Question\\$\to$ Document} & Revisit facts after seeing the request. \end{tabular} \medskip \alert{Extra reading can improve the use of the same state capacity.} \papersource{Teaching interpretation of Secs. 1--3, PDF pp. 1--5. This is not a guarantee of perfect recall.} \end{frame} \begin{frame}{Two interventions give selection more context} \centering \includegraphics[width=\linewidth,height=0.52\textheight,keepaspectratio]{figures/fig-01-overview.png} \medskip \begin{columns}[T,onlytextwidth] \begin{column}{0.47\textwidth} \textbf{\jp}\\ Repeat input using an existing LM. \end{column} \begin{column}{0.49\textwidth} \textbf{\jrt}\\ Train with a non-causal prefix and causal decoding. \end{column} \end{columns} \papersource{Original Fig. 1, PDF p. 2. JRT-Prompt changes the input; JRT-RNN changes the model and training.} \end{frame} \begin{frame}{Causal models are more sensitive to set order} \centering \includegraphics[width=0.98\linewidth,height=0.59\textheight,keepaspectratio]{figures/fig-02-data-order.png} \begin{itemize} \item Input: two sets with one shared token; output that token. \item The right panel measures accuracy(short $A$) $-$ accuracy(long $A$). \item Smaller states show a larger order penalty for causal Based. \end{itemize} \papersource{Original Fig. 2, PDF p. 4; App. B.1, pp. 22--23. Points are best over 3 learning rates and 2 seeds.} \end{frame} \begin{frame}{Re-reading helps, but a memory lower bound remains} Let $k=\min(|A|,|B|)$, $p$ be the number of passes, and $S$ the state in bits. \begin{block}{Communication argument} Alice holds keys; Bob holds queries. Each pass transfers the state. \[ 2pS\ \geq\ \Omega(k) \quad\Longrightarrow\quad S\ \geq\ \Omega(k/p). \] \end{block} \begin{itemize} \item Two passes can store the smaller set, using $O(nk)$ space for $n$-bit elements. \item The Based construction needs input-dependent convolution and kernel assumptions. \item Constant passes cannot solve arbitrarily large recall tasks with a fixed finite state. \end{itemize} \papersource{Prop. 3.1 and Thm. 3.2, p. 5; Thm. G.15, pp. 44--45; Prop. G.29, pp. 52--53.} \end{frame} \begin{frame}{The prompt changes both repetition and question order} \begin{block}{Abstract idea in the main text} \[ \widehat{\mathcal Y} =\mathcal A(\mathcal C,\mathcal Q,\mathcal C,\mathcal Q). \] \end{block} \textbf{Templates in Appendix F} \[ \begin{aligned} \text{Default:}&\quad \text{Document}\ \to\ \text{Cloze cue}\\ \text{Twice:}&\quad \text{Question}\ \to\ \text{Document}\ \to\ \text{Document}\ \to\ \text{Cloze cue}. \end{aligned} \] \begin{itemize} \item No weight updates; the same document information appears twice. \item The reported gain is for this complete prompting strategy. \item These comparisons do not isolate repetition from question-first ordering. \end{itemize} \papersource{Sec. 3.2, p. 5; Table 1, p. 6; actual templates in App. F, pp. 34--39. Last point: reading analysis.} \end{frame} \begin{frame}{JRT-Prompt improves several off-the-shelf models} \centering \begin{tabular}{@{}lrrr@{}} \toprule Model / parameters / training tokens & Default & JRT & $\Delta$ pp\\ \midrule Transformer++ / 1.3B / 50B & 47.8 & 57.8 & +10.0\\ Mamba / 1.3B / 50B & 34.7 & 46.5 & +11.8\\ Based / 1.3B / 50B & 40.3 & 51.8 & +11.5\\ GLA / 1.3B / 100B & 36.7 & 48.9 & +12.2\\ Mamba / 2.8B / 300B & 42.8 & 52.6 & +9.8\\ Mamba-2 / 2.7B / 300B & 48.9 & 58.8 & +9.9\\ \bottomrule \end{tabular} \medskip \begin{itemize} \item Six-task average accuracy; higher is better. \item Default input: up to 1k tokens; JRT: up to 2k. \item Transformers also benefit; some individual task scores decrease. \end{itemize} \papersource{Selected rows from Table 1, PDF p. 6. Differences use rounded averages. Model-count caveat in backup.} \end{frame} \begin{frame}{Linear attention compresses history into two running sums} Use row-vector features $\phi(q_i),\phi(k_i)\in\mathbb R^{1\times\tilde d}$ and values $v_i\in\mathbb R^{1\times d_v}$. \[ \begin{aligned} s_i&=s_{i-1}+\phi(k_i)^\top v_i, &\qquad s_i&\in\mathbb R^{\tilde d\times d_v},\\ z_i&=z_{i-1}+\phi(k_i)^\top, &\qquad z_i&\in\mathbb R^{\tilde d\times1}. \end{aligned} \] \[ \boxed{\displaystyle y_i=\frac{\phi(q_i)s_i}{\phi(q_i)z_i}} \] \begin{itemize} \item Write a key--value outer product, then read with the query. \item The denominator normalizes by accumulated key similarity. \item Per-token state and compute are constant in sequence length, with cost $O(\tilde d d_v)$ per head. \end{itemize} \papersource{Eqs. (1)--(2), PDF p. 7. Row-vector notation made dimensionally consistent for this presentation.} \end{frame} \begin{frame}{Prefix linear attention adds a non-causal state} \[ \begin{aligned} s_\E&=\sum_{j=1}^{M}\phi(k_{e_j})^\top v_{e_j}, &z_\E&=\sum_{j=1}^{M}\phi(k_{e_j})^\top,\\ s_{\D,i}&=\sum_{j=1}^{i}\phi(k_{d_j})^\top v_{d_j}, &z_{\D,i}&=\sum_{j=1}^{i}\phi(k_{d_j})^\top. \end{aligned} \] \[ \boxed{\displaystyle y_i=\frac{\phi(q_i)(s_\E+s_{\D,i})} {\phi(q_i)(z_\E+z_{\D,i})}} \] \begin{itemize} \item Separate encoder and decoder K/V projections. \item The $M$ prefix tokens enter both paths; prefix queries see all of them. \item After prefill, merge states and update only with new decoder K/V. \end{itemize} \papersource{Eqs. (3)--(4), p. 7; implementation, pp. 29--30. The query typo in Eq. (3) is corrected here.} \end{frame} \begin{frame}{JRT-RNN trains both the prefix and the causal suffix} \begin{columns}[T,onlytextwidth] \begin{column}{0.48\textwidth} \begin{block}{Prefix: $M=1024$} Non-causal encoding.\\ Mask 15\% of prefix tokens.\\ Predict their original identities. \end{block} \end{column} \begin{column}{0.48\textwidth} \begin{block}{Suffix: $N-M=1024$} Causal language modeling.\\ Standard next-token prediction.\\ No future suffix access. \end{block} \end{column} \end{columns} \[ \mathcal L= \frac{\mathcal L_{\mathrm{NTP}}+ 0.25\,\mathcal L_{\mathrm{MLM}}}{1.25}. \] \begin{itemize} \item Replace Based linear-attention layers with PLA; retain a hybrid of local attention and short convolutions. \item Train on the Pile, with matched data order across model groups. \item No training-time mask tokens are inserted at inference. \end{itemize} \papersource{Sec. 4.3, Eq. (5), p. 8; App. B, pp. 23--24; Table 15, p. 54. Configuration caveats in notes.} \end{frame} \begin{frame}{The quality benchmark tests document-grounded completion} \begin{tabular}{@{}p{0.26\linewidth}p{0.67\linewidth}@{}} \toprule Tasks & FDA, SWDE, NQ, SQuAD, TriviaQA, DROP\\[2mm] Prompt format & Cloze completion; questions rewritten using Llama-3-70B, with the answer removed\\[2mm] Accuracy & Generated text contains the specified exact answer span\\[2mm] Generation & Stop at a newline or 48 generated tokens\\[2mm] Main average & Nine columns: two lengths each for FDA/SWDE/NQ, plus three Full-document tasks\\ \bottomrule \end{tabular}\par \medskip \alert{This is a processed evaluation suite, with its own metric and weighting.} \papersource{Table 2, p. 9; App. B.5 and Table 7, pp. 24--26. Truncated documents retain the answer span.} \end{frame} \begin{frame}{JRT-RNN closes much of the average recall gap} \centering \begin{tabular}{@{}lrrrr@{}} \toprule Params / tokens & Transformer & Mamba & Based & JRT-RNN\\ \midrule $\sim$360M / 30B & 43.4 & 24.0 & 29.2 & \textbf{42.9}\\ $\sim$1.3B / 10B & 43.9 & 25.2 & 33.7 & \textbf{42.0}\\ $\sim$1.3B / 50B & 51.4 & 36.8 & 42.6 & \textbf{49.5}\\ \bottomrule \end{tabular} \medskip \begin{block}{The reported 99\% and 96\% are score ratios} \[ \frac{42.9}{43.4}\approx98.85\%, \qquad \frac{49.5}{51.4}\approx96.30\%. \] \end{block} \begin{itemize} \item Improvement over Based: \textbf{13.7 pp} and \textbf{6.9 pp}. \item The three 512-token JRT conditions use read-twice padding. \end{itemize} \papersource{Table 2 and Sec. 5.1, PDF p. 9. These are nine-column averages, not universal capability ratios.} \end{frame} \begin{frame}{Read-twice padding contributes to the main quality result} \textbf{JRT-RNN, 1.3B / 50B; documents of 512 tokens} \medskip \begin{center} \begin{tabular}{@{}lrrr@{}} \toprule Inference strategy & FDA & SWDE & NQ\\ \midrule Left-pad & 75.8 & 49.3 & 30.9\\ Read-2$\times$ & \textbf{86.7} & \textbf{49.4} & \textbf{38.3}\\ Iterative encoding & 80.2 & 43.3 & 34.2\\ \bottomrule \end{tabular} \end{center} \begin{itemize} \item Fill the fixed $M=1024$ prefix with another copy of the context. \item At 360M, FDA rises from 61.9 to 82.0: a 20.1 pp change. \item Iterative encoding reprocesses generated text and is more expensive; it does not win here. \end{itemize} \papersource{Table 4 and Sec. 5.1, PDF pp. 9--10. The main table uses the Read-2x rows for these tasks.} \end{frame} \begin{frame}{The 2k setting still leaves a substantial Transformer gap} \textbf{1.3B / 50B; 2k document inputs} \medskip \begin{center} \begin{tabular}{@{}lrrr@{}} \toprule Model & FDA & SWDE & NQ\\ \midrule Transformer & \textbf{79.7} & \textbf{55.5} & \textbf{30.2}\\ Mamba & 21.0 & 29.9 & 23.1\\ Based & 36.1 & 37.7 & 23.4\\ JRT-RNN & 55.2 & 41.4 & 26.2\\ \bottomrule \end{tabular} \end{center} \begin{itemize} \item JRT improves over recurrent baselines, but FDA remains 24.5 pp behind. \item 2k exceeds the \emph{encoder region} ($M=1024$), not the total pretraining length ($N=2048$). \item The small-model rows in Table 3 are labeled 360M/10B, not 30B. \end{itemize} \papersource{Table 3, PDF p. 9; training context in Sec. 5.2, p. 10.} \end{frame} \begin{frame}{JRT helps recall across larger gaps between repeated bigrams} \centering \includegraphics[width=0.90\linewidth,height=0.69\textheight,keepaspectratio]{figures/fig-03b-recall-distance.pdf} \par \alert{Lower perplexity is better. JRT improves on Based and Mamba; attention often remains stronger.} \papersource{Original Fig. 3, top right, p. 10. 360M/30B; bigrams seen fewer than 1000 times in training.} \end{frame} \begin{frame}{Better recall comes with a non-recall perplexity tradeoff} \centering \includegraphics[width=0.88\linewidth,height=0.61\textheight,keepaspectratio]{figures/fig-03c-other-frequency.pdf} \begin{itemize} \item Rare non-recall tokens are harder for JRT in this comparison. \item Evaluate only the final 1024 causal positions of 2048-token sequences. \item Reduced NTP supervision is the authors' explanation, not an isolated cause. \end{itemize} \papersource{Original Fig. 3, bottom, p. 10; protocol: App. B.4, p. 24. 3200 test sequences.} \end{frame} \begin{frame}{General language understanding shows mixed changes} \textbf{SuperGLUE reported averages, $\sim$360M / 30B} \begin{center} \begin{tabular}{@{}l@{\hspace{10mm}}rr@{}} \toprule Model & Zero-shot & Five-shot\\ \midrule JRT-RNN & 39.6 & 45.3\\ Based & 40.9 & 48.1\\ Transformer & 42.0 & 45.3\\ Mamba & 43.7 & 45.5\\ \bottomrule \end{tabular} \end{center} \begin{itemize} \item Strong recall gains do not imply gains on every downstream task. \item The paper does not establish statistical equivalence. \item The 1.3B table omits ReCoRD; do not compare its average across scales. \end{itemize} \papersource{Tables 8--9, PDF p. 27. Official LM-Eval Harness settings; retain the paper's reported averages.} \end{frame} \begin{frame}{Custom CUDA yields large gains for long prefill kernels} \centering \begin{tabular}{@{}lrr@{}} \toprule Implementation & 2048 tokens & 32768 tokens\\ \midrule FlashAttention-2 & 0.5 ms & 107.8 ms\\ Based Triton (FLA) & 1.0 ms & 123.7 ms\\ Based Custom CUDA & 0.3 ms & 4.5 ms\\ JRT-Prompt Custom CUDA & 0.6 ms & 9.0 ms\\ JRT-RNN Custom CUDA & 0.4 ms & 5.6 ms\\ \bottomrule \end{tabular} \medskip \begin{block}{Reported speedups at 32k, batch 16, NVIDIA H100} \jp: \textbf{11.9$\times$}; \quad \jrt: \textbf{19.2$\times$}, relative to FA2. \end{block} \begin{itemize} \item Average of 20 iterations. Prompt repetition is charged at $2\times$ Based time. \item Kernel prefill timing; no end-to-end generation speedup is established. \end{itemize} \papersource{Table 5, PDF p. 11; Sec. 5.3, pp. 11--12. Short-input JRT-Prompt is slower than FA2 here.} \end{frame} \begin{frame}{The strongest claims depend on four boundaries} \begin{enumerate} \item \textbf{Memory.} Re-reading does not remove the finite-state lower bound. Benefits depend on how much relevant information is needed. \medskip \item \textbf{Attribution.} Prompt order, repetition, training objectives, and padding all change. The main tables evaluate complete recipes. \medskip \item \textbf{Length.} Quality is measured on short inputs and around 2k; 32k is a separate efficiency test. \medskip \item \textbf{Scope.} JRT is a small-to-mid-scale pretrained prototype. Generalization to large instruction models is not demonstrated. \end{enumerate} \papersource{Reading synthesis of Secs. 3--5 and Apps. B--G. Evidence boundaries, not additional experiments.} \end{frame} \begin{frame}{The design lesson is to condition compression on the task} \begin{block}{A useful question for efficient language models} When the model compresses a document, does it already know what information will matter? \end{block} \medskip \begin{itemize} \item \textbf{JRT-Prompt:} revisit the input after acquiring more context. \medskip \item \textbf{JRT-RNN:} make the prefix visible non-causally, then retain efficient causal state updates. \medskip \item \textbf{Next experiment:} isolate question placement from repetition, and compare complete systems at matched quality and cost. \end{itemize} \papersource{Paper conclusion, PDF p. 12. Final experiment suggestion is this presentation's analysis.} \end{frame} \begin{frame}{Source and reading scope} \textbf{Just read twice: closing the recall gap for recurrent language models} \medskip Simran Arora, Aman Timalsina, Aaryan Singhal, Benjamin Spector, Sabri Eyuboglu, Xinyi Zhao, Ashish Rao, Atri Rudra, Christopher R\'e. \medskip \href{https://arxiv.org/abs/2407.05483v1}{arXiv:2407.05483v1}, submitted 7 July 2024. \medskip \href{https://github.com/HazyResearch/prefix-linear-attention}{Official repository linked by the paper}. The repository was not executed for this reading. \medskip \begin{itemize} \item Same-version author TeX and 56-page PDF were used. \item Main text, substantive Appendices A--G, and training tables were read. \item Original figures are copied from author assets; selected tables are re-typeset. \item Suggested talk: about 25 minutes through slide 20; backups follow. \end{itemize} \papersource{Companion Chinese notes document detailed source checks and version-specific inconsistencies.} \end{frame} \appendix \begin{frame}{Backup: rare repeated bigrams benefit from prefix processing} \centering \includegraphics[width=0.90\linewidth,height=0.69\textheight,keepaspectratio]{figures/fig-03a-recall-frequency.pdf} \par \alert{Recall-frequency analysis restricts repeated bigrams to gaps of at least 1024 tokens.} \papersource{Original Fig. 3, top left, p. 10; App. B.4, p. 24. 360M/30B; lower perplexity is better.} \end{frame} \begin{frame}{Backup: two passes expose the smaller set before a useful scan} \[ A_1\quad B_1\quad A_2\quad B_2 \] \begin{center} \begin{tabular}{@{}p{0.23\linewidth}p{0.29\linewidth}p{0.35\linewidth}@{}} \toprule Case & Store & Compare against\\ \midrule $|A|\leq |B|$ & $A_2$ & $B_2$\\[2mm] $|B|<|A|$ & $B_1$ & $A_2$\\ \bottomrule \end{tabular} \end{center} \begin{itemize} \item The construction uses the separator and known input length to identify the smaller set. \item Store $k=\min(|A|,|B|)$ elements: $O(nk)$ space for $n$-bit items. \item This is an existence construction, not an optimization guarantee for any pretrained RNN. \end{itemize} \papersource{Teaching reconstruction of Thm. G.15 and Algorithm 3, PDF pp. 44--45.} \end{frame} \begin{frame}{Backup: a matching kernel separates hits from misses} For distinct tokens, assume near-orthogonal features: \[ \langle\phi(a),\phi(b)\rangle= \begin{cases} 1,&a=b,\\ \text{a value in }[-\epsilon,\epsilon],&a\ne b. \end{cases} \] Summing over a stored set $A$ with $\epsilon\leq 1/(3|A|)$: \[ \text{no match: }\rho\leq\tfrac13, \qquad \text{match: }\rho\geq\tfrac23. \] \begin{itemize} \item A thresholding MLP separates the two cases. \item The $O(nk)$ Based construction uses input-dependent reordering and IP features. \item It does not prove that a fixed Taylor map learns exact recall. \end{itemize} \papersource{Prop. G.18, pp. 46--47; Thm. G.22 and Remark G.23, p. 51. Full assumptions and caveats in notes.} \end{frame} \begin{frame}{Backup: prefix processing helps beyond the default hybrid} \textbf{360M / 30B; selected Table 12 columns} \begin{center} \begin{tabular}{@{}lrrr@{}} \toprule Mixer & SWDE-512 & FDA-1024 & FDA-2048\\ \midrule Linear attention & 29.6 & 16.0 & 4.6\\ Prefix linear attention & 36.8 & 48.7 & 8.2\\ Linear + sliding attention & 25.4 & 8.1 & 3.0\\ Prefix linear + sliding & 35.5 & 46.5 & 30.0\\ \bottomrule \end{tabular} \end{center} \begin{itemize} \item Prefix variants improve these reported comparisons. \item Pure prefix linear attention still struggles on FDA-2048. \item The no-MLM Table 11 lacks a matched full-JRT row: it does not isolate the magnitude of the MLM effect. \end{itemize} \papersource{App. C.2, Tables 11--12, PDF p. 28. Table 11 uses 10B training tokens; Table 12 uses 30B.} \end{frame} \begin{frame}{Backup: repetition reduces some errors and creates others} \textbf{Mamba 2.8B, SQuAD; percentage of all test examples} \begin{center} \begin{tabular}{@{}lrr@{}} \toprule Error category & Default & JRT-Prompt\\ \midrule Irrelevant output & 51.3 & 13.8\\ Repetition & 5.5 & 32.9\\ \bottomrule \end{tabular} \end{center} \begin{itemize} \item A base LM may continue a repetition pattern instead of answering. \item Few-shot is not uniformly inferior: on SWDE, Mamba 2.8B scores 51.9 with two examples versus 48.9 with JRT. \item Table 13 trends vary by task; no single error pattern explains all gains. \end{itemize} \papersource{App. E, Tables 13--14, PDF pp. 31--33.} \end{frame} \begin{frame}{Backup: the CUDA path accumulates the encoder state first} \begin{enumerate} \item Apply the Based kernel to encoder K/V, using a sum instead of a causal cumulative sum and skipping query multiplication. \medskip \item Keep the encoder KV state in registers. \medskip \item Run the causal decoder kernel from this initial state, then write the outputs. \end{enumerate} \medskip \[ \text{Taylor features for 16 inputs:}\quad 1+16+16^2=273. \] \alert{Managing this expanded state is essential to realizing the asymptotic benefit.} \papersource{Algorithm 2 and App. D, PDF pp. 29--30; Taylor features in App. B, p. 23.} \end{frame} \begin{frame}{Backup: several v1 claims need careful bookkeeping} \begin{itemize} \item \textbf{Model count.} Table 1 has 16 configurations: 14 recurrent and 2 Transformer rows. Recomputing six-task gains gives about 11.28 pp for the 14 recurrent rows. \medskip \item \textbf{Supervised positions.} The stated mask policy implies $0.5+0.15(0.5)=57.5\%$, while the prose says 65\%. \medskip \item \textbf{Theory.} Prop. G.26's reverse reduction uses intersection membership, more information than a Boolean disjointness oracle. \medskip \item \textbf{Configuration.} Window sizes and layer totals in Tables 15--16 do not fully agree with the prose. \end{itemize} \papersource{Reading checks of Table 1, p. 6; Sec. 5.2, p. 10; Prop. G.26, p. 52; Tables 15--16, pp. 54--55.} \end{frame} \end{document}