% Based on the arxiv-paper-reader Beamer template. Compile with pdflatex.
\documentclass[aspectratio=169,11pt]{beamer}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{lmodern}
\usepackage{amsmath,amssymb,graphicx,booktabs}

\definecolor{AcademicBlue}{RGB}{26,57,86}
\definecolor{AcademicTeal}{RGB}{0,105,112}
\definecolor{WarmOrange}{RGB}{204,82,28}
\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{\fix}{\textcolor{AcademicTeal}{fixed}}
\newcommand{\std}{\textcolor{WarmOrange}{standard}}

\title[A Fix for Softmax and CE]{A Fix for Softmax and Cross-Entropy}
\subtitle{Why tiny gradients disappear -- and how to retain them}
\author{Paper: Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen\\[3pt]
\normalsize Research group discussion}
\date{\href{https://arxiv.org/abs/2610.31642v1}{arXiv:2610.31642v1} \quad 22 Oct 2026}

\begin{document}

\begin{frame}
  \titlepage
  \vspace{-4mm}
  \centering\scriptsize
  Publication not found as of 27 Oct 2026 (\href{https://jonbarron.info/}{author listing: arXiv}).
  Author JAX examples: \href{https://jonbarron.info/softmaxfix/}{project page}; no standalone official repository found.
\end{frame}

\begin{frame}{Two errors at opposite ends of the gradient scale}
  \begin{block}{One-sentence takeaway}
    Stable exponentiation does not imply stable gradients: preserve small probability mass and compute near-equal probability differences from logits.
  \end{block}
  \begin{itemize}
    \item \textbf{Collapse:} high confidence rounds the target gradient to zero.
    \item \textbf{Cancellation:} near teacher--student agreement creates a gradient noise floor.
  \end{itemize}
  \papersource{Paper, Sec. 1, Fig. 1, PDF p. 1.}
\end{frame}

\begin{frame}{The two failure regimes are visible numerically}
  \centering
  \includegraphics[width=0.96\linewidth,height=0.63\textheight,keepaspectratio]{figures/fig1-two-bugs.png}
  \begin{itemize}
    \item Left: $|g_t|$ should decay smoothly as margin $m$ grows.
    \item Right: $\|\nabla_{z_S}\mathcal L\|_2$ should vanish as $z_S\to z_T$.
  \end{itemize}
  \papersource{Paper, Fig. 1, PDF p. 1. Original PDF figure crop.}
\end{frame}

\begin{frame}{Cross-entropy gradients must sum to zero}
  \[
    q=\operatorname{softmax}(z),\qquad
    \mathcal L(z,p)=-\sum_k p_k\log q_k,\qquad
    g=\nabla_z\mathcal L=q-p.
  \]
  \begin{block}{Translation invariance}
    $\mathcal L(z+c\mathbf 1,p)=\mathcal L(z,p)$, hence $\sum_k g_k=0$.
  \end{block}
  \begin{itemize}
    \item Standard code subtracts $\max_k z_k$ to prevent exponential overflow.
    \item Rounding in later operations can still violate the gradient identity.
  \end{itemize}
  \papersource{Paper, Sec. 1, Eq. (1), PDF p. 1; Sec. 2, PDF p. 2.}
\end{frame}

\begin{frame}{Collapse begins when $1+e^{-m}$ rounds to $1$}
  \[
    m=z_t-\log\sum_{j\ne t}e^{z_j},\qquad
    q_t=\frac{1}{1+e^{-m}}.
  \]
  \begin{columns}[T]
    \begin{column}{0.51\linewidth}
      \begin{block}{Standard target gradient}
        $\tilde g_t=q_t-1=0$ after rounding.
      \end{block}
    \end{column}
    \begin{column}{0.47\linewidth}
      \begin{block}{Non-target gradients survive}
        $\tilde g_j\simeq e^{z_j-z_t}>0$.
      \end{block}
    \end{column}
  \end{columns}
  \medskip
  Precision cliffs in margin $m$: bf16 $\approx5.5$, fp32 $\approx16.6$, fp64 $\approx36.7$.
  \papersource{Paper, Sec. 2, Eq. (2)--(4), Fig. 1a, PDF pp. 1--2.}
\end{frame}

\begin{frame}{A three-class example exposes the missing component}
  \begin{block}{fp32, $z=(9,-9,-9)$, target $t=0$}
    $e^{-m}=2e^{-18}\approx 3.046\times 10^{-8}$: small near $1$, representable near $0$.
  \end{block}
  \begin{center}
    \begin{tabular}{lcc}
      \toprule
      & Loss & Gradient $(g_0,g_1,g_2)$ \\
      \midrule
      \std\ CE & $0$ & $(0,\;1.523,\;1.523)\times10^{-8}$ \\
      \fix\ CE & $3.046\times10^{-8}$ & $(-3.046,\;1.523,\;1.523)\times10^{-8}$ \\
      \bottomrule
    \end{tabular}
  \end{center}
  The standard result keeps both non-target entries but loses the target entry.
  \papersource{Paper, Sec. 2, numerical JAX example, PDF p. 3.}
\end{frame}

\begin{frame}{Fix hard-target CE by preserving the small mass}
  \[
    m=-\operatorname{LSE}_{j\ne t}(z_j-z_t),\qquad
    \boxed{\mathcal L(z,t)=\operatorname{softplus}(-m)}
  \]
  \[
    g_j=\frac{e^{z_j-z_t}}{1+e^{-m}}\ (j\ne t),\qquad
    \boxed{g_t=-\sum_{j\ne t}g_j}.
  \]
  \begin{itemize}
    \item Exclude the target from log-sum-exp; never form $q_t-1$.
    \item Autodiff through $z_j-z_t$ recovers the negative sum at the target.
  \end{itemize}
  \papersource{Paper, Sec. 2, Eq. (6)--(7), PDF p. 3.}
\end{frame}

\begin{frame}{The same construction fixes standalone log-softmax}
  Choose $i_{\max}=\arg\max_i z_i$ and define
  \[
    \tilde z_k=\mathbf 1[k\ne i_{\max}](z_k-z_{i_{\max}}),\quad
    m=-\operatorname{LSE}_{j\ne i_{\max}}\tilde z_j.
  \]
  \[
    \boxed{\operatorname{logsoftmax}(z)_k
    =\tilde z_k-\operatorname{softplus}(-m)}
  \]
  \begin{itemize}
    \item Replace the maximum slot by exact zero; exclude it from the LSE.
    \item Then $\operatorname{softmax}(z)=\exp(\operatorname{logsoftmax}(z))$.
    \item A rounded forward probability of $1$ can still have a useful gradient.
  \end{itemize}
  \papersource{Paper, Sec. 2, Eq. (5), JAX listing, PDF p. 3.}
\end{frame}

\begin{frame}{Cancellation loses information before subtraction}
  \[
    p=\operatorname{softmax}(z_T),\quad q=\operatorname{softmax}(z_S),\quad
    \nabla_{z_S}\mathcal L=q-p.
  \]
  \begin{itemize}
    \item Standard soft-target CE evaluates $p$ and $q$ separately, then subtracts.
    \item For $z_S\approx z_T$, rounding probabilities erases the small difference.
    \item Different forward and backward formulas can leave $q-p\ne0$ even at $z_S=z_T$.
  \end{itemize}
  \begin{block}{Better input to the gradient calculation}
    Use the directly available $\delta=z_S-z_T$.
  \end{block}
  \papersource{Paper, Sec. 3, Eq. (8), PDF pp. 3--4.}
\end{frame}

\begin{frame}{Compute $q-p$ from the logit gap}
  \[
    d_k=\delta_k-\sum_i p_i\delta_i,\qquad
    s_k=d_k-\operatorname{log1p}\!\left(\sum_i p_i\operatorname{expm1}(d_i)\right).
  \]
  \[
    s_k=\log(q_k/p_k),\qquad
    u_k=p_k\operatorname{expm1}(s_k),\qquad
    g_k=u_k-q_k\sum_i u_i.
  \]
  \begin{itemize}
    \item $\operatorname{expm1}(x)$ and $\operatorname{log1p}(x)$ retain accuracy near zero.
    \item Centering makes $\sum_i p_i d_i=0$; the final term reduces sum drift.
    \item At $\delta=0$, every intermediate difference vanishes.
  \end{itemize}
  \papersource{Paper, Sec. 3, Eq. (9), PDF p. 4.}
\end{frame}

\begin{frame}{The forward loss is unchanged; the backward path is replaced}
  \begin{block}{Implementation pattern in the paper}
    Compute the usual soft-target CE value, then attach $g$ as its custom gradient using \texttt{stop\_gradient}.
  \end{block}
  \begin{itemize}
    \item A small fp32 teacher--student perturbation produces a wrong sign in one standard gradient entry.
    \item The fixed gradient matches the fp64 reference in the paper's three-class example.
    \item At exactly equal logits, it returns zero and does not push the student away.
  \end{itemize}
  \papersource{Paper, Sec. 3, JAX listing and numerical example, PDF p. 4.}
\end{frame}

\begin{frame}{Collapse creates Slingshot through Adam's memory}
  \begin{enumerate}
    \item After memorization, margins cross the fp32 cliff; target gradients vanish.
    \item Adam's second moment $\hat v$ decays while parameter updates nearly stop.
    \item Some examples recross the cliff: restored gradients meet a small $\sqrt{\hat v}+\epsilon$.
    \item The resulting large step spikes the loss and resets $\hat v$.
  \end{enumerate}
  \begin{block}{Prediction}
    Keeping $g_t\approx-e^{-m}$ smooth should remove this numerical cycle.
  \end{block}
  \papersource{Paper, Sec. 4, PDF pp. 4--5. Mechanistic explanation proposed by the authors.}
\end{frame}

\begin{frame}{The reported Slingshot trajectories support that mechanism}
  \centering
  \includegraphics[width=0.95\linewidth,height=0.72\textheight,keepaspectratio]{figures/fig2-slingshot.png}
  \papersource{Paper, Fig. 2, PDF p. 5. Top: $a+b\bmod31$, $\lambda=1$; bottom: $a^2+b^2\bmod97$, $\lambda=0$. Original PDF figure crop.}
\end{frame}

\begin{frame}{The margin plateau comes from $\epsilon$ and weight decay}
  At an AdamW stationary point, $\hat m_i=g_i$ and $\sqrt{\hat v_i}=|g_i|$:
  \[
    \frac{g_i}{|g_i|+\epsilon}+\lambda\theta_i=0.
  \]
  \begin{itemize}
    \item After memorization, $|g_i|\propto e^{-m}$.
    \item Once $|g_i|\ll\epsilon$, the normalized update scales as $e^{-m}/\epsilon$.
    \item In the studied regime, balance gives $m^*(\epsilon)\approx\log(1/\epsilon)$.
  \end{itemize}
  For $\epsilon=10^{-8}$, $m^*\approx18.4$, just above the fp32 cliff $16.6$.
  \papersource{Paper, Sec. 5, Eq. (10)--(11), Fig. 3, PDF pp. 5--6. Approximation suppresses model-dependent factors.}
\end{frame}

\begin{frame}{The $\epsilon$ sweep separates equilibrium from a precision cliff}
  \centering
  \includegraphics[width=0.90\linewidth,height=0.65\textheight,keepaspectratio]{figures/fig4-epsilon.png}
  \begin{itemize}
    \item With $\lambda=1$, positive-$\epsilon$ margins level near $\log(1/\epsilon)$.
    \item With $\lambda=0$, margins keep growing past that level.
  \end{itemize}
  \papersource{Paper, Fig. 4, PDF p. 6. Original author figure; LaProp updates $\sqrt{\hat v}$ directly.}
\end{frame}

\begin{frame}{Main grokking experiment varies loss, precision, and $\epsilon$}
  \begin{center}
    \begin{tabular}{ll}
      \toprule
      Task & $a+b\bmod31$, 50\% training split \\
      Model & 1-layer width-64 ReLU Transformer \\
      Optimizer & full-batch AdamW, 10,000 epochs \\
      Shared settings & $\eta=10^{-3}$, $\beta=(0.9,0.98)$, $\lambda=1$ \\
      Variables & standard/fixed CE; bf16/fp32/fp64; $\epsilon=10^{-8},10^{-12},10^{-16}$ \\
      Outcomes & test accuracy (higher better); final logit margin \\
      \bottomrule
    \end{tabular}
  \end{center}
  \papersource{Paper, Sec. 6, Fig. 5 and Table 1, PDF pp. 6--7.}
\end{frame}

\begin{frame}{The precision cliff shifts with the number format}
  \centering
  \includegraphics[width=0.94\linewidth,height=0.72\textheight,keepaspectratio]{figures/fig5-precision.png}
  \papersource{Paper, Fig. 5, PDF p. 7. Top: $\epsilon=10^{-8}$; bottom: $10^{-16}$. Original PDF figure crop.}
\end{frame}

\begin{frame}{Fixed CE improves accuracy without changing the loss shape}
  \begin{center}
    \begin{tabular}{lcc}
      \toprule
      Test accuracy (5 seeds) & $\epsilon=10^{-8}$ & $\epsilon=10^{-16}$ \\
      \midrule
      fp32 fixed CE & $100.0\pm0.0$\% & $99.0\pm2.3$\% \\
      fp32 standard CE & $94.3\pm5.4$\% & $96.5\pm4.3$\% \\
      bf16 fixed CE & $100.0\pm0.1$\% & $98.9\pm2.4$\% \\
      bf16 standard CE & $93.3\pm7.4$\% & $94.2\pm8.3$\% \\
      fp64 fixed CE & $100.0\pm0.0$\% & $99.0\pm2.0$\% \\
      fp64 standard CE & $100.0\pm0.0$\% & $97.2\pm3.2$\% \\
      \bottomrule
    \end{tabular}
  \end{center}
  At $\epsilon=10^{-12}$ in fp32: zero-sum projection $48.9\pm2.6$\%, StableMax $95.6\pm2.0$\%, fixed CE $99.6\pm0.7$\%.
  \papersource{Paper, Table 1, PDF p. 7. Accuracy in percent; $a+b\bmod31$, $\lambda=1$.}
\end{frame}

\begin{frame}{The effect repeats on four modular tasks}
  \centering
  \includegraphics[width=0.95\linewidth,height=0.69\textheight,keepaspectratio]{figures/fig6-benchmarks.png}
  \begin{itemize}
    \item Fixed fp32 CE tracks standard fp64 CE; standard fp32 CE repeatedly spikes.
  \end{itemize}
  \papersource{Paper, Fig. 6, PDF p. 8. $\epsilon=10^{-12}$; 1-layer Transformers and 2-layer MLPs. Original PDF figure crop.}
\end{frame}

\begin{frame}{Identical student and teacher make numerical drift measurable}
  \begin{itemize}
    \item Train teacher to 100\% test accuracy on $a+b\bmod53$.
    \item Copy teacher weights to student: $\theta_S^{(0)}=\theta_T$.
    \item Distill for 2,000 Adam steps: $\eta=2\times10^{-3}$, $\lambda=0$, $\epsilon=10^{-8}$, temperature $T=4$.
    \item Measure held-out $\mathrm{KL}(p_T\Vert q_S)$ and test accuracy.
  \end{itemize}
  \begin{block}{The diagnostic}
    Since $z_S=z_T$ initially, the true student gradient is zero. Any move at initialization is numerical.
  \end{block}
  \papersource{Paper, Sec. 7, PDF p. 8. $T=4$ keeps scaled margins below the bf16 collapse cliff.}
\end{frame}

\begin{frame}{Fixed soft-target CE prevents teacher--student drift}
  \centering
  \includegraphics[width=0.94\linewidth,height=0.70\textheight,keepaspectratio]{figures/fig7-distillation.png}
  \begin{itemize}
    \item Standard bf16 CE drives KL upward and degrades accuracy.
    \item Fixed CE stays at KL $=0$ and 100\% accuracy in all three precisions.
  \end{itemize}
  \papersource{Paper, Fig. 7, PDF p. 8. Two of three seeds shown; original PDF figure crop.}
\end{frame}

\begin{frame}{Take-home messages and evidence boundary}
  \begin{enumerate}
    \item Rewrite algebra before precision is lost: $\operatorname{softplus}$ for high confidence, $\operatorname{expm1}/\operatorname{log1p}$ for near equality.
    \item This removes reported Slingshot spikes and distillation drift in the tested modular tasks.
    \item The $m\approx\log(1/\epsilon)$ plateau comes from AdamW plus weight decay, not fp32's representable range.
  \end{enumerate}
  \begin{block}{Scope}
    The paper tests controlled small tasks and an exactly matched student--teacher initialization; it does not measure large-scale training speed or benefit.
  \end{block}
  \papersource{Paper, Secs. 4--8, PDF pp. 4--9.}
\end{frame}

\begin{frame}{Paper and implementation sources}
  \begin{itemize}
    \item Barron, Gnanapragasam, Tschannen. \emph{A Fix for Softmax and Cross-Entropy}. \href{https://arxiv.org/abs/2610.31642v1}{arXiv:2610.31642v1}, 2026.
    \item \href{https://jonbarron.info/softmaxfix/}{Author project page}: paper text, JAX function examples, interactive demo.
    \item \href{https://jonbarron.info/}{Author listing}: labeled arXiv, 2026 as of 27 Oct 2026.
  \end{itemize}
  \medskip
  \begin{block}{Numerical implementation detail}
    The current project-page soft-target snippet includes a fallback for large log-ratios to avoid $0\cdot\infty$; check the intended dynamic range when copying code.
  \end{block}
\end{frame}

\end{document}