[tab] /Users/franzsrambical/Downloads/arXiv-2410.01104v3/example_paper.tex (latex)
%%%%%%%% ICML 2025 EXAMPLE LATEX SUBMISSION FILE %%%%%%%%%%%%%%%%%

\documentclass{article}
\input{math_commands.tex}
% Recommended, but optional, packages for figures and better typesetting:
\usepackage[draft]{minted}
\usepackage{colortbl}
\usepackage{hyperref}
\usepackage{url}
\usepackage{microtype}
\usepackage{graphicx}
\usepackage{subfigure}
\usepackage{wrapfig}
\usepackage{multicol}
\usepackage{booktabs} % for professional tables
\usepackage{xcolor}

\usemintedstyle{default}

% hyperref makes hyperlinks in the resulting PDF.
% If your build breaks (sometimes temporarily if a hyperlink spans a page)
% please comment out the following usepackage line and replace
% \usepackage{icml2025} with \usepackage[nohyperref]{icml2025} above.

\definecolor{olivegreen}{rgb}{0,0.6,0}
\definecolor{mymauve}{HTML}{E60B42}
%\definecolor{camdrk}{RGB}{0,62,114}
\definecolor{echopurple}{HTML}{9400D1}
\definecolor{camdrk}{HTML}{0099CC}
\definecolor{darkgry}{HTML}{333333}
\definecolor{echoblue}{HTML}{0099CC}
\definecolor{echonavy}{HTML}{0054B2}
\definecolor{mygold}{HTML}{B8860B}


% Attempt to make hyperref and algorithmic work together better:

ewcommand{	heHalgorithm}{\arabic{algorithm}}

% Use the following line for the initial blind version submitted for review:
 %\usepackage{icml2025}

% If accepted, instead use the following line for the camera-ready submission:
\usepackage[accepted]{icml2025}

% For theorems and such
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
\usepackage{inconsolata}

% if you use cleveref..
\usepackage[capitalize,noabbrev]{cleveref}


\makeatletter
\def\PYG@reset{\let\PYG@it=
elax \let\PYG@bf=
elax%
    \let\PYG@ul=
elax \let\PYG@tc=
elax%
    \let\PYG@bc=
elax \let\PYG@ff=
elax}
\def\PYG@tok#1{\csname PYG@tok@#1\endcsname}
\def\PYG@toks#1+{\ifx
elax#1\empty\else%
    \PYG@tok{#1}\expandafter\PYG@toks\fi}
\def\PYG@do#1{\PYG@bc{\PYG@tc{\PYG@ul{%
    \PYG@it{\PYG@bf{\PYG@ff{#1}}}}}}}
\def\PYG#1#2{\PYG@reset\PYG@toks#1+
elax+\PYG@do{#2}}

\@namedef{PYG@tok@w}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.73,0.73}{##1}}}
\@namedef{PYG@tok@c}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cp}{\def\PYG@tc##1{	extcolor[rgb]{0.61,0.40,0.00}{##1}}}
\@namedef{PYG@tok@k}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kp}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kt}{\def\PYG@tc##1{	extcolor[rgb]{0.69,0.00,0.25}{##1}}}
\@namedef{PYG@tok@o}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@ow}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.67,0.13,1.00}{##1}}}
\@namedef{PYG@tok@nb}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@nf}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@nc}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@nn}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@ne}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.80,0.25,0.22}{##1}}}
\@namedef{PYG@tok@nv}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@no}{\def\PYG@tc##1{	extcolor[rgb]{0.53,0.00,0.00}{##1}}}
\@namedef{PYG@tok@nl}{\def\PYG@tc##1{	extcolor[rgb]{0.46,0.46,0.00}{##1}}}
\@namedef{PYG@tok@ni}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.44,0.44,0.44}{##1}}}
\@namedef{PYG@tok@na}{\def\PYG@tc##1{	extcolor[rgb]{0.41,0.47,0.13}{##1}}}
\@namedef{PYG@tok@nt}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@nd}{\def\PYG@tc##1{	extcolor[rgb]{0.67,0.13,1.00}{##1}}}
\@namedef{PYG@tok@s}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sd}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@si}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.64,0.35,0.47}{##1}}}
\@namedef{PYG@tok@se}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.67,0.36,0.12}{##1}}}
\@namedef{PYG@tok@sr}{\def\PYG@tc##1{	extcolor[rgb]{0.64,0.35,0.47}{##1}}}
\@namedef{PYG@tok@ss}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@sx}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@m}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@gh}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,0.50}{##1}}}
\@namedef{PYG@tok@gu}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.50,0.00,0.50}{##1}}}
\@namedef{PYG@tok@gd}{\def\PYG@tc##1{	extcolor[rgb]{0.63,0.00,0.00}{##1}}}
\@namedef{PYG@tok@gi}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.52,0.00}{##1}}}
\@namedef{PYG@tok@gr}{\def\PYG@tc##1{	extcolor[rgb]{0.89,0.00,0.00}{##1}}}
\@namedef{PYG@tok@ge}{\let\PYG@it=	extit}
\@namedef{PYG@tok@gs}{\let\PYG@bf=	extbf}
\@namedef{PYG@tok@gp}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,0.50}{##1}}}
\@namedef{PYG@tok@go}{\def\PYG@tc##1{	extcolor[rgb]{0.44,0.44,0.44}{##1}}}
\@namedef{PYG@tok@gt}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.27,0.87}{##1}}}
\@namedef{PYG@tok@err}{\def\PYG@bc##1{{\setlength{\fboxsep}{\string -\fboxrule}\fcolorbox[rgb]{1.00,0.00,0.00}{1,1,1}{\strut ##1}}}}
\@namedef{PYG@tok@kc}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kd}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kn}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@kr}{\let\PYG@bf=	extbf\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@bp}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.50,0.00}{##1}}}
\@namedef{PYG@tok@fm}{\def\PYG@tc##1{	extcolor[rgb]{0.00,0.00,1.00}{##1}}}
\@namedef{PYG@tok@vc}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@vg}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@vi}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@vm}{\def\PYG@tc##1{	extcolor[rgb]{0.10,0.09,0.49}{##1}}}
\@namedef{PYG@tok@sa}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sb}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sc}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@dl}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@s2}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@sh}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@s1}{\def\PYG@tc##1{	extcolor[rgb]{0.73,0.13,0.13}{##1}}}
\@namedef{PYG@tok@mb}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mf}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mh}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mi}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@il}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@mo}{\def\PYG@tc##1{	extcolor[rgb]{0.40,0.40,0.40}{##1}}}
\@namedef{PYG@tok@ch}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cm}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cpf}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@c1}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}
\@namedef{PYG@tok@cs}{\let\PYG@it=	extit\def\PYG@tc##1{	extcolor[rgb]{0.24,0.48,0.48}{##1}}}

\def\PYGZbs{\char`\\}
\def\PYGZus{\char`\_}
\def\PYGZob{\char`\{}
\def\PYGZcb{\char`\}}
\def\PYGZca{\char`\^}
\def\PYGZam{\char`\&}
\def\PYGZlt{\char`\<}
\def\PYGZgt{\char`\>}
\def\PYGZsh{\char`\#}
\def\PYGZpc{\char`\%}
\def\PYGZdl{\char`\$}
\def\PYGZhy{\char`\-}
\def\PYGZsq{\char`\'}
\def\PYGZdq{\char`\"}
\def\PYGZti{\char`\~}
% for compatibility with earlier versions
\def\PYGZat{@}
\def\PYGZlb{[}
\def\PYGZrb{]}
\makeatother

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% THEOREMS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
	heoremstyle{plain}

ewtheorem{theorem}{Theorem}[section]

ewtheorem{proposition}[theorem]{Proposition}

ewtheorem{lemma}[theorem]{Lemma}

ewtheorem{corollary}[theorem]{Corollary}
	heoremstyle{definition}

ewtheorem{definition}[theorem]{Definition}

ewtheorem{assumption}[theorem]{Assumption}
	heoremstyle{remark}

ewtheorem{remark}[theorem]{Remark}

% Todonotes is useful during development; simply uncomment the next line
%    and comment out the line below the next line to turn off comments
%\usepackage[disable,textsize=tiny]{todonotes}
\usepackage[textsize=tiny]{todonotes}


% The \icmltitle you define below is probably too long as a header.
% Therefore, a short form for the running title is supplied here:
\icmltitlerunning{	exttt{Softmax} is not Enough (for Sharp Size Generalisation)}

\begin{document}

	wocolumn[
\icmltitle{	exttt{Softmax} is not Enough (for Sharp Size Generalisation)}

% It is OKAY to include author information, even for blind
% submissions: the style file will automatically remove it for you
% unless you've provided the [accepted] option to the icml2025
% package.

% List of affiliations: The first argument should be a (short)
% identifier you will use later to specify author affiliations
% Academic affiliations should list Department, University, City, Region, Country
% Industry affiliations should list Company, City, Region, Country

% You can specify symbols, otherwise they are numbered in order.
% Ideally, you should not use this facility. Affiliations will be numbered
% in order of appearance and this is the preferred way.
\icmlsetsymbol{equal}{*}

\begin{icmlauthorlist}
\icmlauthor{Petar Veli\v{c}kovi\'{c}}{gdm}
\icmlauthor{Christos Perivolaropoulos}{gdm}
\icmlauthor{Federico Barbero}{oxf,equal}
\icmlauthor{Razvan Pascanu}{gdm}
\end{icmlauthorlist}

\icmlaffiliation{gdm}{Google DeepMind}
\icmlaffiliation{oxf}{University of Oxford}

\icmlcorrespondingauthor{Petar Veli\v{c}kovi\'{c}}{petarv@google.com}

% You may provide any keywords that you
% find helpful for describing your paper; these are used to populate
% the "keywords" metadata in the PDF but will not be shown in the document
\icmlkeywords{Machine Learning, ICML}

\vskip 0.3in
]

% this must go after the closing bracket ] following 	wocolumn[ ...

% This command actually creates the footnote in the first column
% listing the affiliations and the copyright notice.
% The command takes one argument, which is text to display at the start of the footnote.
% The \icmlEqualContribution command is standard text for equal contribution.
% Remove it (just {}) if you do not need this facility.

\printAffiliationsAndNotice{$^\ast$Work performed while the author was at Google DeepMind.}  % leave blank if no need to mention equal contribution
%\printAffiliationsAndNotice{\icmlEqualContribution} % otherwise use the standard text.

\begin{abstract}
A key property of reasoning systems is the ability to make \emph{sharp} decisions on their input data. For contemporary AI systems, a key carrier of sharp behaviour is the 	exttt{softmax} function, with its capability to perform differentiable query-key lookups. It is a common belief that the predictive power of networks leveraging 	exttt{softmax} arises from ``circuits'' which sharply perform certain kinds of computations consistently across many diverse inputs. However, for these circuits to be robust, they would need to generalise well to \emph{arbitrary} valid inputs. In this paper, we dispel this myth: even for tasks as simple as finding the maximum key, any learned circuitry \emph{must disperse} as the number of items grows at test time. We attribute this to a fundamental limitation of the 	exttt{softmax} function to robustly approximate sharp functions with increasing problem size, prove this phenomenon theoretically, and propose \emph{adaptive temperature} as an ad-hoc technique for improving the sharpness of 	exttt{softmax} at inference time.
\end{abstract}

\section{Motivation}\setcounter{footnote}{0}

It is no understatement to say that the $\mathtt{softmax}_	heta : \mathbb{R}^n
ightarrow[0, 1]^n$ function\footnote{Strictly speaking, the proper name for this function should be 	exttt{soft	extbf{arg}max}. We choose to retain the terminology introduced by \citet{bridle1989training}, primarily for reasons of alignment with modern deep learning frameworks.}:
\begin{equation}
    \mathtt{softmax}_	heta(\mathbf{e}) = \begin{bmatrix} \dfrac{\exp (e_1 / 	heta)}{\sum_{k} \exp (e_k / 	heta)} & \dots & \dfrac{\exp (e_n / 	heta)}{\sum_{k} \exp (e_k / 	heta)} \end{bmatrix}
\end{equation}
is one of the most fundamental functions in contemporary artificial intelligence systems.

\begin{figure}
    \centering
    \includegraphics[width=\linewidth]{dispersion.pdf}
    \caption{Illustration of Theorem 
ef{thm:disperse}, one of our key results. Assuming a tokenised input from a fixed vocabulary and a non-zero temperature, for every 	exttt{softmax} attention head inside an architecture comprising only 	extcolor{teal}{\bf MLPs} and 	extcolor{mygold}{\bf 	exttt{softmax} self-attention layers}, it must hold that, given sufficiently many 	extcolor{echonavy}{\bf tokens}, its 	extcolor{olivegreen}{\bf attention coefficients} will \emph{disperse}, even if they were sharp in-distribution.}
    \label{fig:disperse}
\end{figure}

The role of 	exttt{softmax} in deep learning is to convert any vector of \emph{logits}, $\mathbf{e}\in\mathbb{R}^n$, into a \emph{probability distribution}, in a form that is part of the \emph{exponential} family. Further, 	exttt{softmax} allows for application of a \emph{temperature} parameter, $	heta\in\mathbb{R}$, to adjust the amount of probability mass attached to the highest logit---a concept borrowed from the Boltzmann distribution in statistical mechanics.

Initially, the primary utilisation of 	exttt{softmax} in deep learning was within the final layer of \emph{classifiers}. Its influence in this domain vastly expanded after it saw use in the \emph{internal} layers---as a differentiable key-value store \citep{graves2014neural} or a mechanism for \emph{attending} over the most relevant parts of the input \citep{bahdanau2014neural}. This \emph{attentional} framing of 	exttt{softmax} was critical in defining important models for sequences \citep[Transformers]{vaswani2017attention}, images \citep[ViTs]{dosovitskiy2021an} and graphs \citep[GATs]{velickovic2018graph}.

Several efforts attribute the success of 	exttt{softmax} to its capability of modelling computations relevant to reasoning. This can be related to the concept of \emph{circuits} in theoretical computer science \citep{arora2009computational}. Several interpretable pieces of ``circuitry'' \citep{olah2020zoom} have already been discovered in large Transformers, primarily under the umbrella of \emph{mechanistic interpretability} \citep{elhage2021mathematical,olsson2022context,wang2022interpretability}.

Here we study the robustness of such circuitry, especially when going beyond the distribution the models are trained on---a critical regime for \emph{reasoning engines}. We find that, in spite of its many successes, 	exttt{softmax} \emph{is fundamentally unable} to robustly generalise such circuits out of distribution, especially as it provably cannot approximate 	extbf{sharpness} with increasing problem size (Figure 
ef{fig:disperse}). 

Here we call a function taking a variable number of inputs \emph{sharp} if its output value can be expressed using only a \emph{constant} number of these inputs. For example, 	exttt{max} is sharp, as its output value is equal to exactly one of its inputs' values\footnote{In practice, as 	exttt{softmax} coefficients that are \emph{exactly} zero or one are often not easily obtainable, we may consider an item to be \emph{expressed} in the final answer if its coefficient is greater than some $\epsilon > 0$ (which may depend on the problem size). That said, the exact choice of $\epsilon$ does not affect the essence of our results, as we are able to bound the coefficients both above and below.}. The average function is not sharp, as its output value depends on all of its input values (with factor $1/n$ for each of the $n$ items).

\paragraph{Key Theoretical Result} 
We define sharp functions by their behaviour as their number of inputs varies. This directly motivates the \emph{size generalisation} setting we study: generalising to different amounts of inputs. Specifically, when we analyse neural networks that learn sharp functions, we assume that they are trained on problem instances containing no more than $n$ input items, and we take a particular interest in their sharpness on instances with $n' > n$ items; these are considered \emph{out-of-distribution} instances because they go beyond the maximal number of inputs the model had been prepared for. In language modelling, this setting is also known as \emph{length generalisation} \citep{anil2022exploring}; in graph machine learning, it is known as \emph{size generalisation} \citep{yehudai2021local}.

Through one of our key theoretical results (Theorem 
ef{thm:disperse}), we demonstrate that modern deep learning architectures, operating over a fixed vocabulary of input tokens and leveraging the 	exttt{softmax} function, are fundamentally incapable of learning functions that remain sharp under such larger instances. This is due to the fact that the coefficients emitted by the 	exttt{softmax} function must \emph{disperse} as we increase the number of input items. Here by dispersing we mean that, as the number of input items grows, the coefficient attached to each individual item must decay towards zero. This makes it impossible to robustly compute functions that depend on any particular finite amount of input values, such as the aforementioned 	exttt{max}, as we show in Appendix 
ef{app:corollary} (Corollary 
ef{cor:failure} and Remark 
ef{rem:alltrans}).

We hope that our results will encourage future study of alternative attentional functions, in light of the problems we identify, especially for building reasoning engines of the future. That being said, we also believe our findings indicate ways to modify the 	exttt{softmax} function to support sharpness for longer---as one simple example, we propose an \emph{adaptive temperature} mechanism for 	exttt{softmax}.

%\section{Background}
\paragraph{Background} The analysis of attentional coefficients and attempting to attribute interpretable operations to them dates back to the earliest deployments of internal 	exttt{softmax} layers at scale; examples include \citet[Figure 6]{graves2014neural}, \citet[Figure 3]{bahdanau2014neural}, \citet[Figures 3--5]{vaswani2017attention} and \citet[Figure 5]{qiu2018deepinf}. A strong current topic in this space analyses the self-attentional heads of Transformers \citep{voita2019analyzing,jain2019attention}.

With the rise of large language models, mechanistic interpretability has taken charge in detecting and elucidating various circuits in Transformers \citep{elhage2021mathematical}. Some prominent discoveries include induction heads \citep{olsson2022context}, indirect object identification \citep{wang2022interpretability}, multiple-choice heads \citep{lieberum2023does}, successor heads \citep{gould2023successor}, attentional sinks \citep{darcet2023vision}, comparator heads \citep{hanna2024does} and retrieval heads \citep{wu2024retrieval}. Most recently, these efforts have relied on sparse autoencoders \citep{kissane2024interpreting}.

The skills above are quite impressive and span many rules one might hope a robust reasoning system would have, and the discovered heads always appear sharp when inspected on \emph{in-distribution} samples. However, it is also known that many \emph{easy} tasks requiring sharp attention---such as finding minima---are hard to do reliably with LLMs \emph{out-of-distribution}, particularly on larger inputs \citep[Figure 6]{markeeva2024clrs}. More challenging sharp order statistic tasks, such as finding the second minimum \citep{ong2022learnable} may even be hard to learn in-distribution. The discrepancy of such results with the previous paragraph motivate our study, and formalisation of 	exttt{softmax} dispersion.

Certain dispersion effects in 	exttt{softmax}---e.g. as an effect of increasing temperature---are already well-understood in thermodynamics. A core contribution of our work is understanding dispersion in a setting where the 	extbf{amount of logits can vary}, which is relevant for generalisation in Transformers. We are not the first to observe dispersion in this setting empirically; prior works studying the capability of Transformers to execute algorithms \citep{yan2020neural} and perform random-access lookups \citep{ebrahimi2024your} also note dispersion patterns. Our work is the first to rigorously prove these effects, directly attribute them to the 	exttt{softmax} operator, as well as propose ways to improve sharpness empirically within 	exttt{softmax}. The proof technique we will use to demonstrate this is inspired by \citet{barbero2024transformers}, though unlike their work, our key results apply regardless of whether the computational graph is bottlenecked or not.

\paragraph{Primer on Attentional Heads and Transformers}
Within this paper we will primarily study the use of 	exttt{softmax} within \emph{self-attentional} neural network architectures, such as Transformers \citep{vaswani2017attention}. The core building block of such models is the (dot-product) \emph{attentional head}, which operates over a collection of $n$ nodes (or tokens), with features $\mathbf{x}_i^{(n)}\in\mathbb{R}^k$ for node $1\leq i\leq n$, for a given query vector $	ilde{\mathbf{q}}^{(n)}\in\mathbb{R}^k$.

First, the attentional head computes \emph{key}, (updated) \emph{query} and \emph{value} vectors via matrix multiplication:
\begin{equation}
    \mathbf{k}_i^{(n)} = \mathbf{K}\mathbf{x}_i^{(n)} \qquad \mathbf{q}^{(n)} = \mathbf{Q}	ilde{\mathbf{q}}^{(n)} \qquad  \mathbf{v}_i^{(n)} = \mathbf{V}\mathbf{x}_i^{(n)}
\end{equation}
where $\mathbf{K}, \mathbf{Q}, \mathbf{V}\in\mathbb{R}^{k'	imes k}$ are learnable parameter matrices. Then, dot-products between the query and all of the key vectors are taken to compute unnormalised attentional coefficients of each item, also known as \emph{logits}, $e_i^{(n)}\in\mathbb{R}$. These coefficients are normalised using the 	exttt{softmax} function to obtain \emph{attentional coefficients}, $\alpha_{i}^{(n)}\in\mathbb{R}$. Finally, the attentional coefficients are used for a weighted sum of value vectors, which represents the output of the attentional head, $\mathbf{y}^{(n)}\in\mathbb{R}^{k'}$:
\begin{equation}\label{eqn:outatt}
    e_{i}^{(n)} = \left(\mathbf{q}^{(n)}
ight)^	op\mathbf{k}_i^{(n)}, \qquad \alpha_{i}^{(n)} = \mathtt{softmax}_	heta(\mathbf{e}^{(n)})_i
\end{equation}
\begin{equation*}
    \mathbf{y}^{(n)} = \sum_{1\leq i\leq n} \alpha_i^{(n)}\mathbf{v}_i^{(n)}
\end{equation*}
With regard to how attentional heads are used within Transformers, we will mainly analyse two of the most popular strategies: BERT-style \citep{devlin-etal-2019-bert} and GPT-style \citep{radford_improving_2018}. In both cases, each of the input nodes computes its own attentional output, i.e. there is one query vector per node, computed as $\mathbf{q}_i^{(n)} = \mathbf{Q}\mathbf{x}_i^{(n)}$, leading to per-node attention coefficients $\alpha_{ij}$ and outputs $\mathbf{y}_i^{(n)}$ by distributing Equation 
ef{eqn:outatt} across queries. The main difference is in the choice of keys.

In BERT-style self-attention, each node's query vector attends over all of the key vectors, i.e. it is obtained by directly distributing Equation 
ef{eqn:outatt} across all queries:
\begin{equation}
    e_{ij}^{(n)} = \left(\mathbf{q}_i^{(n)}
ight)^	op\mathbf{k}_j^{(n)}, \qquad \alpha_{ij}^{(n)} = \mathtt{softmax}_	heta(\mathbf{e}_i^{(n)})_j
\end{equation}
\begin{equation*}
    \mathbf{y}_i^{(n)} = \sum_{1\leq j\leq n} \alpha_{ij}^{(n)}\mathbf{v}_j^{(n)}
\end{equation*}
In comparison, GPT-style attention (also known as ``causal masking'' or the decoder-only Transformer) only allows information to flow \emph{forwards}; each node's query vector may only attend to the key vectors from nodes that precede it. This yields the following modification:
\begin{equation}
    e_{ij}^{(n)} = \begin{cases} \left(\mathbf{q}_i^{(n)}
ight)^	op\mathbf{k}_j^{(n)} & j \leq i\\
    -\infty & j > i\end{cases}\quad \alpha_{ij}^{(n)} = \mathtt{softmax}_	heta(\mathbf{e}_i^{(n)})_j
\end{equation}
\begin{equation*}
\mathbf{y}_i^{(n)} = \sum_{1\leq j\leq i} \alpha_{ij}^{(n)}\mathbf{v}_j^{(n)}
\end{equation*}
Our key dispersion results hold for both styles of attention---this is mainly due to the fact that all predictions made by GPT-style architectures are dependent on the \emph{final} token embedding, $\mathbf{y}_n^{(n)}$, which will attend over all items, much like any BERT head. The main difference between the two will be in qualitative effects on certain corollaries of the theory (Appendices 
ef{app:corollary}--
ef{app:depthd}).

%\begin{itemize}
%    \item Discuss hard attention mechanisms
%    \item Issues with learning for softmax as another potential problem that is not being explored in this work 
%\end{itemize}
%\begin{itemize}
%    \item Gradient flow analysis relevant for its ``opposite direction'' conclusions in SPARC \citep{bica2024improving}.
%\end{itemize}

%Make a comment about \emph{local attention} in light of this result.

%What truly \emph{fixes} this? Sigmoid attention, or any kind of unnormalised attention, should work in principle \citep{ramapuram2024theory}. Hard attention is also fair game. Even linear attention \citep{schmidhuber1992learning}. The result about singular values presumes dot-product attention is used to compute coefficients. Using more expressive attention such as GATv2 \citep{brody2022how} can alleviate this.

\section{Dispersion in 	exttt{softmax} and Transformers}

\begin{figure}
    \centering
    \includegraphics[width=\linewidth]{attentions.png}
    \caption{Visualising the attentional head for the 	exttt{max} retrieval task for a batch of $32$ randomly-sampled input sets (each represented by one of the rows), over the $16$ items with largest key (columns). If the head operates correctly, it must allocate sharp attention to the rightmost item. From left to right, in each frame we \emph{double} the number of items the head has to process (starting from $16$ items).}
    \label{fig:soft_disp}
\end{figure}

\begin{figure}
  \begin{center}
    \includegraphics[width=0.49	extwidth]{entropies.png}
  \end{center}
  \caption{Entropy of attention heads in the first block of Gemma 2B with prompt 	exttt{"What is the maximum in the following sequence: \{seq\}? The maximum is:"} and varying the number of elements in 	exttt{seq}. Each curve is one attentional head; the blue shaded curve is the mean and standard deviation across all of them.}
  \label{fig:ent_gemma}
\end{figure}
To motivate our theory, we train a simple architecture including a single attentional head to predict a feature of the \emph{maximum} item in a set. Each item's features are processed with a deep MLP before attending, and the output vector of the attention is passed to a deep MLP predictor (see Appendix 
ef{app:experiment} for experimental details). We train this model using sets of $\leq 16$ items, and in Figure 
ef{fig:soft_disp} we visualise the head's attentional coefficients, computed over sets of varying size at inference time.

While the model indeed attributes focus sharply and cleanly on the maximum item, this only holds true on the problem sizes that the model was trained on. As we simulate a generalisation setting where the problem size increases (without changing the value distribution), the attentional coefficients eventually disperse towards the uniform distribution.

This effects manifests in the attention heads of Transformers as well---we visualise the \emph{entropy} (a proxy for sharpness) of Gemma 2B \citep{team2024gemma}'s heads when answering a similar maximisation task in Figure 
ef{fig:ent_gemma}.

In fact, we can show that this effect is \emph{inevitable} in 	exttt{softmax} using the following Lemma:
\begin{lemma}[	exttt{softmax} must disperse]{\label{lem:disperse}Let $\mathbf{e}^{(n)}\in\mathbb{R}^{n}$ be a collection of $n$ logits going into the $\mathtt{softmax}_	heta$ function with temperature $	heta > 0$, bounded above and below s.t. $m\leq e^{(n)}_k\leq M$ for some choice of constants $m, M\in\mathbb{R}$. Then, as more items are added ($n
ightarrow +\infty$), it must hold that, for each item $1\leq k\leq n$, $\mathtt{softmax}_	heta(\mathbf{e}^{(n)})_k=\Theta(\frac{1}{n})$.}
\end{lemma}
The computed attention coefficients hence must 	extbf{disperse} for all items with increasing problem size.
\begin{proof}

Let us denote the attentional coefficient assigned to $k$ by $\alpha^{(n)}_k = 	exttt{softmax}_	heta(\mathbf{e}^{(n)})_k\in [0, 1]$. Then we can bound $\alpha^{(n)}_k$ above as:
\begin{equation}
    %\alpha^{(n)}_k = 
    \frac{\exp (e^{(n)}_k/	heta)}{\sum_l \exp (e^{(n)}_l/	heta)} \leq \frac{\exp (M/	heta)}{n\exp (m/	heta)} = \frac{1}{n}\exp \left(\frac{M-m}{	heta}
ight)
\end{equation}
Similarly, we can bound $\alpha^{(n)}_k$ below as:
\begin{equation}
    %\alpha^{(n)}_k = 
    \frac{\exp (e^{(n)}_k/	heta)}{\sum_l \exp (e^{(n)}_l/	heta)} \geq \frac{\exp (m/	heta)}{n\exp (M/	heta)} = \frac{1}{n}\exp \left(\frac{m-M}{	heta}
ight)
\end{equation}
Hence, if we let $\delta = (M - m)$
\begin{equation}
   \frac{1}{n}\exp -\frac{\delta}{	heta}\leq \alpha^{(n)}_k \leq \frac{1}{n}\exp\frac{\delta}{	heta}
\end{equation}
Which implies $\alpha_k^{(n)}=\Theta(\frac{1}{n})$ as $\delta$ and $	heta$ are both constants.
\end{proof}
Lemma 
ef{lem:disperse} relies on bounding logits. The difference of these bounds (the \emph{spread}, $\delta = \max_{i} e^{(n)}_i - \min_j e^{(n)}_j$) controls the rate of dispersion. In modern Transformer LLM architectures operating over a vocabulary of possible token values, we can actually bound the logits in every single attentional layer---implying that dispersion \emph{must} happen everywhere in a Transformer for sufficient problem sizes:
\begin{theorem}[	exttt{softmax} in Transformers over vocabularies must disperse]
\label{thm:disperse}
Let $\mathcal{X}\subset\mathbb{R}^m$ be a set of possible $m$-dimensional input features, and let $\mathbf{X}^{(n)}\in\mathcal{X}^n$ be a matrix of input features for $n$ items. Further, assume that input features come from a 	extbf{finite} set of possible values, i.e. $|\mathcal{X}| < |\mathbb{N}|$. Let $e_j^{(n)} = (\mathbf{q}^{(n)})^	op\mathbf{k}_j^{(n)}$ where $\mathbf{q}^{(n)} = \phi(\mathbf{x}^{(n)}_1,\dots,\mathbf{x}^{(n)}_n)$ and $\mathbf{K}^{(n)} = \kappa(\mathbf{x}^{(n)}_1,\dots,\mathbf{x}^{(n)}_n)$, where $\phi : \mathcal{X}^n
ightarrow\mathbb{R}^k$ and $\kappa : \mathcal{X}^n
ightarrow\mathbb{R}^{n	imes k}$ are continuous functions, each expressible as a composition of $L$ layers $g_L \circ f_L \circ \dots \circ g_1 \circ f_1$ where each layer contains a feedforward component $f_i(\mathbf{z}_1, \dots, \mathbf{z}_n)_k = f_i(\mathbf{z}_k)$ or a self-attentional component $g_i(\mathbf{z}_1, \dots, \mathbf{z}_n)_k = \sum_{1\leq l\leq n}\alpha_{lk} v_i(\mathbf{z}_l)$ where $\alpha_{lk}\in [0, 1]$ are 	exttt{softmax}-normalised attention coefficients and $v_i$ is a feedforward network.
Then, for any $	heta > 0$ and $\epsilon > 0$, there must exist an $n\in\mathbb{N}$ such that $\mathtt{softmax}_	heta(\mathbf{e}^{(n)})_k < \epsilon$ for all $1\leq k\leq n$.
\end{theorem}
That is, attention coefficients must 	extbf{disperse} in all global Transformer heads if the input vocabulary is finite.
\begin{proof}
    Firstly, note that since $\mathcal{X}$ is a finite set of $m$-dimensional vectors, then it is also part of a \emph{compact} space spanning all convex combinations of those vectors. Then, all feedforward layers, $f_i$ and $v_i$, being continuous functions, move inputs from a compact set to another compact set. Similarly, every self-attentional layer, $g_i$, computes a convex combination of the outputs of $v_i$, and as such, if outputs of $v_i$ are on a compact space, the outputs of $g_i$ remain on the same compact space. Therefore, if the input space of $\phi$ and $\kappa$ is compact, then the output space of $\phi$ and (each row of) $\kappa$ on $\mathbb{R}^k$ must be compact as well,  regardless of the choice of $n$. Further, the dot product of two vectors $(\mathbf{q}^{(n)})^	op\mathbf{k}_j^{(n)}$ coming from compact spaces must be compact as well. Hence, the logits must be bounded by $m \leq e^{(n)}_k \leq M$ for constant $m$ and $M$. Then, letting $\delta=M-m$, we know (Lemma 
ef{lem:disperse}) that $\mathtt{softmax}_	heta(\mathbf{e}^{(n)})_k\leq \frac{1}{n}\exp\left(\delta/	heta
ight)$, so for all $n > \dfrac{\exp\left(\delta/	heta
ight)}{\epsilon}$ this value will be below $\epsilon$.
\end{proof}
It might seem intuitive that attention head dispersion is a potentially destructive event, which forces the Transformer into misclassifying certain inputs. We prove this intuition in Appendix 
ef{app:corollary}. We also discuss the rate at which dispersion occurs at various model depths in Appendix 
ef{app:depthd}.

\paragraph{Practical Values of $\delta$}
It is evident that the logit spread, $\delta$, is the key controller of dispersion in our results. If we were to na\"{i}vely apply our theory (to the maximal and minimal values representable by the types used to represent floating point numbers such as 	exttt{bfloat16}), it would result in fairly loose bounds which are not particularly informative.

However, it can easily be observed that \emph{empirical} values of $\delta$ do not reach those theoretical limits. For example, when feeding Gemma 2B and 7B models with an entire code sample of Gemma's Modules file\footnote{\url{https://github.com/google-deepmind/gemma/blob/main/gemma/modules.py}} ($\sim4,000$ tokens), the empirically observed values of $\delta$ across all of their attention heads are substantially more contained:
\begin{itemize}
    \item Gemma 2B: $\delta\in [2.28, 14.78]$ (average $5.69\pm 2.05$);
    \item Gemma 7B: $\delta\in [0.09, 32.74]$ (average $5.82\pm 2.61$).
\end{itemize}
There are several possible reasons why practical pre-trained models keep logit spreads contained; for example, any label noise prevents the model from learning arbitrarily large parameter values, or the derivative needed to make off-target logits arbitrarily negative may vanish when the 	exttt{softmax} output is already sufficiently sharp.


\section{Adaptive Temperature}

Since we now know dispersion is inevitable, are there any ways we can leverage our theory's findings to make 	exttt{softmax} sharper? One obvious constraint our theory rests on is the assumption that $	heta > 0$, i.e. that our temperature is nonzero. While zero temperature---also known as \emph{hard attention} \citep{denil2012learning,ranzato2014learning,mnih2014recurrent,pmlr-v37-xuc15}---guarantees sharpness, training large-scale Transformers with it tends to not work well in practice \citep{bica2024improving}.

What about applying $	heta=0$ to an \emph{already-trained} Transformer? We can show this is also problematic since, for any attention head where the Transformer has learnt to induce sharpness, it \emph{necessarily} did so by increasing magnitude of its weights (see Appendix 
ef{app:proposition} for a proof and numerical validation):
\begin{figure*}
    \centering
    \begin{Verbatim}[commandchars=\\\{\}]
\PYG{k}{def} \PYG{n+nf}{adaptive\PYGZus{}temperature\PYGZus{}softmax}\PYG{p}{(}\PYG{n}{logits}\PYG{p}{):}
  \PYG{n}{original\PYGZus{}probs} \PYG{o}{=} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{nn}\PYG{o}{.}\PYG{n}{softmax}\PYG{p}{(}\PYG{n}{logits}\PYG{p}{)}

  \PYG{n}{poly\PYGZus{}fit} \PYG{o}{=} \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{array}\PYG{p}{([}\PYG{o}{\PYGZhy{}}\PYG{l+m+mf}{0.037}\PYG{p}{,} \PYG{l+m+mf}{0.481}\PYG{p}{,} \PYG{o}{\PYGZhy{}}\PYG{l+m+mf}{2.3}\PYG{p}{,} \PYG{l+m+mf}{4.917}\PYG{p}{,} \PYG{o}{\PYGZhy{}}\PYG{l+m+mf}{1.791}\PYG{p}{])}  \PYG{c+c1}{\PYGZsh{} see Figure 6}
  \PYG{n}{entropy} \PYG{o}{=} \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{sum}\PYG{p}{(}\PYG{o}{\PYGZhy{}}\PYG{n}{original\PYGZus{}probs} \PYG{o}{*} \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{log}\PYG{p}{(}\PYG{n}{original\PYGZus{}probs} \PYG{o}{+} \PYG{l+m+mf}{1e\PYGZhy{}9}\PYG{p}{),}
                    \PYG{n}{axis}\PYG{o}{=\PYGZhy{}}\PYG{l+m+mi}{1}\PYG{p}{,} \PYG{n}{keepdims}\PYG{o}{=}\PYG{k+kc}{True}\PYG{p}{)}  \PYG{c+c1}{\PYGZsh{} compute the Shannon entropy}
  \PYG{n}{beta} \PYG{o}{=} \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{where}\PYG{p}{(}  \PYG{c+c1}{\PYGZsh{} beta = 1 / theta}
      \PYG{n}{entropy} \PYG{o}{\PYGZgt{}} \PYG{l+m+mf}{0.5}\PYG{p}{,}  \PYG{c+c1}{\PYGZsh{} don\PYGZsq{}t overcorrect low\PYGZhy{}entropy heads}
      \PYG{n}{jnp}\PYG{o}{.}\PYG{n}{maximum}\PYG{p}{(}\PYG{n}{jnp}\PYG{o}{.}\PYG{n}{polyval}\PYG{p}{(}\PYG{n}{poly\PYGZus{}fit}\PYG{p}{,} \PYG{n}{entropy}\PYG{p}{),} \PYG{l+m+mf}{1.0}\PYG{p}{),}  \PYG{c+c1}{\PYGZsh{} never increase entropy}
      \PYG{l+m+mf}{1.0}\PYG{p}{)}

  \PYG{k}{return} \PYG{n}{jax}\PYG{o}{.}\PYG{n}{nn}\PYG{o}{.}\PYG{n}{softmax}\PYG{p}{(}\PYG{n}{logits} \PYG{o}{*} \PYG{n}{beta}\PYG{p}{)}
\end{Verbatim}
\caption{Our implementation of adaptive temperature in JAX.} \label{fig:jax}
\end{figure*}


\begin{proposition}[Sharpness in Transformers necessitates large weights]\label{prop:weights}
Let $\mathbf{e}^{(n)} \in \mathbb{R}^n$ be a collection of $n$ logits, computed using dot product attention; i.e. $e^{(n)}_k = \left \langle \mathbf{Q}\mathbf{y}, \mathbf{K}\mathbf{x}_k 
ight 
angle$, where $\mathbf{y}\in\mathbb{R}^m$ is a query vector and $\mathbf{Q},\mathbf{K}\in\mathbb{R}^{m'	imes m}$ are parameters. Let $\delta = \max\limits_{1\leq i\leq n} e^{(n)}_i - \min\limits_{1\leq j\leq n} e^{(n)}_j$ be their maximum difference. Then $\delta$ is upper bounded as $\delta \leq 2 \sigma_\mathrm{max}^{(Q)} \sigma_\mathrm{max}^{(K)} \| \mathbf{y} \| \max_{1\leq i\leq n} \| \mathbf{x}_i \|$, where $\sigma_\mathrm{max}^{(Q)}, \sigma_\mathrm{max}^{(K)}\in\mathbb{R}$ are the largest singular values of $\mathbf{Q}$ and $\mathbf{K}$.
\end{proposition}
That is, the sharpness of the softmax in Transformers depends on the norm of its parameters.
Note that there is a common practice of leveraging operators such as \emph{layer normalisation} \citep{ba2016layer} extensively within Transformer architectures, which clamps $\|\mathbf{x}_i\|$ and $\|\mathbf{y}\|$ if applied right before the query-key mechanism, accentuating the impact of $\mathbf{Q}$ and $\mathbf{K}$'s singular values.

However, forcing large parameters promotes overfitting, and the likelihood that the \emph{incorrect} item gets the largest logit---see Figure 
ef{fig:soft_disp}. Setting temperature to zero will then \emph{degrade} accuracy---we might prefer to make the coefficients sharper while making sure that the chosen item is not left behind. This motivates our use of 	extbf{adaptive temperature}, where we vary $	heta$ depending on the \emph{entropy} in the input coefficients. Adaptive temperature can be elegantly motivated by the fact that decreasing the temperature must monotonically decrease the entropy, which is well-known in thermodynamics:
\begin{proposition}[Decreasing temperature decreases entropy]\label{prop:entropy}
    Let $\mathbf{e}^{(n)}\in\mathbb{R}^n$ be a collection of $n$ logits. Consider the Boltzmann distribution over these $n$ items, $p_i\propto\exp(-\beta e_i^{(n)})$ for $\beta\in\mathbb{R}$, and let $H = -\sum_i p_i\log p_i$ be its Shannon entropy. Then, as $\beta$'s magnitude increases, $H$ must monotonically decrease.
\end{proposition}
\begin{figure}
\centering
\includegraphics[width=\linewidth]{temperature_power_series.png}
  \caption{Entropy of $\mathtt{softmax}_	heta(\mathbf{a})$ for 10 elements of a power series $a_i = \lambda^i$, split into four regions depending on range of $(\lambda, 	heta)$. Degenerate cases: near $\lambda = 0$ and $\lambda = 1$ (all logits equal).}% Entropy increases with temperature but the rate at which it increases is heavily dependent on the attention logit distribution. Degenerate cases: near $\lambda = 0$ and $\lambda = 1$ all logits are the same, leading to highest entropy.}
  \label{fig:emp_entropy}
\end{figure}
\begin{figure}
    \includegraphics[width=0.49	extwidth]{polyfits.png}
    \caption{The polynomial fit used to derive our adaptive formula for $	heta$ as a function of the Shannon entropy, $H$. The fit degree-4 function was $	heta \approx 1 / (- 1.791 + 4.917H - 2.3H^2 + 0.481H^3 - 0.037H^4)$. We do not apply the correction to $	heta$ if predicted greater than $1$.}
    \label{fig:polyfit}
\end{figure}
Thus, since $\beta\propto\frac{1}{	heta}$ where $	heta$ is the temperature in $\mathtt{softmax}_	heta$, decreasing the temperature must monotonically decrease the entropy. We provide a full proof in Appendix 
ef{app:entropy}. To supplement Proposition 
ef{prop:entropy} empirically, we also provide---in Figure 
ef{fig:emp_entropy}---a visualisation of how the Shannon entropy varies with temperature, for a $10$-logit input with varying spread between the logits.

To compute the approximate temperature value as a function of entropy, we generate a dataset of inputs to our model where the maximal items do not obtain the highest logit (see Appendix 
ef{app:scheme} for a detailed overview). For each such input, we find the ``optimal'' value of $	heta$ that would maximise its probability. Then we fit an inverse degree-4 polynomial to this data---see Figure 
ef{fig:polyfit}---and use it to predict temperatures to use at inference time. Note we do not wish to increase entropy; as such, we do not correct $	heta$ to values greater than $1$.

The JAX \citep{jax2018github} implementation of our adaptive-$	heta$ 	exttt{softmax} is provided in Figure 
ef{fig:jax}, and we use it as a drop-in replacement for 	exttt{jax.nn.softmax} in all of our experiments.

While this approach requires two calls to 	exttt{jax.nn.softmax} in place of one, as well as computing several additional intermediate tensors, we are able to implement it in a way that allows the entropy correction computation to be fully \emph{streamed}, and hence compatible with efficient, scalable approaches like Flash Attention \citep{dao2022flashattention} that uses $O(n)$ rather than $O(n^2)$ memory to compute attention. We provide the derivation of our streamed algorithm in Appendix 
ef{app:stream}.

Note we are not the first to propose dynamically adapting temperature---\citet{neumann2018relaxed,radford2021learning} do this in the classification layer (and hence do not have to handle an ever-increasing amount of items), whereas \citet{chiang2022overcoming,cao2024mvsformer} perform it over intermediate attentional heads, but in a way that only depends on problem size (e.g. multiplying logits by $\log n$), hence not taking into account initial logit sharpness. It is important to also call out AERO \citep{jha2024aero}, a method which introduces \emph{learnable} temperature, and Entropix \citep{xjdr2024entropix}, a notable library for (var)entropy-based LLM sampling. 


\section{Experimental Results}

To validate the utility of our proposed adaptive temperature scheme, we evaluate it on both our previously-mentioned 	exttt{max} retrieval task---which allows us a pristine environment for evaluating whether adaptive temperature leads to more useful attention heads---as well as the CLRS-Text algorithmic reasoning benchmark \citep{markeeva2024clrs}, which represents a challenging reasoning task for decoder-only Transformers, and is hence likely to require low-entropy behaviour.

\setlength{	abcolsep}{0.7pt}

\begin{table*}[]
    \caption{Improvements observed when applying adaptive temperature on the 	exttt{max} retrieval task (without changing the parameters), averaged over ten seeds. $p$-values computed using a paired $t$-test.}
    \centering
    \begin{tabular}{cccccccccccc}
        	oprule
        &\cellcolor{green!15} {\bf ID size} & \multicolumn{10}{c}{\bf Out-of-distribution sizes}\\
        {\bf Model} & \cellcolor{green!15} $16$ & $32$ & $64$ & $128$ & $256$ & $512$ & $1,024$ & $2,048$ & $4,096$ & $8,192$ & $16,384$\\
        \midrule
        Baseline & \cellcolor{green!15} {$\mathbf{98.6}\%$} & $\mathbf{97.1}\%$ & ${94.3}\%$ & $89.7\%$ & $81.3\%$ & $70.1\%$ & $53.8\%$ & $35.7\%$ & $22.6\%$ & $15.7\%$ & $12.4\%$\\
        Adaptive $	heta$ & \cellcolor{green!15} $\mathbf{98.6}\%$ & $\mathbf{97.1}\%$ & $\mathbf{94.5}\%$ & $\mathbf{89.9}\%$ & $\mathbf{82.1}\%$ & $\mathbf{72.5}\%$ & $\mathbf{57.7}\%$ & $\mathbf{39.4}\%$ & $\mathbf{24.9}\%$ & $\mathbf{17.5}\%$ & $\mathbf{14.0}\%$\\
        \midrule
        $p$-value & \cellcolor{green!15} $0.4$ & $0.4$ & $\underline{0.002}$ & $\underline{2\cdot 10^{-5}}$ & $\underline{2\cdot 10^{-4}}$ & $\underline{3\cdot 10^{-5}}$ & $\underline{10^{-4}}$ & $\underline{6\cdot 10^{-4}}$ & $\underline{0.02}$ & $\underline{10^{-3}}$ & $\underline{4\cdot10^{-3}}$\\
        \bottomrule 
    \end{tabular}
    \label{tab:temp_results}
\end{table*}

\subsection{	exttt{Max} Retrieval} For this task, we first train our single attention head architecture as described in Appendix 
ef{app:experiment}; then, we evaluate it at various numbers of input items, with and without applying adaptive temperature to its sole 	exttt{softmax} function call. Note that this is a ``pure'' inference time adjustment---no modifications to the learned parameters are performed! 

The results---averaged over ten seeds and with statistical significance tests applied---are summarised in Table 
ef{tab:temp_results}. As is evident, applying adaptive temperature leads to a more performant retrieval head on larger inputs, with statistical significance ascertained via a paired $t$-test.

\begin{figure}
    \centering
    \includegraphics[width=\linewidth]{attentions.png}
    \includegraphics[width=\linewidth]{attentions_fixed.png}
    \caption{Visualising the attentional head for the 	exttt{max} retrieval task with (	extbf{below}) and without (	extbf{above}) adaptive temperature, for the same batch and model as in Figure 
ef{fig:soft_disp}. Note the increased sharpness in the coefficients, especially as the amount of items increases.}
    \label{fig:attn_temp}
\end{figure}

These results are further supplemented by a qualitative comparison of the 	exttt{softmax} coefficients before and after applying the temperature adaptation. As can be seen in Figure 
ef{fig:attn_temp}, our proposed adaptive temperature adaptation indeed leads to sharper coefficients on larger inputs and higher attention being directed to the desired item, even in situations where it did not receive the largest logit.

We have now successfully validated the predictions of our theory in a controlled environment. What about a more challenging benchmark with a baseline model comprising \emph{many} attentional heads?

\subsection{CLRS-Text}

\begin{figure*}
    \centering
    \includegraphics[width=\linewidth]{entropic.png}
    \caption{Resampling test results on CLRS-Text of variants of Gemma 2B, fine-tuned with and without adaptive temperature applied, on various problem sizes. Each point on the $x$ axis corresponds to a particular problem size in the corresponding algorithmic task. For example, on sorting tasks, this corresponds to the number of items being sorted; for graph tasks, it corresponds to the number of nodes in the graph. The 	extcolor{echonavy}{\bf blue} curves represent the accuracy of the baseline fine-tuned Gemma 2B model, whereas the 	extcolor{mymauve}{\bf red} curves represent the accuracy of that same model, fine-tuned with adaptive temperature. Both Gemma 2B variants were explicitly trained on CLRS-Text tasks---the training set sizes are denoted by red dots---and are evaluated zero-shot. Note that we limit our sample length to $2,048$ tokens, and only show performance metrics for sizes where the answer fits in this constraint.}
    \label{fig:entropic}
\end{figure*}

In this benchmark, we follow the protocol established by \citet{markeeva2024clrs} and fine-tune Gemma 2B models \citep{team2024gemma} on the thirty algorithmic execution tasks in CLRS-Text, plotting their performance profiles in- and out-of-distribution at various problem sizes.

While it may be tempting to directly re-apply our learned adaptive temperature function from Figure 
ef{fig:polyfit} solely at inference time---the same way we did in the 	exttt{max} retrieval experiments---this approach does not empirically work well in the CLRS-Text regime. This is due to the fact that CLRS-Text inputs are often textual representations of \emph{floating-point} numbers and therefore individual numbers often span \emph{multiple} tokens. It is therefore insufficient and inappropriate to aim for entropy levels where all the focus would be on \emph{one} token only, as was desirable in the 	exttt{max} retrieval task.

One follow-up on this could be to perform exactly the same polynomial fit exercise leading up to Figure 
ef{fig:polyfit}, only this time focussing on ``optimal'' values of temperature for Gemma's attentional heads. However, in this regime, we argue this exercise is substantially less trivial to do---as we are now dealing with a system spanning many attentional heads across many layers, it is not easy to even discover relevant attentional heads' behaviours, and even less so to ascertain that the model's robustness depends on those specific heads in those ways. As briefly discussed before, any such individual endeavour typically leads to a brand-new research project in mechanistic interpretability, and we do not find this to be in-scope of our paper.

That being said, there is an alternate route to make the Gemma model still benefit from our adaptive temperature module exactly as-is (i.e., with exactly the same polynomial fit as in Figure 
ef{fig:polyfit}); it just has to directly \emph{learn} how to leverage it. As such, in our CLRS-Text ablation we apply adaptive temperature both during fine-tuning and at inference time. What this means is, we replace all instances of 	exttt{jax.nn.softmax} within all the attentional heads of Gemma 2B with our 	exttt{adaptive\_temperature\_softmax} function, both during fine-tuning of the model and during inference. This allows the model to learn how to compute key/query embeddings that can maximally exploit the temperature adaptation.

These final comparative results may be found in Figure 
ef{fig:entropic}, and they demonstrate a significant advantage of the adaptive temperature-backed model on nearly all of the thirty algorithms under study. While we cannot make strong claims about where the underperformance on outlying algorithms comes from, a unifying property of some of them (Heapsort, MST Kruskal and Bubble Sort) in CLRS-Text is that they all occupy relatively large chunks of Gemma 2B's context window, which stretches further beyond the largest contexts over which the polynomial fit of our adaptive temperature was derived; this might cause unintended distribution shifts.
 
Our result indicates that, even in a complex system with many interactions between attentional heads, it is possible to extract benefits from the simple idea of dynamically adapting the temperature---and we hope our result paves the way for more involved future investigation of such approaches.

\section{Conclusions}

\emph{``Energy continuously flows from being concentrated\\
To becoming 	extbf{dispersed}, spread out, wasted and useless.''}---The 2nd Law: Unsustainable, by Muse

In this paper, we have provided extensive theoretical and empirical evidence that the 	exttt{softmax}---a key function in the design of modern frontier architectures---is leveraged within those architectures in a way that is fundamentally unable to sustain robust reasoning behaviours across all possible inputs. This is due to the fact its output coefficients are necessarily dispersing provided sufficient input elements. Our evidence relies on fundamental properties of the 	exttt{softmax}, as well as specific popular design decisions as tokenisation, global attention, and similar.

Beyond illustrating and proving these dispersion effects, we also attempted to use our theoretical framework to propose an \emph{adaptive temperature} approach that is able---at least to a certain extent---to hold the dispersion effect at bay. It is our opinion that the favourable results we observe with adaptive temperature warrant further investigation, and indicate that such adaptive layers are a strategy worth dedicating further attention to in future work. 

We conclude by remarking, once again, that adaptive temperature is merely an \emph{ad-hoc} method and it does not escape the conclusions of our theory! The key takeaway of our paper is \emph{not} the adaptive temperature proposal; it is the fact that we find it worthwhile to more seriously invest in research of hybrid architectures that will not fully rely on the 	exttt{softmax} function, at least within the confines of the assumptions of our theory. To name a few possibilities:

\begin{itemize}
    \item Any kind of unnormalised attention, such as \emph{linear} \citep{schmidhuber1992learning}, \emph{sigmoidal} \citep{ramapuram2024theory} or \emph{stick-breaking} attention \citep{tan2025scaling} does not have the dispersion issues presented here. That being said, it becomes substantially harder to meaningfully \emph{rank} items using them, see e.g. the GATv2 paper \citep{brody2022how}.
    \item Explicitly \emph{removing} excess attention could be another attractive alternative, as is robustly done in \emph{selective attention} \citep{leviathan2025selective} in a way that is capable of alleviating dispersion. The \emph{Differential Transformer} \citep{ye2025differential} follows a similar idea, though the exact way it is realised -- subtraction of two distributions \emph{after} the application of 	exttt{softmax} -- may not effectively evade dispersion, assuming it already happened in the individual 	exttt{softmax} calls.
    \item Similarly, forcing the attention to be \emph{hard} or \emph{local} \citep{martins2016softmax,correia-etal-2019-adaptively,peters-etal-2019-sparse,vitvitskyi2025makes} would also escape the confines of our theory. We already briefly discussed the challenges of learning with hard attention---local attention provides a very interesting alternative, but it must be stressed that ``out-of-distribution'' behaviours for certain heads may appear even at highly ``local'' scales; OOD here refers to going outside \emph{the largest problem size the head saw at training time}, 	extbf{not} the largest context deployed at training time.
    \item Lastly, our key Theorem relies on the model being built out of \emph{continuous} building blocks. Inserting \emph{discontinuities} in the feedforward layers---perhaps using approaches like \citet{dudzik2024asynchronous} as inspiration---would also break the assumptions of our theory, though it comes with obvious challenges to learning at scale.
\end{itemize}

While such approaches haven't seen as much success at scale as the ``vanilla'' Transformer, we hope our results inspire future work into making them stable, especially for constructing reasoning systems.

% Acknowledgements should only appear in the accepted version.
\section*{Acknowledgements}

We would like to deeply thank Daniel Johnson, the author of Penzai \citep{johnson2024penzai}---without it, our exploratory analyses would not be nearly as fruitful. Further, we thank Alex Matthews, Andrew Dudzik and Jo\~{a}o Ara\'{u}jo for helping us with the proof of one of our propositions, and Arthur Conmy, Neel Nanda and Csaba Szepesv\'{a}ri for reviewing the paper prior to submission and numerous highly useful comments and references. Lastly, we show deep appreciation to Alex Vitvitskyi, Olga Kozlova and Larisa Markeeva, for their endless engineering support with CLRS-Text.

\section*{Impact Statement}

This paper presents a theoretical and empirical study into the numerical properties of the 	exttt{softmax} function, placing into the spotlight several issues that models leveraging it will experience when generalising to longer inputs at test time. This is a setting that has clear implications to research in long context generalisation, which is an important frontier for reasoning in AI systems. Hence, any societal consequences of our work would potentially be contained within the consequences for general-purpose research in either of those established areas.


% In the unusual situation where you want a paper to appear in the
% references without citing it in the main text, use 
ocite

\bibliography{example_paper}
\bibliographystyle{icml2025}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% APPENDIX
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

ewpage
\appendix
\onecolumn
\input{append-math}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


\end{document}


% This document was modified from the file originally made available by
% Pat Langley and Andrea Danyluk for ICML-2K. This version was created
% by Iain Murray in 2018, and modified by Alexandre Bouchard in
% 2019 and 2021 and by Csaba Szepesvari, Gang Niu and Sivan Sabato in 2022.
% Modified again in 2023 and 2024 by Sivan Sabato and Jonathan Scarlett.
% Previous contributors include Dan Roy, Lise Getoor and Tobias
% Scheffer, which was slightly modified from the 2010 version by
% Thorsten Joachims & Johannes Fuernkranz, slightly modified from the
% 2009 version by Kiri Wagstaff and Sam Roweis's 2008 version, which is
% slightly modified from Prasad Tadepalli's 2007 version which is a
% lightly changed version of the previous year's version by Andrew
% Moore, which was in turn edited from those of Kristian Kersting and
% Codrina Lauth. Alex Smola contributed to the algorithmic style files.