\documentclass[aspectratio=169,11pt]{beamer}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{lmodern}
\usepackage{amsmath,amssymb,graphicx,booktabs}
\pdfstringdefDisableCommands{\def\newline{ }\def\translate#1{#1}}
\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}}
\newcommand{\fullfig}[2]{\centering\includegraphics[width=\linewidth,height=#2,keepaspectratio]{figures/fig-#1.png}\par}
\title[Mixture-of-Depths]{Mixture-of-Depths}
\subtitle{Dynamically allocating compute in\newline transformer-based language models}
\author[Raposo et al.]{Paper by David Raposo, Sam Ritter, Blake Richards,\newline Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro}
\institute{Google DeepMind; McGill University \& Mila\par\medskip Independent paper discussion}
\date{arXiv:2404.02258v1 $\cdot$ 2 April 2024\par Reading prepared: 8 September 2026}
\begin{document}

% Main presentation: approximately 25 minutes, followed by backup material.
\begin{frame}
\titlepage
\end{frame}

\begin{frame}{Fix the budget; learn which tokens use it}
\begin{block}{Research question}
Can a decoder spend less compute per forward pass while matching the best dense model at the same total training budget?
\end{block}
\begin{itemize}
\item Give each routed block exactly $k$ token slots during training.
\item Learn which tokens receive attention and MLP updates.
\item Let other tokens pass unchanged to later blocks.
\end{itemize}
\par\medskip
\alert{Two useful operating points:} a larger model for lower loss, or a cheaper model for comparable quality.
\papersource{Paper, Secs. 1 and 4.1, PDF pp. 1--2 and 7--10.}
\end{frame}

\begin{frame}{MoD changes effective depth at each token}
\renewcommand{\arraystretch}{1.4}
\begin{tabular}{@{}p{0.20\linewidth}p{0.34\linewidth}p{0.38\linewidth}@{}}
\toprule
Architecture & Routing decision & Consequence \\
\midrule
Dense & Execute every block & Uniform block participation \\
Early exit & Stop processing a token & Cannot re-enter later blocks \\
Typical MoE & Choose MLP experts & More parameter choices per token \\
MoD & Execute block or bypass & Can skip, then re-enter; attention also changes \\
\bottomrule
\end{tabular}
\par\medskip
The choice affects both \alert{which tokens get updated} and \alert{which tokens supply keys and values} in that block.
\papersource{Paper, Secs. 1--2, PDF pp. 2--3; conceptual comparison, not a benchmark of all methods.}
\end{frame}

\begin{frame}{Skipping a middle block does not end computation}
\begin{columns}[c,onlytextwidth]
\column{0.69\textwidth}
\fullfig{01}{5.5cm}
\column{0.28\textwidth}
\begin{itemize}
\item Purple: use block.
\item Orange: bypass.
\item Read vertically to follow one token.
\end{itemize}
\par\medskip
All token states remain in the residual stream.
\end{columns}
\papersource{Paper, Fig. 1, PDF p. 3. Original author figure.}
\end{frame}

\begin{frame}{Capacity controls expensive tensor sizes}
For sequence length $S$ and capacity $C=k$, define $\rho=C/S$.
\[
\underbrace{S\times d}_{\text{all token states}}
\;\longrightarrow\;
\underbrace{C\times d}_{\text{selected states}}
\;\longrightarrow\;
\underbrace{S\times d}_{\text{scatter updates back}}
\]
\begin{itemize}
\item Token projections and MLP work scale roughly with $\rho$.
\item Token-pair attention work scales roughly with $\rho^2$.
\item Routing, data movement, and other layers still cost compute.
\end{itemize}
\begin{block}{Best reported configuration}
Alternate full blocks with routed blocks; $S=2048$, $C=256$.
Only the routed blocks bypass 87.5\% of tokens.
\end{block}
\papersource{Paper, Secs. 3.1--3.2, PDF pp. 4--5; Fig. 3, p. 8. Scaling explanation is a teaching derivation.}
\end{frame}

\begin{frame}{Expert-choice fills a fixed number of slots}
\fullfig{02}{4.9cm}
\par\medskip
One expensive path selects top-$k$ tokens; every other token takes the residual path. No load-balancing loss is needed for this partition.
\papersource{Paper, Sec. 3.3 and Fig. 2, PDF pp. 5--6. Original author figure.}
\end{frame}

\begin{frame}{The routing weight also carries a gradient}
\[
r_i^l=w_\theta^T x_i^l,\qquad \beta=1-C/S
\]
\[
x_i^{l+1}=\begin{cases}
x_i^l+r_i^l f_i(\tilde X^l), & r_i^l>P_\beta(R^l),\\
x_i^l, & r_i^l<P_\beta(R^l).
\end{cases}
\]
\begin{itemize}
\item $x_i^l\in\mathbb R^d$: token state; $r_i^l$: scalar router score.
\item $\tilde X^l$: the $C$ selected states; $f$: attention and MLP update.
\item Multiplying by $r_i^l$ lets the language loss train the router.
\end{itemize}
\alert{The top-$k$ index selection itself is still discrete.}
\papersource{Paper, Eq. (1), Sec. 3.4, PDF pp. 6--7. Equality/tie-breaking is omitted in the paper's equation.}
\end{frame}

\begin{frame}{Selected tokens keep their original causal order}
\centering
\renewcommand{\arraystretch}{1.4}
\setlength{\tabcolsep}{7pt}
\begin{tabular}{@{}l*{8}{c}@{}}
\toprule
Position & 1 & 2 & 3 & 4 & 5 & 6 & 7 & 8 \\
Score & .1 & \alert{.9} & .2 & .3 & \alert{.8} & .4 & .05 & .6 \\
\bottomrule
\end{tabular}
\par\medskip
\raggedright
For $C=2$, positions 2 and 5 enter this block.
\begin{itemize}
\item Position 2 cannot attend to position 5; position 5 can attend to 2.
\item Other positions keep their input states at this block.
\item A later router may select position 3: bypassing is temporary.
\end{itemize}
\par\medskip
Preserve original positions when gathering tokens and constructing masks.
\papersource{Teaching example, not experimental data. Based on Secs. 3.2--3.4 and 5, PDF pp. 4--7 and 12.}
\end{frame}

\begin{frame}{Full-sequence top-$k$ is non-causal}
\begin{block}{A future token can change an earlier decision}
An early score of 0.9 may be in the top-2. If two future scores are 1.0 and 1.1, the early token loses its slot.
\end{block}
\begin{itemize}
\item Causal attention masks do not remove this ranking dependency.
\item During generation, future router scores are unavailable.
\item Training-time top-$k$ therefore needs a causal substitute.
\end{itemize}
\par\medskip
\alert{This is a property of the routing operation, even before considering the attention output.}
\papersource{Paper, Sec. 3.5, PDF p. 7. Numeric illustration is a teaching example.}
\end{frame}

\begin{frame}{Learn a local decision that imitates top-$k$}
\begin{columns}[t,onlytextwidth]
\column{0.48\textwidth}
\begin{block}{Auxiliary loss on the router}
Use top-$k$ membership as a binary target and $r_i^l$ as the logit.
\[
p_i=\sigma(r_i^l)
\]
At generation: threshold the local probability at 0.5.
\end{block}
Reported effect on the main objective: approximately 0.2--0.3\%.
\column{0.48\textwidth}
\begin{block}{Separate MLP predictor}
Predict top-$k$ membership from the same token state, with a stopped input gradient.
\[
\hat p_i=g_\phi(\operatorname{sg}(x_i^l))
\]
At generation: use the local predictor.
\end{block}
Authors report no effect on the main objective and little step-speed overhead.
\end{columns}
\par\medskip
Local predictions do \alert{not} enforce exactly $C$ selections per sequence.
\papersource{Paper, Sec. 3.5, PDF p. 7. Predictor notation and capacity observation are explanatory.}
\end{frame}

\begin{frame}{isoFLOP compares equal total training compute}
\renewcommand{\arraystretch}{1.2}
\begin{tabular}{@{}p{0.36\linewidth}p{0.59\linewidth}@{}}
\toprule
Setting & Reported value \\
\midrule
Batch / sequence length & 128 / 2048 \\
Schedule & Cosine over the full training run \\
Total FLOP budgets & $6\times10^{18}$, $2\times10^{19}$, $10^{20}$ \\
Model size sweep & 60M--3B parameters \\
Main control & Best dense model at each total FLOP budget \\
Causal evaluation data & 256,000 held-out sequences; about 500M tokens \\
\bottomrule
\end{tabular}
\par\medskip
Equal total compute does not imply equal parameter counts, steps, or training tokens.
\papersource{Paper, Sec. 3.6, p. 7; Fig. 4, p. 9; Sec. 4.2, p. 10. Data corpus and full optimizer settings are not reported.}
\end{frame}

\begin{frame}{Learned routing and interleaving preserve quality}
\fullfig{03}{4.8cm}
\par\medskip
At $6\times10^{18}$ FLOPs: alternating 12.5\%-capacity blocks works best; random routing performs much worse.
\papersource{Paper, Sec. 4.1 and Fig. 3, PDF pp. 7--8. Random control shown at 50\% capacity. Original author figure.}
\end{frame}

\begin{frame}{Faster steps can still mean similar total training time}
\renewcommand{\arraystretch}{1.4}
\begin{tabular}{@{}p{0.35\linewidth}p{0.56\linewidth}@{}}
\toprule
Fig. 3 models \#1 and \#3 & Dense baseline vs. MoD \\
\midrule
Parameters & Both approximately 220M \\
Training budget & Both $6\times10^{18}$ FLOPs \\
Language loss & Comparable; text describes MoD as slightly better \\
Training steps per second & MoD: +66\% in the figure caption \\
Total training wall-clock & Approximately the same \\
\bottomrule
\end{tabular}
\par\medskip
A 66\% step-rate gain means $t_{\mathrm{MoD}}/t_{\mathrm{dense}}\approx1/1.66=0.602$: about 39.8\% less time per step.
\papersource{Paper, Sec. 4.1 / Fig. 3, PDF pp. 7--8; same-hardware comparison, axis labeled TPUv5. Time ratio is derived.}
\end{frame}

\begin{frame}{MoD offers lower-loss, lower-cost operating points}
\fullfig{04}{4.9cm}
\par\medskip
Left: optimal MoD models tend to be larger. Right: points below and left of $(1,1)$ improve both loss and forward FLOPs.
\papersource{Paper, Fig. 4, PDF p. 9. Right panel: middle training budget; bubble size denotes parameters. Original author figure.}
\end{frame}

\begin{frame}{The lowest-loss model is not the cheapest-step model}
\begin{columns}[t,onlytextwidth]
\column{0.48\textwidth}
\begin{block}{Optimize loss}
Scale MoD up until forward FLOPs are roughly those of the optimal dense baseline.
\end{block}
The resulting MoD can have more parameters and a lower loss.
\column{0.48\textwidth}
\begin{block}{Optimize step cost}
Choose a smaller MoD operating point while matching the best dense model's quality.
\end{block}
Smaller here need not mean fewer parameters than the dense baseline.
\end{columns}
\par\medskip
\begin{itemize}
\item Authors summarize up to 1.5\% improvement in the language objective.
\item Adding depth rather than width worked better in their experiments.
\end{itemize}
\papersource{Paper, Sec. 1, p. 2; Sec. 4.1, pp. 8--9. Empirical design rules, not optimality guarantees.}
\end{frame}

\begin{frame}{Routing is structured, but its meaning is not settled}
\fullfig{05}{4.3cm}
\par\medskip
Full blocks form dense horizontal bands; some tokens recur across routed blocks. The authors report a preliminary correlation with prediction entropy.
\par\medskip
\alert{This does not establish a causal measure of semantic importance.}
\papersource{Paper, Fig. 5 / Sec. 4.1, PDF pp. 9--10. Right-hand Count axis is logarithmic. Original author figure.}
\end{frame}

\begin{frame}{A causal predictor largely preserves held-out quality}
\fullfig{06}{4.7cm}
\par\medskip
Predictor and top-$k$ curves nearly overlap. Routing accuracy exceeds 97\% early; the plotted plateau is about 98\% (Sec. 3.5 summarizes 99\%).
\papersource{Paper, Sec. 4.2 / Fig. 6, PDF pp. 10--11. Left axis: normalized log-likelihood; lower is better here. Original author figure.}
\end{frame}

\begin{frame}{MoD can also combine with expert routing}
\begin{columns}[c,onlytextwidth]
\column{0.65\textwidth}
\fullfig{07}{5.0cm}
\column{0.32\textwidth}
\textbf{Staged MoDE}\par
Select tokens before attention, then route to MLP experts.
\par\medskip
\textbf{Integrated MoDE}\par
Add a no-op option to MLP expert routing; attention still runs first.
\par\medskip
MoDE's curve improves on MoE here.
\end{columns}
\papersource{Paper, Sec. 4.3 and Fig. 7, PDF pp. 10--12. The plot does not label separate staged/integrated curves. Original figure.}
\end{frame}

\begin{frame}{The evidence supports a compute trade-off, with limits}
\begin{itemize}
\item \textbf{Reproducibility:} corpus, tokenizer, and full optimizer/model configurations are not reported.
\item \textbf{Evaluation scope:} language modeling; no broad downstream task suite or multi-seed uncertainty analysis.
\item \textbf{Deployment:} forward FLOPs do not fully specify decoding latency, batching, or KV-cache behavior.
\item \textbf{Causal routing:} high classification accuracy does not guarantee identical decisions or exact per-sequence capacity.
\end{itemize}
\par\medskip
At 12.5\% capacity, always predicting ``bypass'' already gives 87.5\% accuracy. Quality evaluation remains essential.
\papersource{Reader analysis of reporting and method boundaries; Secs. 3.5--4.3, PDF pp. 7--12. The 87.5\% baseline is derived.}
\end{frame}

\begin{frame}{A token may be useful as memory, even if easy to predict}
\begin{block}{Two reasons to spend compute at a position}
Update its own representation; provide useful keys and values to later positions.
\end{block}
\begin{itemize}
\item This complicates an ``easy token vs. hard token'' interpretation.
\item The authors propose decoupling query, key, and value routing.
\item They also suggest memory and tool-use paths as future extensions.
\end{itemize}
\par\medskip
\alert{The paper does not implement or validate those memory/tool systems.}
\papersource{Paper, Sec. 5, PDF pp. 12--13. Proposed extensions, not experimental results.}
\end{frame}

\begin{frame}{Take-home: allocate depth, then choose an operating point}
\begin{enumerate}
\item \textbf{Mechanism:} a fixed training capacity with learned token participation reduces attention and MLP work.
\item \textbf{Evidence:} interleaved full/routed blocks improve the reported isoFLOP loss--compute frontier.
\item \textbf{Generation:} local predictors approximate non-causal top-$k$ decisions with little observed quality loss.
\end{enumerate}
\par\medskip
\begin{block}{Discussion question}
Can fully causal routing retain both exact budget control and the quality benefits of full-sequence competition?
\end{block}
\papersource{Synthesis of Secs. 3--5, PDF pp. 3--13.}
\end{frame}

\begin{frame}{Source and reading coverage}
\textbf{David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro.}
\par\medskip

\emph{Mixture-of-Depths: Dynamically allocating compute in transformer-based language models.} 2024.
\par\medskip

\href{https://arxiv.org/abs/2404.02258v1}{arXiv:2404.02258v1} \quad (retrieved 8 September 2026)
\begin{itemize}
\item Read the author TeX and the same-version 14-page PDF.
\item Covered all substantive sections, Eq. (1), and Figs. 1--7; no appendix.
\item Figures are unchanged copies from the author source package.
\item No author code was executed; no training experiments were reproduced.
\end{itemize}
\papersource{Companion Chinese notes contain the detailed evidence map, metric caveats, and derivations.}
\end{frame}

\appendix
\begin{frame}{Backup: binary supervision for local routing}
Let $z_i=\mathbf{1}\{i\in\operatorname{TopK}(R^l,C)\}$ and $p_i=\sigma(r_i^l)$.
\[
\mathcal L_{\mathrm{aux}}=-\frac1S\sum_{i=1}^S
\left[z_i\log p_i+(1-z_i)\log(1-p_i)\right].
\]
\begin{itemize}
\item This equation formalizes the prose in Sec. 3.5; it is not a numbered paper equation.
\item Probability threshold $p_i>0.5$ corresponds to raw logit $r_i^l>0$.
\item The paper's Fig. 5 uses a 0--1 ``Router Weight'' axis; do not silently interpret that as the raw linear score.
\item The separate predictor uses a stopped input gradient; full loss coefficients and predictor dimensions are not reported.
\end{itemize}
\papersource{Paper, Sec. 3.5, PDF p. 7; Fig. 5, p. 10. Equation and notation clarification are explanatory.}
\end{frame}

\begin{frame}{Backup: interleaving changes the savings calculation}
For half full blocks and half routed blocks, with $\rho=1/8$:
\[
\text{Token-wise work ratio}\approx\tfrac12(1+\rho)=56.25\%,
\]
\[
\text{Token-pair work ratio}\approx\tfrac12(1+\rho^2)=50.78125\%.
\]
\begin{itemize}
\item These are component-level estimates at fixed architecture.
\item They omit routing, gathering/scattering, output layers, and hardware effects.
\item Incremental decoding requires a separate cache and latency analysis.
\end{itemize}
\papersource{Teaching derivation from Sec. 3.2, PDF pp. 4--5, and the configuration in Fig. 3, p. 8. Not measured speedups.}
\end{frame}

\begin{frame}{Backup: keep the metric denominators separate}
\renewcommand{\arraystretch}{1.25}
\begin{tabular}{@{}p{0.26\linewidth}p{0.36\linewidth}p{0.30\linewidth}@{}}
\toprule
Claim / metric & Comparison or denominator & Evidence \\
\midrule
Up to 1.5\% gain & Relative language objective, not accuracy points & Sec. 1, p. 2 \\
+66\% steps/s & Training models \#3 vs. \#1 on the same hardware & Fig. 3, p. 8 \\
Normalized FLOPs & Forward pass of isoFLOP-optimal dense baseline & Fig. 4, p. 9 \\
Predictor quality & Same 256,000-sequence held-out set & Sec. 4.2, p. 10 \\
97--99\% accuracy & Top-$k$ membership prediction; reports vary & Sec. 3.5 / Fig. 6 \\
\bottomrule
\end{tabular}
\papersource{Abstract also reports sampling step-speed gains above 50\%; Fig. 6 does not provide a full per-configuration decoding latency table.}
\end{frame}
\end{document}