Replacing softmax with ReLU in Vision Transformers

ArXiv: 2309.08586

🎯 Pitch

Vision transformers don’t need softmax to scale—just applying ReLU to attention scores and dividing by sequence length matches softmax accuracy across model sizes. This simple tweak eliminates exponentials and cross-sequence normalization, unlocking new parallelization strategies.

---

1. Executive Summary

This paper studies whether the softmax operation in vision transformer attention can be replaced with a pointwise activation without sacrificing accuracy scaling behavior. The experiments train small to large vision transformers on ImageNet-21k and evaluate on ImageNet-1k. The core mechanism is ReLU-attention — applying ReLU to the query-key dot products and dividing by sequence length $L$ — which the paper finds is critical for matching softmax performance, as prior work omitting this scaling observed accuracy degradation. Across model sizes S/32, S/16, S/8, B/32, B/16, and L/16, ReLU-attention with $L^{-1}$ scaling approaches or matches the compute-accuracy scaling trend of softmax attention, establishing that a non-normalizing pointwise activation can substitute for softmax only when sequence-length normalization is applied to preserve the expected $O(L^{-1})$ magnitude of attention weights at initialization.

2. Context and Motivation

The Core Problem: Softmax Is a Computational Bottleneck in Attention

The transformer architecture depends on an attention mechanism that computes a weighted sum over input tokens, where the weights are produced by a softmax operation. Softmax performs two expensive operations that create practical bottlenecks for efficient deployment: an exponentiation on every query-key dot product, and a sum over the entire sequence length to normalize the weights into a probability distribution. The sum over sequence length is particularly problematic because it forces the attention computation to gather values from all positions before producing any output — it is a global synchronization barrier that makes parallelization across the sequence length dimension challenging.

The paper characterizes this cost directly. The exponentiation step requires evaluating $\exp(x)$ for every element of the attention matrix $QK^\top$, which has dimensions $L \times L$ for sequence length $L$. The subsequent normalization sum $\sum_{j=1}^L \exp(q_i \cdot k_j)$ for each query $i$ similarly touches all $L$ key positions, preventing efficient sharding along the sequence axis. This is a well-recognized bottleneck in the systems literature — works like FlashAttention (Dao et al., 2022) and Rabe and Staats (2021) cite these operations as the primary obstacles to memory-efficient and parallelizable attention implementations.

The paper asks a direct question: can we simply remove softmax and replace it with something computationally cheaper? A pointwise activation like ReLU avoids both the exponentiation and the normalization sum. ReLU is a simple threshold operation ($\max(0, x)$) that operates independently on each element with no interaction between positions. If ReLU could substitute for softmax without degrading model accuracy, attention could be parallelized across the sequence length dimension with fewer gather operations — each output position would depend only on its own query-key interactions, not on a global normalization constant computed from all positions.

However, prior work had already observed that naive ReLU-for-softmax substitution degrades accuracy. The paper cites Shen et al. (2023) and Hua et al. (2022), which found that simply replacing softmax with ReLU or squared ReLU produced worse results than standard attention. The field thus faced a tension: the computational motivation for removing softmax was clear, but the empirical evidence suggested softmax was functionally necessary for good performance. This paper's core contribution is identifying why those naive substitutions failed and providing a simple fix — dividing by sequence length — that closes the gap.

Why This Problem Matters: Computational Efficiency and Scaling

The practical motivation is straightforward. Vision transformers process images as sequences of patches, with sequence lengths that grow quadratically in resolution. A $224 \times 224$ image at a typical patch size of $16 \times 16$ produces $L = 196$ tokens. Higher-resolution images or video inputs can push $L$ into the thousands. The $O(L^2)$ memory and compute cost of attention is already the dominant factor for large models, and the softmax's gather operation — which requires reading every key position before computing any attention output — compounds this in distributed settings by forcing communication between devices sharding different parts of the sequence.

Pointwise attention without normalization eliminates this gather. Each element $\alpha_{ij} = \phi(q_i^\top k_j / \sqrt{d})$ depends only on the single pair $(i, j)$. In a distributed setting where different heads or different sequence blocks are assigned to different devices, this independence means fewer cross-device communication steps. The paper frames this explicitly as the primary benefit: "ReLU-attention can be parallelized over the sequence length dimension with fewer gather operations than traditional attention."

There is also a more subtle scaling concern. As transformers scale to larger models and longer sequences, the softmax's exponentiation becomes numerically fragile. The standard fix — subtracting the maximum before exponentiating — still requires computing that maximum across the sequence, another gather operation. A pointwise activation sidesteps both the numerical instability and the mitigation cost.

Prior Approaches and Their Shortcomings

The paper identifies several lines of prior work that attempted to modify or remove attention's softmax, each with limitations:

Naive pointwise replacements (Shen et al., 2023; Hron et al., 2020; Hua et al., 2022). These works substituted softmax with ReLU or squared ReLU directly, without any sequence-length normalization. The result was accuracy degradation relative to softmax attention. The paper's experiments confirm this: when the scaling exponent $\alpha$ is set to 0 (meaning no sequence-length scaling, just the raw activation), accuracy consistently underperforms softmax (Figure 2). These prior works correctly identified the computational motivation but missed the crucial normalization insight — the expected magnitude of attention weights changes when softmax is removed, and model hyperparameters tuned for softmax's $O(L^{-1})$ per-weight scale no longer work properly.

Linear attention (Katharopoulos et al., 2020; Lu et al., 2021; Koohpayegani and Pirsiavash, 2022). A second line of work removes the activation function entirely, making attention linear so that the order of matrix multiplications can be rearranged from $(QK^\top)V$ to $Q(K^\top V)$. This changes the computational complexity from $O(dL^2)$ to $O(d^2L)$, which is beneficial when $L \gg d$. The paper explicitly notes, however, that "removing the activation entirely reduced accuracy" in their experiments (Figure 2, $\alpha = 0$ with identity activation). The identity function produces unconstrained attention weights that are likely too large or poorly distributed relative to what the rest of the network expects, causing a performance drop.

Scale-invariant architectures (Li et al., 2022). This work replaced softmax while still requiring normalization over the sequence length axis to ensure the attention weights sum to one. The paper points out that this "retains the downside of requiring a gather" — the normalization sum is exactly the operation that prevents efficient parallelization. So while this approach changed the functional form, it preserved the computational bottleneck, defeating much of the purpose of removing softmax.

Gated attention units (Hua et al., 2022). Another approach adds a gating mechanism — an extra projection followed by elementwise multiplication with the attention output — without scaling by sequence length. The paper investigates this combination (Figure 4) and finds that while gating helps slightly, the best accuracy is still achieved with sequence-length scaling, whether or not a gate is present. Notably, gating increases computational cost by roughly 9.3% for the S/8 model with ReLU, making it a less attractive efficiency tradeoff.

Theoretical work on softmax-free attention (Bai et al., 2023; Fu et al., 2023). After writing the initial version of this paper, the authors learned that the specific variant of ReLU-attention they study had been independently explored with theoretical motivation. Bai et al. (2023) studied what a single attention layer can learn and used ReLU-normalized-by-sequence-length as part of a theoretical analysis. This independent discovery reinforces the idea that the $L^{-1}$ scaling is not merely an ad-hoc trick but emerges naturally from the structure of the problem.

How This Paper Positions Itself

The paper's positioning is modest and empirical. It does not claim to invent a novel attention mechanism — the idea of using ReLU in attention existed, and the specific $L^{-1}$ scaling appeared in concurrent theoretical work. Instead, the paper's contribution is demonstrating that this simple modification, which prior practical work had overlooked, is sufficient to close the gap with softmax in vision transformers at scale. The paper fills a gap between the theoretical motivation (Bai et al., 2023; Fu et al., 2023) and the practical implementation (Shen et al., 2023; Hua et al., 2022), showing via systematic experiments on ImageNet-21k that the $L^{-1}$ scaling factor is the missing ingredient that makes pointwise attention work.

The paper also positions itself as raising questions rather than providing final answers. The conclusion explicitly states: "This report leaves many open questions. In particular, we are unsure why the factor $L^{-1}$ improves performance or if this term could be learned. Moreover, it is likely that there is a better activation function that we do not explore." This framing suggests the paper is a proof-of-concept rather than a definitive solution — it establishes that softmax-free attention can work, opening the door for further optimization of the specific activation and normalization choices.

The analytical motivation for the $L^{-1}$ scaling is deliberately brief and approximate. The paper observes that under softmax, $\sum_j \alpha_{ij} = 1$, which implies $\mathbb{E}_j[\alpha_{ij}] = L^{-1}$ if the weights are viewed as expectation-taking over positions. While the paper does not claim this is a necessary condition, it hypothesizes that "preserving this condition may alleviate the need to change other hyperparameters when replacing softmax." The argument then notes that at initialization, $q$ and $k$ elements are $O(1)$, so $q_i^\top k_j / \sqrt{d}$ is $O(1)$, and ReLU preserves this order of magnitude. Therefore, a factor of $L^{-1}$ ensures $\mathbb{E}_j[\alpha_{ij}]$ remains $O(L^{-1})$, matching the scale that downstream components of the transformer expect. This is not a rigorous derivation but rather a heuristic justification for why prior work's omission of sequence-length scaling caused failures: those naive replacements produced attention weights of $O(1)$ rather than $O(L^{-1})$, changing the effective learning dynamics and requiring different hyperparameters.

Importantly, the paper tests this hypothesis by sweeping over the exponent $\alpha$ in the scaling factor $L^{-\alpha}$ for $\alpha \in [0, 1]$ (Figure 2). The consistent finding — across different activation functions, model patch sizes, and datasets — is that $\alpha \approx 1$ is optimal. Values of $\alpha$ near 0 (the setting of prior work) underperform, and values of $\alpha$ near 1 approach or match softmax. This systematic sweep across the key hyperparameter is what converts the observation from speculation into evidence, and it is the central empirical contribution that distinguishes this work from both the prior experiments (which didn't try sequence-length scaling) and the prior theory (which didn't validate at scale on ImageNet-21k).

3. Technical Approach

3.1 Reader Orientation

This paper proposes a simple modification to the standard transformer attention mechanism: replace the softmax operation with a pointwise ReLU activation followed by division by the sequence length. The system being built is not a new architecture but rather a drop-in replacement for one operation inside existing vision transformer architectures, with no changes to hyperparameters, training recipes, or model structure beyond the attention computation itself. The problem this solves is the computational bottleneck imposed by softmax — specifically its exponentiation and its global normalization sum — which prevents efficient parallelization across sequence positions. The "shape" of the solution is a single formula substitution: wherever the transformer computes $\text{softmax}(QK^\top / \sqrt{d})$, it instead computes $L^{-1} \cdot \text{ReLU}(QK^\top / \sqrt{d})$, where $L$ is the sequence length. The paper's insight is that this seemingly minor change works only when the sequence-length division is included, and that prior work's failure to include it explains why naive ReLU-for-softmax substitution degraded accuracy.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three logical components, though two of them are simply the standard vision transformer pipeline operating under a modified attention computation:

1. Standard Vision Transformer Backbone (ViT) — the full model architecture from Dosovitskiy et al. (2021) using the configurations from the BigVision codebase (Beyer et al., 2022). This includes patch embedding, positional encoding, transformer blocks with multi-head attention and MLP layers, and a classification head. All components remain identical to the standard ViT; only the attention computation inside each transformer block is modified.

2. Modified Attention Computation (ReLU-attention) — inside every attention head, the standard two-step procedure applies: (a) compute attention weights from queries and keys using the modified function $\alpha_{ij} = L^{-1} \cdot \text{ReLU}(q_i^\top k_j / \sqrt{d})$, then (b) compute output values as the weighted sum $\sum_j \alpha_{ij} v_j$. The modification replaces softmax with a pointwise activation and sequence-length scaling, eliminating the normalization sum and exponentiation.

3. Optional Architectural Variants (qk-layernorm, gating) — two additional components that can be combined with the modified attention: query-key LayerNorm (Dehghani et al., 2023), which normalizes queries and keys before computing attention weights, and gated attention units (Hua et al., 2022), which add an extra projection and elementwise multiplication after attention. These are not part of the core proposal but are studied in ablation experiments to understand their interaction with sequence-length scaling.

Information flows identically to a standard ViT: an input image is patchified and embedded → the sequence of patch tokens passes through a series of transformer blocks → inside each block, the modified attention computes a weighted sum of values using ReLU-based weights → the output passes through an MLP → after all blocks, the classification token is projected to class logits. The only change is in the attention weight computation inside step two, where softmax is replaced by ReLU divided by sequence length.

3.3 Roadmap for the Deep Dive

• First, the standard attention formulation (Equation 1) that the paper modifies. Understanding the baseline is necessary to see what changes and why the modification matters computationally.
• Second, the ReLU-attention formula and the general scaled pointwise attention family it belongs to. This covers the exact mathematical substitution, the family of activation functions tested, and the sequence-length scaling parameter $\alpha$ that is the paper's key hyperparameter.
• Third, the analytical motivation for sequence-length scaling. The paper provides a brief heuristic argument for why $L^{-1}$ matters; we will unpack this argument step by step, explaining what goes wrong without it and why preserving the $O(L^{-1})$ expected weight magnitude at initialization is important.
• Fourth, the experimental configurations and training setup. This covers the models, datasets, training recipes, hyperparameters, and evaluation protocols — all the concrete details that make the empirical claims reproducible.
• Fifth, the ablation dimensions studied: the choice of activation function $h$, the exponent $\alpha$, the role of qk-layernorm, and the interaction with gated attention units. Each ablation answers a specific question about what components are necessary for ReLU-attention to match softmax.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an empirical investigation paper whose core idea is that replacing softmax in vision transformer attention with ReLU divided by sequence length can match softmax's accuracy scaling behavior, and that the sequence-length division — omitted in prior work — is the critical ingredient that makes this substitution work.

---

Standard Attention Formulation (The Baseline Being Modified)

The paper begins from the standard definition of attention in transformers. Attention operates on three matrices derived from the input sequence: queries $Q$, keys $K$, and values $V$, all of dimension $L \times d$ for sequence length $L$ and hidden dimension $d$. The standard two-step procedure (Equation 1) is:

Step 1 — Compute attention weights:

$\alpha_{ij} = \phi\left( \frac{1}{\sqrt{d}} \left[ q_i^\top k_1, q_i^\top k_2, ..., q_i^\top k_L \right] \right)_j$

where $\phi$ is an activation function applied to the vector of scaled dot products between query $i$ and all keys. The subscript $j$ indicates selecting the $j$-th element of the output vector — the weight for key $j$ when computing the output for query $i$. In standard transformers, $\phi$ is the softmax function, defined as:

$\text{softmax}(x)_j = \frac{\exp(x_j)}{\sum_{k=1}^L \exp(x_k)}$

where each element $x_j$ is the scaled dot product $q_i^\top k_j / \sqrt{d}$. The softmax produces a probability distribution: each $\alpha_{ij} \geq 0$ and $\sum_{j=1}^L \alpha_{ij} = 1$ for every query $i$.

Step 2 — Compute output values:

$o_i = \sum_{j=1}^L \alpha_{ij} v_j$

where $v_j$ is the value vector for position $j$ and $o_i$ is the output for query position $i$. This is a weighted sum of all value vectors, with weights determined by the attention mechanism.

The computational cost of softmax comes from two operations. The exponentiation $\exp(x_j)$ requires evaluating the exponential function for every element of the $L \times L$ attention matrix — $L^2$ operations, each more expensive than a simple multiply or add. The normalization sum $\sum_k \exp(x_k)$ requires, for each of the $L$ queries, summing over all $L$ key positions — another $L^2$ operations that additionally forces a gather: the sum cannot be computed until all exponentiated values are available, which means all positions must synchronize before any attention output can be produced. This gather is the parallelization bottleneck because it prevents computing different parts of the attention matrix independently on different devices without communication.

The scaling factor $1/\sqrt{d}$ is standard in transformers and serves to keep the dot products from growing with dimension. At initialization, if query and key components are independent with mean 0 and variance 1, then $q_i^\top k_j = \sum_{m=1}^d q_{im} k_{jm}$ has mean 0 and variance $d$, so the dot product is $O(\sqrt{d})$. Dividing by $\sqrt{d}$ brings it back to $O(1)$, preventing the softmax from saturating (producing near-one-hot weights) at initialization. This scaling is preserved in the paper's modified attention.

---

ReLU-Attention and the Scaled Pointwise Attention Family

The paper replaces the softmax function $\phi$ with a family of functions parameterized by two choices: the activation function $h$ and the sequence-length scaling exponent $\alpha$. The general form is:

$\phi(x) = L^{-\alpha} \cdot h(x)$

where $L$ is the sequence length, $\alpha \in [0, 1]$ is the scaling exponent, and $h$ is a pointwise activation function applied independently to each element. "Pointwise" means $h$ operates on each input element $x_j$ independently — the output $h(x_j)$ depends only on $x_j$, not on any other $x_k$ for $k \neq j$. This is in contrast to softmax, which is not pointwise: the output $\text{softmax}(x)_j$ depends on all elements of $x$ through the denominator $\sum_k \exp(x_k)$.

When applied to the attention weight computation, this gives:

$\alpha_{ij} = L^{-\alpha} \cdot h\left( \frac{q_i^\top k_j}{\sqrt{d}} \right)$

ReLU-attention is the specific case where $h$ is the ReLU (Rectified Linear Unit) function and $\alpha = 1$:

$\alpha_{ij} = L^{-1} \cdot \text{ReLU}\left( \frac{q_i^\top k_j}{\sqrt{d}} \right)$

where ReLU is defined as:

$\text{ReLU}(x) = \max(0, x)$

ReLU outputs zero for negative inputs and passes positive inputs through unchanged. In the context of attention, this means query-key pairs with negative similarity (after scaling) receive zero weight, while pairs with positive similarity receive weight proportional to their similarity, divided by $L$.

The full family of activation functions tested. The paper explores seven activation functions for $h$, all applied pointwise (Figure 2):

• relu: $\max(0, x)$ — the standard rectifier
• squared relu ($\text{relu}^2$): $(\max(0, x))^2$ — squares the positive outputs, penalizing small positive values more severely
• gelu: the Gaussian Error Linear Unit (Hendrycks and Gimpel, 2016), a smooth approximation to ReLU that weights inputs by their probability of being positive under a Gaussian
• softplus: $\log(1 + \exp(x))$ — a smooth approximation to ReLU that is never exactly zero
• identity: $h(x) = x$ — no activation at all, just the raw scaled dot product
• relu6: $\min(6, \max(0, x))$ — ReLU clipped at a maximum value of 6
• sigmoid: $1 / (1 + \exp(-x))$ — the classic S-shaped function that outputs values in $(0, 1)$

These span a range of behaviors: hard thresholding (ReLU, ReLU6), smooth thresholding (GELU, Softplus), unbounded scaling (identity), and bounded outputs (Sigmoid). The paper systematically sweeps the scaling exponent $\alpha$ from 0 to 1 for each activation on three model configurations (S/32, S/16, S/8) and two training datasets (ImageNet-21k and ImageNet-1k). The key finding (visible across all panels of Figure 2) is that all activations perform best when $\alpha$ is close to 1, and at $\alpha \approx 1$, there is no clearly superior activation — they all approach similar accuracy. The paper selects ReLU for the main experiment "for its speed," since ReLU is the computationally cheapest of the tested options (a simple comparison and conditional assignment, versus the exponential and division in softplus or sigmoid, or the additional multiply in squared ReLU).

---

Analytical Motivation for Sequence-Length Scaling

The paper provides a brief heuristic argument for why scaling by $L^{-1}$ (i.e., $\alpha = 1$) is important. While the paper states this argument is not a rigorous derivation, it provides the conceptual foundation for understanding why prior work — which omitted this scaling — observed accuracy degradation.

The softmax baseline constraint. Under softmax, the attention weights for query $i$ satisfy:

$\sum_{j=1}^L \alpha_{ij} = 1$

This is a hard constraint enforced by the normalization denominator. A direct consequence is that the expected value of $\alpha_{ij}$ over the key positions $j$, treating the distribution as uniform across positions, is:

$\mathbb{E}_j[\alpha_{ij}] = \frac{1}{L} \sum_{j=1}^L \alpha_{ij} = \frac{1}{L}$

In words: if the softmax produces weights that sum to 1 over $L$ positions, the average weight per position is $L^{-1}$. This is not a claim about any individual $\alpha_{ij}$ — some will be much larger, some much smaller — but rather about the scale of the typical weight that downstream components of the transformer expect to receive.

What happens to the expected weight when softmax is removed. If softmax is replaced by a pointwise activation $h(x)$ without any sequence-length scaling, each $\alpha_{ij} = h(q_i^\top k_j / \sqrt{d})$ is independent of $L$. At initialization, the elements of $q$ and $k$ are typically drawn from distributions with $O(1)$ variance (e.g., Gaussian or uniform initialization scaled to preserve variance through layers), so the scaled dot product $q_i^\top k_j / \sqrt{d}$ is also $O(1)$. For activations like ReLU that preserve this order of magnitude (since $\max(0, x)$ for $x = O(1)$ is also $O(1)$), each $\alpha_{ij}$ is $O(1)$. The expected value over positions is therefore:

$\mathbb{E}_j[\alpha_{ij}] = \mathbb{E}_j\left[h\left(\frac{q_i^\top k_j}{\sqrt{d}}\right)\right] = O(1)$

This is fundamentally different from the softmax case, where $\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$. The attention weights are now $L$ times larger in expectation than what the rest of the network was designed for. This matters because the output $o_i = \sum_j \alpha_{ij} v_j$ sums over $L$ terms, each of weight $O(1)$ instead of $O(L^{-1})$, producing outputs of magnitude $O(L)$ instead of $O(1)$. The downstream LayerNorm and residual connections, which were tuned for softmax-scale outputs, now receive inputs of a very different magnitude, potentially disrupting training dynamics.

The $L^{-1}$ correction. Dividing by $L$ restores the expected order:

$\mathbb{E}_j[\alpha_{ij}] = \mathbb{E}_j\left[L^{-1} \cdot h\left(\frac{q_i^\top k_j}{\sqrt{d}}\right)\right] = L^{-1} \cdot O(1) = O(L^{-1})$

Now the expected per-position weight matches the softmax case. The output $o_i = \sum_j \alpha_{ij} v_j$ becomes $O(1)$ again, since the $L$ terms each of size $O(L^{-1})$ sum to a constant. The paper states this explicitly: "Preserving this condition may alleviate the need to change other hyperparameters when replacing softmax." The key word is "alleviate" — the authors are not claiming that $O(L^{-1})$ expected weight is necessary for good performance, only that matching this property of softmax reduces the risk of disrupting other parts of the training pipeline that were tuned under the softmax assumption.

Why $\alpha = 1$ rather than some other value near 1. The empirical sweep in Figure 2 answers this question. For all activations and configurations tested, the accuracy curve as a function of $\alpha$ peaks at or very near $\alpha = 1$. Values of $\alpha$ near 0.5 or 1.5 are consistently worse. The paper does not attempt to derive why $\alpha = 1$ is optimal — it presents the $O(L^{-1})$ argument as motivation and lets the data confirm that this is the right scaling. The analytical motivation explains the left side of the curve (why $\alpha = 0$ fails) but does not predict the shape beyond that.

A nuance: squared ReLU is not $O(1)$. The paper notes a small technical detail: "With the exception of squared ReLU." For squared ReLU, $h(x) = (\max(0, x))^2$, so if $x = O(1)$, then $h(x) = O(1)$ as well (the square of an $O(1)$ quantity is $O(1)$). The paper's point is that squared ReLU does not increase the asymptotic order — it doesn't make $O(1)$ into $O(L)$, for example — so the $L^{-1}$ scaling argument still applies. The exception phrasing simply acknowledges that the variance or constant factors differ between activations, but the scaling order argument is unchanged.

---

Experimental Configurations and Training Setup

All experiments use vision transformer (ViT) architectures as introduced by Dosovitskiy et al. (2021) and refined by Beyer et al. (2022) in the BigVision codebase. The paper uses existing training configurations without modifying hyperparameters — this is a deliberate choice to test whether ReLU-attention works as a drop-in replacement without requiring per-task tuning.

Models tested. The paper trains six vision transformer variants spanning a range of scales:

• S/32: Small variant with patch size $32 \times 32$. For a $224 \times 224$ input image, this gives $L = (224/32)^2 = 7^2 = 49$ tokens (plus a classification token, making $L = 50$ total sequence length).
• S/16: Small variant with patch size $16 \times 16$, giving $L = (224/16)^2 = 14^2 = 196$ tokens (plus CLS token, $L = 197$).
• S/8: Small variant with patch size $8 \times 8$, giving $L = (224/8)^2 = 28^2 = 784$ tokens (plus CLS token, $L = 785$).
• B/32: Base variant, patch size 32, $L \approx 50$.
• B/16: Base variant, patch size 16, $L \approx 197$.
• L/16: Large variant, patch size 16, $L \approx 197$.

The small (S), base (B), and large (L) variants differ primarily in the number of transformer blocks and the embedding dimension, following standard ViT sizing. The patch size controls sequence length, with smaller patches producing longer sequences (more tokens) at higher computational cost. Testing across both model scale and sequence length tests whether the $L^{-1}$ scaling works for varying $L$, which is important because the scaling factor itself depends on $L$.

Training datasets and recipes. Two training configurations are used, drawn directly from the BigVision codebase without modification:

• ImageNet-21k (i21k) training: The model is trained on ImageNet-21k (Deng et al., 2009), a dataset with approximately 14 million images spanning 21,841 classes. Training runs for 30 epochs with a specific learning rate schedule, optimizer, and augmentation policy from the BigVision i21k configuration. The paper specifies that this produces roughly $9 \times 10^5$ training steps.
• ImageNet-1k (i1k) training: The model is trained on ImageNet-1k, with approximately 1.28 million images in 1,000 classes. Training runs for 300 epochs using the BigVision i1k configuration, also producing roughly $9 \times 10^5$ steps. The longer epoch count compensates for the smaller dataset, yielding similar total optimization steps.

The matching step count between the two configurations (roughly 900,000 steps) allows a fair comparison of how dataset size and diversity affect the softmax-vs-ReLU-attention tradeoff.

Default architectural choices. All models use qk-layernorm (Dehghani et al., 2023) by default, as "this was previously observed to be necessary to prevent instability when scaling model size." QK-layernorm applies Layer Normalization (Ba et al., 2016) to the queries and keys before they enter the attention computation:

$q_i \leftarrow \text{LayerNorm}(q_i), \quad k_j \leftarrow \text{LayerNorm}(k_j)$

This normalization reduces the variance of the dot products and was found in prior work to be important for training stability in large models. The paper includes an ablation (Figure 3) testing whether qk-layernorm is necessary specifically for ReLU-attention.

Evaluation protocol. The paper uses two evaluation modes:

1. Direct ImageNet-1k accuracy for i21k-trained models: Models trained on ImageNet-21k are evaluated on ImageNet-1k by "taking the top class among those that are in ImageNet-1k, without fine-tuning." This means the model predicts all 21,841 classes, but only the 1,000 classes that correspond to ImageNet-1k are considered when computing accuracy — the model's raw performance on the ImageNet-1k subset without any task-specific adaptation.

2. 10-shot linear transfer: For evaluating representation quality on downstream tasks, a linear classifier is trained on top of the frozen ViT features using 10 labeled examples per class from each of eight downstream datasets: Caltech Birds (Welinder et al., 2010), Caltech-101 (Fei-Fei et al., 2004), Stanford Cars (Krause et al., 2013), CIFAR-100 (Krizhevsky et al., 2009), DTD (Cimpoi et al., 2014), ColHist (Kather et al., 2016), Pets (Parkhi et al., 2012), and UC Merced (Yang and Newsam, 2010). The linear probe is trained three times with different random seeds — corresponding to different selections of the 10 examples per class — and results are averaged. This tests whether the learned representations transfer to diverse visual recognition tasks, not just ImageNet classification.

Compute measurement. The x-axis in the main result (Figure 1) uses TPU core hours, which is the total wall-clock time multiplied by the number of TPU cores used. This accounts for the actual computational cost of each attention variant, including any overhead from the activation function. The paper notes that gating "increases the core hours required for the experiment by roughly 9.3% for the S/8 model with ReLU," indicating that compute is measured empirically from actual training runs, not estimated from FLOP counts.

---

Ablation Dimensions and Design Choices

The paper systematically investigates four dimensions that could affect whether ReLU-attention matches softmax performance.

Activation function choice (Figure 2). The primary sweep tests seven activation functions $h$ at various scaling exponents $\alpha$. The experiment uses the S/32, S/16, and S/8 models trained on both i21k and i1k. For each combination, the x-axis sweeps $\alpha$ from 0 to 1 in increments (the figure shows continuous curves), and the y-axis shows ImageNet-1k accuracy.

The universal finding is that all activations perform poorly at $\alpha = 0$ and improve as $\alpha$ approaches 1, with peak accuracy typically at $\alpha \approx 1$. At $\alpha = 1$, the different activations produce similar accuracy — there is "no clear best non-linearity." This has two implications. First, it confirms that the sequence-length scaling is the dominant factor, not the choice of activation. Second, it justifies using ReLU for the main experiment: ReLU is the fastest activation among those tested (a simple max operation with no exponentials, divisions, or special functions), so if all activations perform similarly, choosing the fastest maximizes the computational benefit.

The identity activation (no activation at all) performs worst across all settings, confirming that "removing the activation entirely reduced accuracy." Without any nonlinearity, the attention weights can be arbitrarily positive or negative and are not constrained to any particular scale or sign pattern, which may cause training instability or poor gradient flow.

Sequence length scaling exponent (Figure 2). The sweep over $\alpha$ answers the question: is $L^{-1}$ specific or would any sequence-length-dependent scaling work? The result is that $\alpha \approx 1$ is consistently optimal across activations, model sizes, and datasets. Values near $\alpha = 0$ (the prior work setting) underperform substantially; values near $\alpha = 1$ approach or match softmax; intermediate values produce intermediate performance. This monotonic relationship is not a foregone conclusion — one could imagine that $\alpha = 0.5$ (scaling by $1/\sqrt{L}$) might be optimal for some theoretical reason — but the data clearly favor $\alpha = 1$.

The paper does not test $\alpha > 1$, which would mean dividing by more than $L$ (e.g., $L^{-2}$ would make $\mathbb{E}_j[\alpha_{ij}] = O(L^{-2})$). The sweep is restricted to $[0, 1]$, presumably because beyond $\alpha = 1$ the expected attention weights become smaller than in the softmax case, which would likely harm performance by making the attention output too small relative to the residual stream.

QK-layernorm ablation (Figure 3). The paper tests whether qk-layernorm is necessary for ReLU-attention to work by comparing ReLU and squared ReLU with and without qk-layernorm, across the $\alpha$ sweep and the S/32, S/16, S/8 models trained on i21k. The result is that "qk-layernorm does not have a large effect for these models." The curves with and without qk-layernorm are close together for both activations. However, the paper includes an important caveat: "this may change at scale" — larger models than those tested (B/32, B/16, L/16 appear in Figure 1 but not in this ablation) might require qk-layernorm for stability. The paper retains qk-layernorm in the main experiments as a conservative choice based on prior scaling work (Dehghani et al., 2023).

The fact that qk-layernorm is not critical for ReLU-attention at the tested scales is noteworthy because it suggests the $L^{-1}$ scaling already provides sufficient normalization. Without qk-layernorm, the $q_i^\top k_j / \sqrt{d}$ terms have variance determined by the weight initialization and the network depth; with qk-layernorm, they are explicitly normalized to have approximately unit variance. The insensitivity to this choice suggests that ReLU-attention is tolerant of a range of dot-product scales, perhaps because the ReLU threshold acts as a scale-independent filter (only the sign matters for whether a weight is zero or positive) and the $L^{-1}$ factor handles the overall magnitude.

Gated attention unit interaction (Figure 4). Hua et al. (2022) proposed replacing softmax with a gating mechanism and did not scale by sequence length. The paper tests whether adding such a gate eliminates the need for $L^{-1}$ scaling. The gated attention unit adds an extra projection after the attention computation:

$o_i' = o_i \odot \sigma(W_g x_i)$

where $o_i$ is the standard attention output (from ReLU-attention in this case), $W_g$ is a learned projection matrix, $\sigma$ is the sigmoid function, $\odot$ is elementwise multiplication, and $x_i$ is the input to the attention block. This gate modulates the attention output elementwise based on the input — it can amplify or suppress individual feature dimensions of the attention output.

The experiment tests all four combinations: {ReLU, squared ReLU} $\times$ {with gate, without gate} across the $\alpha$ sweep. The result is that "the best accuracy is still achieved with sequence length scaling, with or without the gate." Gating shifts the curves slightly — performance at low $\alpha$ improves a bit — but the peak remains at $\alpha \approx 1$, and the gated version does not outperform the ungated version at the optimal $\alpha$. The paper also reports that "gating increases the core hours required for the experiment by roughly 9.3% for the S/8 model with ReLU," meaning the additional computation from the gate is not justified by accuracy gains at the optimal $\alpha$. This ablation establishes that the gating approach from prior work and the sequence-length-scaling approach from this paper address different aspects of the softmax removal problem, and that the $L^{-1}$ scaling is the more effective intervention.

Summary of Design Choices and Their Justifications

• Pointwise activation over softmax: Eliminates exponentiation and the normalization sum, enabling better parallelization across sequence positions. The motivation is purely computational — same accuracy with fewer gather operations.
• ReLU as the activation: Selected over six alternatives tested (squared ReLU, GELU, softplus, identity, ReLU6, sigmoid) because all perform similarly at $\alpha \approx 1$, and ReLU is the fastest (a simple max operation with no exponentials, divisions, or special-function evaluations).
• Division by sequence length $L$: Ensures $\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$ at initialization, matching the scale that softmax produces and preventing the training dynamics disruption that naive pointwise replacement causes. The exponent $\alpha = 1$ is chosen empirically from a sweep over $[0, 1]$ and matches the heuristic motivation.
• Retaining qk-layernorm: Not necessary at the scales tested but included as a conservative choice based on prior scaling work showing it prevents instability in larger models.
• No gating: The gated attention unit from Hua et al. (2022) adds 9.3% computational overhead without improving accuracy at the optimal $\alpha$, so it is not included in the main ReLU-attention configuration.
• Unmodified training hyperparameters: The entire experiment uses existing BigVision configurations without any tuning, demonstrating that ReLU-attention is a drop-in replacement that does not require per-task hyperparameter search. This is a deliberate design choice to isolate the effect of the attention modification from confounding hyperparameter changes.

4. Key Insights and Innovations

Innovation 1: Sequence-Length Scaling as the Missing Ingredient That Makes Pointwise Activations Viable in Attention

The paper's most distinctive conceptual contribution is diagnostic rather than inventive: it identifies why prior attempts to replace softmax with pointwise activations failed, and shows that the fix is a single scaling factor — division by sequence length $L$. This is not a new activation function or a new attention mechanism, but rather the recognition that a seemingly minor implementation detail (the $L^{-1}$ factor) bridges the gap between a known idea (ReLU in attention) and competitive performance at scale.

What prior work missed. The paper cites multiple precedents that tried ReLU or squared ReLU as softmax replacements (Shen et al., 2023; Hron et al., 2020; Hua et al., 2022; Koohpayegani and Pirsiavash, 2022). These works all observed accuracy degradation. The dominant response in the field was to develop more elaborate solutions: gating mechanisms (Hua et al., 2022), scale-invariant normalization (Li et al., 2022), or full linear attention formulations that reorganize the computation (Katharopoulos et al., 2020; Lu et al., 2021). The implicit assumption was that pointwise activations simply could not reproduce the functional properties that make softmax attention effective.

This paper's diagnostic is that the problem was never about the activation function's shape — it was about the expected magnitude of the attention weights. Softmax enforces $\sum_j \alpha_{ij} = 1$, which implies $\mathbb{E}_j[\alpha_{ij}] = L^{-1}$. A naive pointwise activation produces $\mathbb{E}_j[\alpha_{ij}] = O(1)$, making each attention weight roughly $L$ times larger than what the rest of the network — LayerNorm, residual connections, learning rates — was designed to handle. The paper's insight is that this magnitude mismatch, not the shape of the activation, was responsible for prior failures. The evidence for this interpretation is Figure 2: when the scaling exponent $\alpha$ is swept, performance rises monotonically from $\alpha = 0$ (the prior work setting) to $\alpha \approx 1$ (the softmax-matched scale). The fact that this pattern holds across seven different activation functions with qualitatively different shapes — from hard-thresholding ReLU to smooth sigmoid to unbounded identity — strongly supports the claim that scale, not shape, is the dominant factor.

Why this reframes the problem. Before this paper, the question was "can we find an alternative to softmax that works?" After this paper, the question becomes "what scale properties must an attention function satisfy, independent of its functional form?" The $L^{-1}$ factor is not presented as a derived optimum (the paper is explicit about lacking a theoretical explanation) but as an empirical regularity that points toward a deeper principle: the transformer's training dynamics are tuned for attention weights of expected magnitude $O(L^{-1})$, and any violation of this scale — regardless of the activation function — will require compensatory hyperparameter changes. This reframing matters because it suggests that future work on alternative attention mechanisms should consider expected weight magnitude as a first-class design constraint, not an afterthought.

The paper positions this as a fundamental insight with modest scope: it is not claiming to have solved all problems with softmax-free attention, only to have identified the critical missing piece that prior applied work overlooked. The concurrent theoretical work by Bai et al. (2023) and Fu et al. (2023), which independently studied ReLU-attention with $L^{-1}$ scaling from a statistical learning perspective, reinforces the claim that the $L^{-1}$ factor is not an ad-hoc trick but emerges naturally from the structure of attention.

---

Innovation 2: Demonstrating That Softmax Is Functionally Replaceable at Scale Without Architectural Changes

The paper's second contribution is empirical proof-of-concept: it demonstrates that a pointwise activation can approach or match softmax attention's scaling behavior across a range of vision transformer sizes, from S/32 through L/16, on ImageNet-21k training, without modifying any hyperparameters from the standard training recipe. This is significant not because it introduces a new method, but because it establishes a baseline that changes the cost-benefit calculus for future attention research.

What the field assumed before this work. The transformer architecture has been remarkably resistant to modifications of its core attention mechanism. Changes to the attention function — whether linear attention, sparse attention patterns, or activation substitutions — typically required either accepting an accuracy penalty, modifying other architectural components (gating, normalization), or retuning hyperparameters extensively. The de facto assumption was that softmax's specific properties — positivity, normalization to sum-to-one, exponential sharpening — were functionally necessary for training stability and representational power. Works that achieved competitive performance by modifying attention (e.g., Hua et al., 2022) did so by adding new components (gates), not by simplifying the attention computation.

This paper challenges that assumption directly. The main result (Figure 1) shows that ReLU-attention with $L^{-1}$ scaling produces scaling curves — accuracy as a function of compute — that are essentially indistinguishable from softmax attention across six model configurations. The models range from S/32 (approximately 50 tokens, small capacity) to L/16 (approximately 197 tokens, large capacity), spanning roughly an order of magnitude in both sequence length and model size. The fact that the same $L^{-1}$ factor works across this range, with no per-model tuning, suggests that the scaling is robust to changes in $L$ and model depth.

Why this is more than a metric-equivalence claim. The paper is careful not to overclaim. It does not assert that ReLU-attention is better than softmax — only that it can "approach or match" softmax. The computational advantage (fewer gather operations for parallelization) is stated as an "opportunity," not a realized gain in the experiments (the paper reports core hours, not throughput or latency improvements at deployment). The contribution is therefore primarily epistemic: it demonstrates that softmax is not computationally irreducible. The exponential and the normalization sum — the two operations that make softmax expensive and hard to parallelize — are not functionally necessary for transformers to learn effective visual representations at the scales tested.

This distinguishes the paper from the linear attention literature (Katharopoulos et al., 2020; Lu et al., 2021), which also removes softmax but changes the computational structure (reordering matrix multiplications to achieve $O(d^2L)$ rather than $O(dL^2)$ complexity). Those methods face a different tradeoff: they gain asymptotic efficiency at the cost of representational limitations (linear attention cannot express certain attention patterns that quadratic attention can). ReLU-attention makes no such representational compromise — it computes the full $QK^\top$ matrix, just with a different activation — so the only potential loss is from the activation change, which the paper shows is recoverable via the $L^{-1}$ scaling.

The significance of hyperparameter-free substitution. The paper emphasizes that it uses existing BigVision training configurations "without modifying hyperparameters." This is a deliberate methodological choice that strengthens the claim. If the paper had achieved matching performance only after tuning learning rates, weight decay, or augmentation parameters specifically for ReLU-attention, the result would be less convincing — the gain might be attributable to improved hyperparameter tuning rather than the attention modification itself. By demonstrating drop-in compatibility, the paper shows that ReLU-attention does not require special treatment during training, which substantially lowers the barrier to adoption and makes the claim of functional equivalence more credible.

---

Innovation 3: Empirically Separating Activation Shape, Sequence-Length Scaling, and Auxiliary Components as Independent Factors

The paper's ablation structure (Figures 2–4) represents a methodological contribution: it disentangles three factors that prior work had conflated — the choice of nonlinearity $h$, the presence of sequence-length scaling $L^{-\alpha}$, and auxiliary architectural components (qk-layernorm, gating). By sweeping the scaling exponent $\alpha$ independently for each activation and each auxiliary component configuration, the paper produces a 2D map of the design space that reveals which factors matter and which do not.

What makes this distinctive. Prior work on softmax replacement typically tested specific combinations: Hua et al. (2022) tested squared ReLU with gating but without sequence-length scaling; Shen et al. (2023) tested ReLU and softmax variants but without systematic scaling sweeps; Koohpayegani and Pirsiavash (2022) replaced softmax with normalization but kept the gather operation. From these studies, one could not tell whether the failure of a pointwise activation was due to the activation itself, the lack of scaling, or the absence of auxiliary components. The field lacked a systematic map of which interventions compensated for which properties of softmax.

This paper provides that map. The key finding — visible in Figure 2 — is that the scaling exponent $\alpha$ dominates: performance curves for all seven activations rise together as $\alpha \to 1$, and at $\alpha \approx 1$, the activations converge to similar accuracy. This establishes a hierarchy of importance: get the scaling right first (activate at $\alpha \approx 1$), then choose any reasonable activation (ReLU for speed), and only then consider auxiliary components (qk-layernorm is neutral, gating is unnecessary and costly). The paper does not claim this hierarchy is exhaustive or general — it tests only vision transformers on ImageNet-scale data — but within its scope, it provides a clear prioritization.

The gating ablation (Figure 4) as a negative result with positive implications. The finding that gating "does not remove the need for sequence length scaling" and adds 9.3% computational overhead for the S/8 model is a cleanly demonstrated negative result. Prior work (Hua et al., 2022) had proposed gating as the mechanism for making softmax-free attention work. This paper shows that gating addresses a different problem than the one that actually matters — the $L^{-1}$ scaling handles the magnitude mismatch, and once that is fixed, gating provides no additional benefit that justifies its cost. This is the kind of result that prevents the field from pursuing red herrings: if you are working on softmax-free attention, focus on the scale properties, not on adding learned modulation mechanisms.

The qk-layernorm ablation (Figure 3) as a robustness check. The finding that removing qk-layernorm does not substantially affect ReLU-attention's performance at the tested scales is important for understanding why the $L^{-1}$ scaling works. QK-layernorm reduces the variance of the dot products entering the activation, which controls the fraction of dot products that pass the ReLU threshold (positive vs. negative). The insensitivity to qk-layernorm suggests that ReLU-attention is tolerant of a range of positivity fractions — the $L^{-1}$ scaling adapts the overall magnitude regardless of how many weights survive the ReLU threshold. This robustness is practically valuable because it means the attention mechanism does not depend on precise initialization or normalization choices to function.

Limitations of this disentanglement. The paper explicitly notes that "this may change at scale" for the qk-layernorm ablation. At larger model sizes (beyond L/16) or longer sequences (beyond 197 tokens), the dot product distribution may shift sufficiently that qk-layernorm becomes necessary even for ReLU-attention. The paper does not test this regime, so the hierarchy it establishes should be understood as validated only for the tested configuration range. Similarly, the convergence of activation functions at $\alpha \approx 1$ is tested on ImageNet classification; tasks with different signal characteristics (e.g., retrieval, where cosine similarity matters more than magnitude) might favor different activations.

---

Innovation 4: A Simplicity Result That Converts a Computational Bottleneck Into an Engineering Problem

The paper's most practically significant contribution is the demonstration that a single division — $L^{-1}$ — is sufficient to close the softmax-vs-pointwise performance gap. This is a simplicity result: it shows that the solution to a long-standing problem (softmax is expensive) is not a complex new mechanism but a simple scaling factor that was overlooked because the field focused on the wrong axis of variation (activation shape rather than weight magnitude).

Why simplicity is a contribution. In machine learning, the simplest effective solution is often the most impactful because it lowers the barrier to adoption and reduces the surface area for bugs, tuning, and unintended interactions. Prior approaches to softmax-free attention — gating (Hua et al., 2022), linear attention reformulations (Katharopoulos et al., 2020), scale-invariant normalization (Li et al., 2022) — each added complexity to the architecture or training procedure. This paper's solution adds exactly one multiply per attention head: the $L^{-1}$ factor. It does not change the computation graph structure, does not introduce new learned parameters, and does not require any training recipe changes. This minimality means ReLU-attention is not a new architecture to implement, debug, and tune — it is a one-line code change in an existing transformer implementation.

The significance extends beyond vision. Transformers are used across modalities — language, audio, reinforcement learning — and softmax is a bottleneck in all of them. A solution that is specific to vision (e.g., relying on 2D patch structure) would have limited impact. A solution that works because of a general property of attention (the expected per-position weight scale) is modality-agnostic. The paper does not test non-vision domains, but the heuristic motivation — that $\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$ under softmax and this should be preserved — makes no assumptions about the input modality. The $L^{-1}$ factor depends only on sequence length, which is domain-universal.

The open questions as part of the contribution. The paper's conclusion enumerates what it does not answer: "why the factor $L^{-1}$ improves performance," "if this term could be learned," and whether "there is a better activation function." This explicit acknowledgment of incompleteness is not a weakness — it is a contribution of its own, because it defines the research frontier. Before this paper, the open question was "can anything replace softmax without hurting accuracy?" After this paper, the open questions are more specific and actionable: "what is the optimal scaling exponent for ultra-long sequences?" "Can the $L^{-1}$ factor be learned from data as a function of sequence position rather than being a fixed constant?" "Is there an activation function that outperforms ReLU at $\alpha = 1$ by a margin that justifies its additional cost?" By providing a simple, working baseline and a set of systematically explored ablations, the paper gives future work a clear starting point and a well-defined set of questions to pursue, rather than an open-ended design space to explore from scratch.

What this is not. The paper is explicitly a report ("this report"), not a full conference paper with theoretical analysis, comprehensive benchmarks, and downstream deployment measurements. It does not claim to have optimized ReLU-attention for maximum throughput, nor does it compare against the full space of efficient attention methods (FlashAttention-style implementations, sparse attention, linear attention optimizations). Its contribution is narrower and more fundamental: establishing that the activation can be changed without architectural upheaval, and identifying the one factor ($L^{-1}$) that makes the difference. This is an enabling result — it opens a direction rather than solving it completely — and the paper's modest framing ("This report leaves many open questions") accurately reflects this scope.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The primary experiments use ImageNet-21k (Deng et al., 2009) for training, consisting of approximately 14 million images across 21,841 classes. Ablation experiments additionally use ImageNet-1k (approximately 1.28 million images, 1,000 classes). The paper uses the standard splits from the BigVision codebase (Beyer et al., 2022) without modification.

• Base model(s). All experiments use the Vision Transformer (ViT) architecture (Dosovitskiy et al., 2021) as implemented in the BigVision codebase. Six configurations are tested spanning model scale and patch size: S/32, S/16, S/8, B/32, B/16, and L/16, where "S," "B," and "L" denote small, base, and large variants, and the number indicates patch size (e.g., S/16 uses 16×16 patches). These span sequence lengths from approximately 50 tokens (S/32, B/32) to 785 tokens (S/8).

• Metrics. The primary metric is ImageNet-1k top-1 classification accuracy (%), measured in two ways: (1) for ImageNet-21k-trained models, accuracy is computed by taking "the top class among those that are in ImageNet-1k, without fine-tuning" — effectively measuring zero-shot transfer performance within the 1,000-class subset; (2) for ImageNet-1k-trained models, standard top-1 accuracy on the validation set. A secondary metric is average 10-shot linear transfer accuracy across eight downstream datasets (Caltech Birds, Caltech-101, Stanford Cars, CIFAR-100, DTD, ColHist, Pets, UC Merced), computed by training a linear classifier on frozen ViT features with 10 labeled examples per class, averaged over three random seeds corresponding to different selections of training examples.

• Baselines. The sole baseline is standard softmax attention as used in the original vision transformer (Dosovitskiy et al., 2021). The paper compares ReLU-attention against this softmax baseline across the full range of model sizes and both evaluation metrics.

• Generation budget / compute accounting. The paper measures computational cost using TPU core hours — the product of wall-clock time and the number of TPU cores used during training. This is reported on the x-axis of Figure 1 and accounts for actual training cost including any overhead from the attention mechanism. The paper quantifies that gated attention "increases the core hours required for the experiment by roughly 9.3% for the S/8 model with ReLU," confirming that compute is measured empirically from training runs rather than estimated from theoretical FLOP counts.

• Cross-validation / statistical protocol. No cross-validation is used. The paper evaluates on the standard ImageNet-1k validation set (or the ImageNet-21k-to-ImageNet-1k zero-shot transfer, which uses the ImageNet-1k validation set). For the 10-shot linear transfer experiments, results are averaged over three random seeds corresponding to different random selections of the 10 training examples per class. No error bars or confidence intervals are reported in any figure.

Main Quantitative Results

ReLU-Attention Matches Softmax Scaling Behavior Across Model Sizes (Figure 1)

The central result of the paper appears in Figure 1, which plots accuracy against TPU core hours for both softmax attention and ReLU-attention (using $L^{-1} \cdot \text{ReLU}$) across six model configurations (S/32, S/16, B/32, B/16, L/16 on ImageNet-1k accuracy; the same plus S/8 on 10-shot linear transfer).

On ImageNet-1k accuracy (Figure 1, left panel), the curves for softmax and ReLU-attention largely overlap. The smallest model (S/32) achieves approximately 0.68–0.70 accuracy for both methods at ~$10^2$ TPU core hours. As compute increases through S/16, B/32, B/16, and L/16, both methods scale upward in near-lockstep, with the L/16 model reaching approximately 0.80–0.82 accuracy. The paper reports that ReLU-attention "can approach or match the performance of softmax-attention in terms of scaling behavior as a function of compute."

On the 8-dataset 10-shot linear transfer metric (Figure 1, right panel), the same pattern holds. Softmax and ReLU-attention produce nearly identical scaling curves, from approximately 0.72–0.73 average accuracy at S/32 and S/16 to approximately 0.83–0.84 at L/16. The S/8 model (highest compute point, at roughly $10^{2.2}$ TPU core hours) achieves the highest transfer accuracy for both methods.

The paper states these results qualitatively: "ReLU-attention matches the scaling trends for softmax attention for ImageNet-21k training." There is no quantitative claim of outperformance — only that the performance is comparable. The figure does not report numerical values at specific compute points, so exact differences between softmax and ReLU-attention at each model size cannot be determined from the plot alone. The visual impression is of near-indistinguishable scaling behavior.

Sequence-Length Scaling Exponent Sweep: $\alpha \approx 1$ Is Optimal Across All Activations (Figure 2)

Figure 2 presents a systematic sweep over the scaling exponent $\alpha$ in the general form $\phi = L^{-\alpha} \cdot h$, testing seven activation functions ($h \in \{\text{relu}, \text{relu}^2, \text{gelu}, \text{softplus}, \text{identity}, \text{relu6}, \text{sigmoid}\}$) on the S/32, S/16, and S/8 models trained on both ImageNet-21k and ImageNet-1k. The x-axis sweeps $\alpha$ from 0 to 1 (with some sweeps extending to 2.0), and the y-axis shows ImageNet-1k accuracy.

Across all 36 sub-panels (7 activations × 3 model configurations × 2 training datasets, minus some combinations), the following pattern is consistent:

• At $\alpha = 0$ (no sequence-length scaling), all activations substantially underperform softmax. The softmax baseline is plotted as a horizontal reference line (where applicable). For the S/16 model trained on i21k with ReLU, accuracy at $\alpha = 0$ is roughly 0.65 versus approximately 0.76 for softmax. For identity activation at $\alpha = 0$, accuracy drops to approximately 0.64 on the same configuration — the worst of all tested activations.

• As $\alpha$ increases toward 1, accuracy rises monotonically for all activations. The curves show a smooth upward trend, with some activations (ReLU, GELU, softplus) approaching the softmax reference line at $\alpha \approx 1$. The paper states: "We typically observe the best results when $\alpha$ is close to 1."

• At $\alpha \approx 1$, all activations (except identity) converge to similar accuracy levels. For the i21k-trained S/16 model at $\alpha = 1$, ReLU, squared ReLU, GELU, softplus, ReLU6, and sigmoid all achieve roughly 0.75–0.76, close to the softmax reference. Identity activation lags significantly, remaining around 0.65–0.70 even at $\alpha = 1$ on most configurations.

• The relationship between $\alpha$ and accuracy is not perfectly monotonic for all activations. For squared ReLU and sigmoid on some configurations, the curves show a slight dip or flattening before continuing to rise toward $\alpha = 1$, suggesting some sensitivity to the interaction between activation shape and scaling exponent.

The key interpretive finding is stated explicitly: "There is no clear best non-linearity at $\alpha \approx 1$, so we use ReLU in our main experiment for its speed." This empirical convergence across activation shapes at the optimal scaling exponent is the paper's central evidence that scale, not shape, is the dominant factor.

QK-LayerNorm Has Minimal Effect on ReLU-Attention Performance at Tested Scales (Figure 3)

Figure 3 ablates the effect of qk-layernorm (Dehghani et al., 2023) for ReLU and squared ReLU activations across the $\alpha$ sweep, using the S/32, S/16, and S/8 models trained on i21k.

For both activations, the curves with and without qk-layernorm are closely overlapping across the full $\alpha$ range. For the S/16 model with ReLU, the with-qk-layernorm and without-qk-layernorm curves are nearly indistinguishable, both rising from roughly 0.65 at $\alpha = 0$ to approximately 0.76 at $\alpha = 1$. For the S/8 model with ReLU, a small separation appears at intermediate $\alpha$ values (roughly 0.5–0.8), but the curves re-converge at $\alpha = 1$.

The paper concludes: "qk-layernorm does not have a large effect for these models, but this may change at scale." The caveat is important — the largest model tested in this ablation is S/8 (highest sequence length but still "small" in capacity), not the B/16 or L/16 configurations from Figure 1 where qk-layernorm might matter more.

Gated Attention Units Do Not Eliminate the Need for Sequence-Length Scaling (Figure 4)

Figure 4 tests the interaction between gated attention units (Hua et al., 2022) and sequence-length scaling for ReLU and squared ReLU activations. The experiment uses the S/32, S/16, and S/8 models trained on i21k, sweeping $\alpha$ from 0 to 1 for four conditions: ReLU with and without gating, and squared ReLU with and without gating.

The finding is that "the best accuracy is still achieved with sequence length scaling, with or without the gate." For the S/16 model with ReLU, gating shifts the curve upward slightly at low $\alpha$ values (e.g., with gate at $\alpha = 0$ achieves roughly 0.68 vs. 0.65 without gate), but at $\alpha = 1$, both curves converge to approximately the same accuracy (roughly 0.76). The same pattern holds for squared ReLU and across model configurations.

The practical implication is that gating does not substitute for the $L^{-1}$ scaling — it provides a small benefit when scaling is absent but no benefit when scaling is correctly set. The paper further notes that "gating increases the core hours required for the experiment by roughly 9.3% for the S/8 model with ReLU," making it an unfavorable cost-accuracy tradeoff given the null benefit at $\alpha = 1$.

ImageNet-1k Training Confirms the Pattern Generalizes Across Datasets

The right column of Figure 2 and the corresponding Figure 3/4 sub-panels for i1k training show qualitatively identical patterns to the i21k results. At $\alpha = 0$, activations underperform softmax; at $\alpha \approx 1$, they approach softmax performance; and at $\alpha \approx 1$, activation choice makes little difference (except identity, which consistently underperforms). The i1k-trained models achieve lower absolute accuracy than i21k-trained models at the same $\alpha$ (roughly 0.68–0.72 vs. 0.74–0.78 for S/16 depending on activation), reflecting the smaller training set, but the relative pattern of scaling exponent and activation function effects is preserved.

This cross-dataset consistency strengthens the paper's claim that the $L^{-1}$ scaling is not an artifact of a particular training distribution or dataset size — it generalizes across both a large-scale classification dataset (21k classes, 14M images) and a smaller one (1k classes, 1.28M images).

Ablation Studies and Robustness Checks

• Activation function choice (Figure 2): Seven activation functions tested. At the optimal $\alpha \approx 1$, all nonlinear activations (ReLU, squared ReLU, GELU, softplus, ReLU6, sigmoid) achieve similar accuracy, approximately matching softmax. The identity activation (no nonlinearity) consistently underperforms by a substantial margin, confirming that some nonlinearity is necessary but its precise form matters little once scaling is correct.

• Scaling exponent sensitivity (Figure 2): The relationship between $\alpha$ and accuracy is smooth and monotonic across $\alpha \in [0, 1]$ for most activation-model-dataset combinations. There is no sharp phase transition — performance degrades gradually as $\alpha$ moves away from 1. At $\alpha = 0.5$, accuracy is typically intermediate between the $\alpha = 0$ and $\alpha = 1$ extremes, losing roughly half the gap to softmax. This smoothness suggests the scaling factor can be set to 1 without precise tuning — small deviations cause small losses.

• QK-LayerNorm removal (Figure 3): Removing qk-layernorm has minimal impact on ReLU-attention and squared ReLU-attention at the tested scales (S/32, S/16, S/8). The curves with and without qk-layernorm closely track each other across the full $\alpha$ range. This is a non-obvious result because qk-layernorm was introduced specifically to stabilize training of larger models — its irrelevance here suggests the $L^{-1}$ scaling already provides sufficient normalization for the attention weights at these scales.

• Gating interaction (Figure 4): Adding a gated attention unit does not substitute for sequence-length scaling. Gating provides a small accuracy improvement at low $\alpha$ (where scaling is absent) but none at $\alpha = 1$ (where scaling is correct). The computational cost of gating (~9.3% additional core hours) is not justified by accuracy gains at the optimal scaling exponent. This result invalidates the hypothesis from Hua et al. (2022) that gating alone can replace softmax's functional role.

• Model scale and sequence length robustness (Figures 1–4): The $L^{-1}$ scaling is tested across sequence lengths from approximately 50 tokens (S/32, B/32) to 785 tokens (S/8) and model capacities from small (S) to large (L). The optimal $\alpha \approx 1$ holds across this range, suggesting the scaling factor generalizes to varying $L$ without needing to be adapted. This is important because $L^{-1}$ explicitly depends on sequence length — the fact that the same formula works at $L = 50$ and $L = 785$ without modification supports the heuristic that $\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$ is the relevant invariant.

• Dataset size robustness (Figure 2, right vs. left columns): The pattern of results transfers from ImageNet-21k (14M images) to ImageNet-1k (1.28M images). Absolute accuracies are lower on i1k (due to less training data), but the relationship between $\alpha$, activation choice, and relative performance vs. softmax is preserved.

• Downstream transfer robustness (Figure 1, right panel): The compute-accuracy scaling equivalence between softmax and ReLU-attention holds not only for ImageNet-1k accuracy (which could be influenced by the training objective) but also for 10-shot linear transfer across eight diverse downstream datasets. This suggests that ReLU-attention does not produce representations that are qualitatively different from softmax attention — they transfer similarly well to new tasks.

Critical Assessment

Does the Paper Demonstrate Matching Scaling Behavior?

The central claim is that ReLU-attention "can approach or match the performance of softmax-attention in terms of scaling behavior as a function of compute." Figure 1 provides visual evidence for this claim: the softmax and ReLU-attention curves largely overlap across six model configurations. However, the claim is demonstrated only visually — the paper does not report numerical accuracy values at specific compute points, does not compute differences between softmax and ReLU-attention at each model size, and does not provide confidence intervals or error bars. A reader cannot determine whether the gap at any particular model size is 0.1%, 1%, or larger. The curves "look" close together, but "approaches or matches" is a qualitative judgment based on visual inspection of log-scale plots, which can compress visible differences.

This matters because the computational advantage of ReLU-attention (fewer gather operations) is an implementation-level benefit that depends on deployment hardware and software optimization. The paper measures compute as "TPU core hours" for training, but this does not directly measure the specific operations that ReLU-attention is supposed to eliminate (exponentiation, gather for normalization sum). Training time is dominated by many factors beyond the attention softmax — matrix multiplications, data loading, communication — so equivalence in total training core hours does not directly imply equivalence in the attention operation's cost profile. An experiment measuring wall-clock time per forward pass or per attention operation, at deployment scale, would strengthen the computational motivation more directly than training time curves.

Is the Hyperparameter-Free Claim Fully Substantiated?

The paper emphasizes that it uses existing BigVision configurations "without modifying hyperparameters." This is a strength — it shows ReLU-attention works as a drop-in replacement in a well-tuned training recipe. However, the existing hyperparameters were tuned for softmax attention. If softmax attention were replaced with ReLU-attention and then hyperparameters were re-tuned for the new mechanism, it is possible (though not demonstrated) that ReLU-attention could outperform softmax. Alternatively, the existing hyperparameters might be suboptimal for ReLU-attention in ways that only appear at larger scales, different architectures, or different tasks.

The claim should therefore be understood as: ReLU-attention can replace softmax without degrading performance when other hyperparameters remain fixed at their softmax-optimized values. It does not demonstrate that ReLU-attention is at its best under these hyperparameters, nor that softmax would not regain an advantage under per-method hyperparameter optimization. The paper's framing as a simplicity result makes this limitation acceptable — the point is that no special tuning is needed, not that further tuning wouldn't help.

What the Paper Does Not Test

No latency or throughput measurements. The paper's compute metric (TPU core hours) measures total training cost but does not report inference-time latency, throughput, or memory usage. The stated advantage of ReLU-attention — "parallelized over the sequence length dimension with fewer gather operations" — pertains to deployment efficiency, not training. This advantage is never quantified. A reader knows that ReLU-attention should be faster in principle but does not learn how much faster at what batch sizes and hardware configurations.

No comparison to optimized softmax implementations. The paper does not compare ReLU-attention against softmax attention as implemented in systems like FlashAttention (Dao et al., 2022), which fuse the softmax computation with the attention matrix multiply to reduce memory overhead. Optimized softmax implementations already address some of the parallelism and memory bottlenecks that motivate ReLU-attention. Without this comparison, the practical compute advantage of ReLU-attention over a well-engineered softmax baseline is unknown.

No test beyond ImageNet-scale vision tasks. The models are vision transformers trained on ImageNet-21k/1k and evaluated on standard transfer datasets. The paper does not test language modeling, multimodal tasks, reinforcement learning, or any other domain where transformers are used. The heuristic motivation ($\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$ under softmax) is domain-agnostic, so the approach should in principle transfer, but the specific scaling properties and the optimal $\alpha$ might differ. Without cross-domain evidence, the result is bounded to vision transformers on image classification.

No test at larger model scales. The largest model is L/16, which is "large" by 2023 ViT standards but far from the frontier scale of models with billions of parameters. The qk-layernorm ablation explicitly caveats that findings "may change at scale," and the same caution should apply to the overall equivalence claim. At scales where training instability becomes the dominant challenge, the absence of softmax's bounded, normalized outputs might interact differently with other stability mechanisms.

No test of sequence lengths beyond 785 tokens. The longest sequence tested is S/8 (785 tokens). Many applications that would benefit most from parallelized attention — video, high-resolution images, long documents — involve sequence lengths of thousands or tens of thousands of tokens. At such lengths, the $L^{-1}$ scaling factor becomes very small (e.g., $10^{-4}$ for $L = 10,000$), and the ReLU threshold might eliminate a different fraction of attention weights than at $L = 200$. The behavior of ReLU-attention in this regime is entirely unexplored.

Are the Ablations Comprehensive?

The ablation structure is well-designed for the paper's scope: it disentangles scaling exponent, activation choice, qk-layernorm, and gating. However, several ablations that would strengthen the central claim are absent:

• No ablation over the scaling function form. The paper tests only $L^{-\alpha}$ scaling. Alternative scaling functions — such as dividing by $\sum_j h(q_i^\top k_j / \sqrt{d})$ (a normalization that doesn't require a gather if computed incrementally), or scaling by a learned parameter, or using a different function of $L$ — are not explored. The claim that $L^{-1}$ is the "right" scaling is supported by the $\alpha$ sweep but not by comparison to alternative normalization strategies.

• No ablation over initialization or optimizer settings. The heuristic argument for $L^{-1}$ relies on initialization-scale reasoning ($q$, $k$ elements are $O(1)$). The paper does not test whether different initialization schemes (e.g., varying the standard deviation of weight initialization) change the optimal $\alpha$, which would directly test whether the $O(L^{-1})$ expected weight argument is correct.

• No ablation over number of attention heads. Multi-head attention divides the embedding dimension across heads, changing the effective $d$ per head. The paper uses standard head counts from the BigVision configurations but does not test whether the $L^{-1}$ scaling is robust to varying the number of heads and thus the per-head dimension.

• No training trajectory analysis. The paper evaluates only final accuracy. It does not show training curves (loss or accuracy over steps) to demonstrate that ReLU-attention trains as stably as softmax attention. The $O(L^{-1})$ argument is about initialization-time scale, but the relevant dynamics occur throughout training. Analyzing the evolution of attention weight statistics during training for both methods would provide stronger evidence for the scale-matching hypothesis.

Consistency Between Heuristic and Evidence

The paper's analytical motivation — preserving $\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$ at initialization — predicts that $\alpha = 1$ should be optimal. The empirical sweep in Figure 2 confirms this prediction. This consistency between heuristic and data is a genuine strength. However, the heuristic does not predict that all activations converge at $\alpha = 1$ — that is an additional empirical finding that strengthens the case for scale dominance but is not explained by the simple expected-value argument. The heuristic also does not explain why identity activation fails even at $\alpha = 1$: identity preserves $O(1)$ and with $L^{-1}$ scaling achieves $O(L^{-1})$ expected weight, so by the heuristic it should work. That identity consistently underperforms suggests additional functional requirements beyond expected weight magnitude — perhaps the ability to produce zero or near-zero weights for irrelevant token pairs, or the importance of non-negativity, or the need for a bounded range below some threshold. The paper does not investigate this discrepancy.

Overall Assessment

The experiments demonstrate that, for vision transformers at the tested scales on ImageNet, a pointwise activation with $L^{-1}$ scaling can approximately match the accuracy scaling behavior of softmax attention. The case for the $L^{-1}$ factor as the critical missing ingredient is well-supported by the systematic $\alpha$ sweep showing monotonic improvement from the prior-work setting ($\alpha = 0$) to the softmax-matched setting ($\alpha \approx 1$). The finding that activation choice matters little at the optimal scaling exponent is robustly demonstrated across seven activations, three model configurations, and two datasets.

However, the practical claim — that this enables better parallelization — is not directly tested. The paper demonstrates accuracy equivalence under a training-time compute metric, not throughput or latency improvements under deployment conditions. The gap between "matching scaling behavior" (demonstrated) and "enabling faster inference through improved parallelization" (claimed as motivation) requires experiments not present in this report. The paper's modest framing ("This report leaves many open questions") correctly acknowledges this gap, but a reader should understand that the primary demonstrated contribution is diagnostic (identifying $L^{-1}$ scaling as the missing ingredient) rather than engineering (demonstrating wall-clock speedups from ReLU-attention in a deployment setting).

6. Limitations and Trade-offs

Limitation 1: The Practical Computational Advantage Is Never Measured — Only Training Cost, Not Inference Throughput or Latency

The assumption or constraint. The paper's entire motivation is computational: softmax is costly "due to an exponent calculation and a sum over sequence length which makes parallelization challenging," and ReLU-attention "can be parallelized over the sequence length dimension with fewer gather operations than traditional attention" (Section 1, Figure 1 caption). The paper measures training cost in TPU core hours (Figure 1 x-axis) but never measures the quantity that matters for the stated motivation — inference-time throughput, latency, or memory usage in a deployment setting. The paper does not report wall-clock time per forward pass, tokens processed per second, or peak memory consumption for either softmax-attention or ReLU-attention.

The consequence. A practitioner evaluating whether to adopt ReLU-attention cannot determine how much faster or more memory-efficient it actually is — only that it costs roughly the same to train to equivalent accuracy. Training TPU core hours conflate the cost of the attention operation with the cost of data loading, MLP layers, communication, and many other factors. The parallelization benefit of ReLU-attention (fewer gather operations across the sequence length dimension) is a property of the forward pass computational graph, not of total training time. It is entirely possible that ReLU-attention provides negligible throughput improvement at commonly used batch sizes and sequence lengths because: (1) the softmax may not be the dominant cost in the forward pass relative to the matrix multiplications in the attention and MLP layers; (2) the gather operation may already be well-optimized in accelerator hardware; or (3) the benefit may only manifest at inference time with specific sharding strategies, not during data-parallel training. Without throughput or latency measurements, the paper demonstrates functional equivalence but not the practical advantage that justifies replacing softmax in the first place.

What evidence exists in the paper. The paper provides only one concrete number related to computational cost beyond the TPU core hour curves: gating "increases the core hours required for the experiment by roughly 9.3% for the S/8 model with ReLU" (Section 4, Figure 4 caption). This measures the training-time overhead of the gating mechanism, but it does not decompose how much of the remaining ~90.7% is the ReLU-attention operation itself versus other components of training. There is no analogous number comparing softmax-attention training time to ReLU-attention training time for the same model configuration — the training time equivalence is presented graphically in Figure 1 without numerical quantification. There is no experiment that measures throughput or latency at inference time for any batch size or hardware configuration.

Mitigation status. The paper does not attempt to address this limitation. It does not report any inference benchmarks, does not compare against optimized softmax implementations (e.g., FlashAttention, which fuses the softmax computation to reduce memory overhead), and does not characterize the conditions (batch size, sequence length, parallelism strategy) under which the gather-operation savings would translate to wall-clock speedups. The paper frames the computational advantage as an "opportunity" (Section 1: "This result presents new opportunities for parallelization") rather than a demonstrated benefit, which is an honest framing but leaves the central practical question unanswered.

---

Limitation 2: The $L^{-1}$ Scaling Factor Is Empirically Motivated but Not Theoretically Understood, and the Paper Explicitly States It Does Not Know Why It Works

The assumption or constraint. The paper's analytical motivation for dividing by sequence length is a heuristic: softmax enforces $\sum_j \alpha_{ij} = 1$, implying $\mathbb{E}_j[\alpha_{ij}] = L^{-1}$; a pointwise activation produces $\alpha_{ij} = O(1)$; dividing by $L$ restores the $O(L^{-1})$ expected weight magnitude at initialization (Section 3). However, the paper explicitly acknowledges this is not a rigorous justification. The conclusion states directly: "we are unsure why the factor $L^{-1}$ improves performance or if this term could be learned" (Section 5). The heuristic explains why $\alpha = 0$ fails (weights are too large) but does not explain why $\alpha = 1$ is specifically optimal — it does not rule out $\alpha = 0.8$ or $\alpha = 1.2$ as equally valid on theoretical grounds, and it does not explain why the identity activation with $\alpha = 1$ (which also satisfies $\mathbb{E}_j[\alpha_{ij}] = O(L^{-1})$ in expectation) consistently underperforms while ReLU succeeds.

The consequence. Without understanding why $L^{-1}$ works, a practitioner cannot predict when it will stop working. The heuristic depends on initialization-scale reasoning: at initialization, $q$ and $k$ elements are $O(1)$, the scaled dot product is $O(1)$, and ReLU preserves $O(1)$. But during training, queries and keys evolve — their scale may grow or shrink depending on the optimizer and the training dynamics. If the dot products become systematically larger or smaller than $O(1)$ during training, the $L^{-1}$ scaling may no longer preserve the softmax-matching expected weight magnitude, potentially causing instability or accuracy degradation partway through training. The paper evaluates only final accuracy, not training trajectory statistics, so a practitioner cannot know whether this is a real concern. Furthermore, if the $L^{-1}$ factor is compensating for a property that could be satisfied in other ways (e.g., by learning a per-head scaling parameter, or by modifying initialization rather than the attention formula), then the current approach may be suboptimal — but without understanding the mechanism, alternatives cannot be rationally designed.

What evidence exists in the paper. The $\alpha$ sweep in Figure 2 demonstrates empirically that $\alpha \approx 1$ is optimal, but this is a descriptive result, not an explanatory one. The paper does not measure attention weight statistics (mean, variance, sparsity) during training to test whether the $O(L^{-1})$ expected weight property actually holds throughout training or only at initialization. It does not test alternative explanations for why $L^{-1}$ works — for instance, whether the benefit comes from matching the expected magnitude of the output $o_i = \sum_j \alpha_{ij} v_j$ rather than the expected per-weight magnitude, or whether it prevents gradient vanishing/exploding in the attention softmax replacement. The failure of identity activation at $\alpha = 1$ (Figure 2) reveals that the $O(L^{-1})$ expected weight argument is incomplete — identity with $L^{-1}$ scaling should satisfy the heuristic but consistently underperforms, implying additional constraints the heuristic does not capture (e.g., the importance of non-negativity, or sparsity, or a bounded activation range). The paper does not investigate this discrepancy.

Mitigation status. The paper is fully transparent about this limitation. The conclusion states: "This report leaves many open questions. In particular, we are unsure why the factor $L^{-1}$ improves performance or if this term could be learned." This is an honest acknowledgment that the work is empirical and preliminary. The paper does not attempt to resolve the theoretical gap or to test whether a learned scaling factor (rather than the fixed $L^{-1}$) would perform better. The limitation is acknowledged but not mitigated.

---

Limitation 3: Results Are Demonstrated Only on Image Classification with a Single Model Architecture at Moderate Scale — No Evidence for Cross-Domain, Cross-Architecture, or Large-Scale Generalization

The assumption or constraint. All experiments use vision transformers (ViT) on ImageNet-21k and ImageNet-1k image classification (Section 4). The largest model tested is L/16, which is "large" for a ViT but far from the scale of frontier models with billions of parameters. The longest sequence length tested is 785 tokens (S/8), which is modest compared to sequence lengths in video, high-resolution imagery, or long-document language modeling where the parallelization benefit of ReLU-attention would be most impactful. The paper tests only one model architecture family (ViT with qk-layernorm) and makes no claims about cross-modal transfer to language, speech, or other domains where transformers are dominant.

The consequence. A practitioner considering ReLU-attention for a different domain, architecture, or scale faces substantial uncertainty. The $L^{-1}$ heuristic depends on the expected weight magnitude at initialization, which is influenced by the initialization scheme, the embedding dimension per head, the number of layers, and the presence of other normalization mechanisms — all of which differ across architectures and domains. In language models, sequence lengths routinely exceed 2,000 tokens (and can reach 100,000+ in long-context settings), and the attention patterns are qualitatively different from vision (causal masking, different sparsity patterns). At such lengths, $L^{-1}$ becomes very small (e.g., $5 \times 10^{-4}$ for $L = 2,000$, or $10^{-5}$ for $L = 100,000$), and the ReLU threshold may eliminate a much larger fraction of attention weights than at $L = 200$. Whether the $L^{-1}$ scaling remains optimal at these lengths is completely untested.

At larger model scales, the paper's own caveat about qk-layernorm applies more broadly: "this may change at scale" (Figure 3 caption). The L/16 model is the largest tested, and the paper does not report results for models with hundreds of millions or billions of parameters. Training instability — which qk-layernorm was introduced to prevent — is a dominant challenge at such scales. ReLU-attention removes softmax's bounded, normalized outputs, which could interact adversely with other stability mechanisms at scale. The paper's finding that qk-layernorm is unnecessary for ReLU-attention at the tested scales (Figure 3) may not hold at larger scales, but no evidence is provided either way.

What evidence exists in the paper. The paper tests six model configurations (S/32, S/16, S/8, B/32, B/16, L/16) spanning a moderate range of model capacities and sequence lengths (approximately 50 to 785 tokens). The consistent results across these configurations suggest robustness within this range, but the range itself is narrow relative to the full space of transformer deployments. The paper provides 10-shot linear transfer results on eight downstream datasets (Figure 1, right panel), which tests generalization beyond ImageNet classification but still within the domain of visual recognition. There are no experiments on language modeling, multimodal tasks, reinforcement learning, or any non-classification visual task (detection, segmentation, retrieval). The paper does not compare ViT variants (e.g., Swin, DeiT, MAE) to test whether the architectural details of the transformer influence the effectiveness of ReLU-attention.

Mitigation status. The paper does not attempt to address this limitation. It does not test language modeling, does not evaluate on tasks beyond image classification and linear transfer, and does not provide results for models larger than L/16. The domain scope is inherent in the paper's framing as a study of "vision transformers" (title), so the limitation is one of scope rather than oversight — but a practitioner deploying transformers outside vision cannot extrapolate from these results with confidence.

---

Limitation 4: The Paper Provides No Statistical Quantification of the Performance Gap Between ReLU-Attention and Softmax — Only Visual Comparisons Without Error Bars or Numerical Differences

The assumption or constraint. All main results are presented as line plots without error bars, confidence intervals, or numerical reporting of accuracy differences at specific compute points. Figure 1 shows softmax and ReLU-attention curves that visually overlap, but the paper never states the actual accuracy difference at any model size. Figure 2 sweeps $\alpha$ and shows many curves, but the softmax reference line (where present) is plotted without indication of variance, and the paper never reports whether the difference between the best ReLU-attention configuration and softmax is statistically significant or smaller than the run-to-run variance from random initialization.

The consequence. A practitioner evaluating the cost of switching from softmax to ReLU-attention needs to know whether the accuracy change is zero (within measurement noise), slightly negative (a small but consistent degradation), or occasionally positive. Visual inspection of log-scale plots can mask small but systematic differences. If ReLU-attention consistently underperforms softmax by 0.5–1.0% in absolute accuracy across configurations, this might be an acceptable tradeoff for some deployment scenarios (where throughput gains dominate) but unacceptable for others (where state-of-the-art accuracy is paramount). Without numerical differences, the practitioner cannot make this judgment. The qualitative claim "approaches or matches" is ambiguous — it could mean "statistically indistinguishable" or "within 2%," which have very different practical implications.

The absence of error bars also makes it impossible to assess whether the optimal $\alpha \approx 1$ finding is robust to random seed variation. The curves in Figure 2 are smooth, suggesting they may represent single runs or averages whose variance is unknown. If the peak at $\alpha = 1$ is only 0.3% above neighboring $\alpha$ values and run-to-run variance is 0.5%, the optimal $\alpha$ is not well-identified. The 10-shot linear transfer results are averaged over three random seeds (Section 4), which provides some variance reduction, but no standard deviations or confidence intervals are reported for these averages.

What evidence exists in the paper. The paper reports one quantitative cost number: gating adds "roughly 9.3%" core hours for the S/8 model with ReLU (Figure 4 caption). No analogous quantitative comparison is provided for softmax vs. ReLU-attention on either accuracy or computational cost. The paper states that the i21k and i1k training configurations both use "a roughly similar number of steps of around 9e5" (Section 4), which is a qualitative precision estimate rather than an exact step count. The figures contain multiple data points per curve, suggesting systematic sweeps, but the exact number of runs per point and the variance across runs are not specified.

Mitigation status. The paper does not address this limitation. It reports no confidence intervals, no standard deviations, no statistical tests, and no numerical accuracy differences at specific compute points. The 10-shot linear transfer averaging over three seeds is the only nod to statistical variation, but the variance across these seeds is not reported. The paper's framing as a report ("this report") rather than a full experimental paper may explain the informal presentation, but a practitioner's decision to adopt the method depends on the magnitude and reliability of any performance difference. Without statistical quantification, the evidence is suggestive rather than conclusive.

---

Limitation 5: The Results Are Bounded by the Model's Capability — ReLU-Attention Matches Softmax but Does Not Exceed It, and the Paper Does Not Explore Whether Hyperparameter Re-Tuning Could Unlock Gains

The assumption or constraint. The paper's experimental design tests ReLU-attention as a drop-in replacement for softmax: it uses existing BigVision training configurations "without modifying hyperparameters" (Section 4). This demonstrates that ReLU-attention can be substituted without degrading performance, but it does not test whether ReLU-attention could surpass softmax if hyperparameters were optimized for the new attention mechanism. The paper explicitly frames the result as "approach or match" (Section 1, Figure 1 caption) — equivalence, not superiority.

The consequence. The practical value proposition of ReLU-attention rests entirely on its computational advantage (fewer gather operations) at matched accuracy. If the accuracy is merely matched, the benefit must come from throughput or latency improvements at deployment. However, as discussed in Limitation 1, these improvements are never measured. If ReLU-attention could outperform softmax under properly tuned hyperparameters — for example, by enabling higher learning rates, different regularization, or more aggressive data augmentation due to the changed attention dynamics — then the value proposition would be stronger: ReLU-attention would offer both computational and accuracy benefits. Conversely, if softmax were re-optimized for the specific hyperparameters used, it might pull ahead of ReLU-attention by a margin that outweighs the computational savings. The one-sided experimental design (ReLU-attention is dropped into a softmax-optimized recipe; softmax is not re-optimized for the recipe) makes it impossible to distinguish these scenarios.

The paper also does not explore whether ReLU-attention enables qualitatively different capabilities that softmax attention does not — for example, whether the sparsity induced by ReLU (which sets negative dot products to exactly zero) produces attention patterns that are more interpretable, more robust to noise, or better at handling irrelevant tokens. These potential advantages remain unexplored because the paper evaluates only final accuracy on standard benchmarks.

What evidence exists in the paper. The paper provides a single ablation that hints at unexplored potential: removing qk-layernorm has minimal effect on ReLU-attention (Figure 3), whereas prior work found qk-layernorm was necessary for softmax-attention stability at scale (Dehghani et al., 2023). This suggests ReLU-attention is less sensitive to certain architectural choices than softmax attention, which could translate to simpler architectures or training recipes — but this direction is not pursued. The paper also shows that gating provides no benefit at $\alpha = 1$ (Figure 4), indicating that mechanisms developed for softmax-free attention under the prior no-scaling paradigm become redundant once scaling is correct. This simplifies the architecture, but the paper does not quantify whether the simplified architecture (ReLU-attention without qk-layernorm, without gating) can be trained with a simpler or faster recipe than the full softmax pipeline.

Mitigation status. The paper does not attempt to address this limitation. It accepts equivalence as sufficient and does not run hyperparameter sweeps, learning rate searches, or architecture optimizations for ReLU-attention beyond the ablations shown. The paper's framing as a demonstration of viability rather than a state-of-the-art pursuit makes this an acceptable scope limitation, but it leaves open whether ReLU-attention's full potential exceeds the drop-in replacement results shown.

---

Limitation 6: No Investigation of Training Dynamics — The $L^{-1}$ Argument Operates at Initialization, but Training-Time Evolution of Attention Weight Statistics Is Unmeasured

The assumption or constraint. The paper's heuristic for the $L^{-1}$ scaling relies on initialization-time reasoning: at initialization, query and key elements are $O(1)$, the scaled dot products are $O(1)$, and so the expected attention weight without scaling is $O(1)$. Dividing by $L$ restores the $O(L^{-1})$ expected magnitude that softmax produces (Section 3). However, this reasoning applies only at the start of training. During training, the queries and keys evolve under gradient descent — their magnitudes, their correlations, and the distribution of their dot products all change. The ReLU threshold (which passes only positive dot products) interacts with these evolving statistics: as the model learns, certain query-key pairs may develop consistently positive dot products while others become consistently negative, changing the effective number of non-zero attention weights and thus the actual expected weight magnitude (as opposed to the initialization-time expectation over independent random vectors).

The consequence. Three potential failure modes are invisible in the paper's analysis because only final accuracy is reported:

1. Training instability. If the dot products grow in magnitude during training (a common phenomenon as models learn sharper representations), the ReLU output grows proportionally. The $L^{-1}$ scaling is a fixed constant that does not adapt to evolving dot-product scales, so the effective attention weight magnitude may drift upward during training, potentially causing the outputs $o_i$ to grow and destabilizing the residual stream. Softmax's normalization ($\sum_j \exp(x_j)$) automatically compensates for scale changes because it divides by the sum of exponentiated values — if all dot products double, the softmax output becomes sharper but its sum remains 1. ReLU-attention has no such adaptive normalization.

2. Changing sparsity patterns. At initialization, approximately half of query-key dot products are positive (assuming symmetric distributions around zero), so roughly half the attention weights survive the ReLU threshold. As training proceeds, the fraction of positive dot products may change — the model might learn to make most dot products negative (producing very sparse attention) or predominantly positive (producing dense attention). The $L^{-1}$ scaling treats all patterns identically, but the effective attention behavior depends on how many weights are non-zero. Softmax handles this automatically because the normalization denominator adapts to the number of large vs. small weights.

3. Untested interaction with optimizers and learning rate schedules. The $L^{-1}$ scaling might interact with adaptive optimizers (like Adam) differently than softmax's normalization. In softmax, the gradient with respect to each query-key dot product includes a term from the normalization denominator, creating competition among attention weights — increasing one weight decreases others through the softmax normalization. ReLU-attention lacks this competitive normalization, so the gradient dynamics are fundamentally different. The paper does not analyze whether this leads to different convergence speeds, different sensitivity to learning rate, or different propensity for overfitting.

What evidence exists in the paper. The paper provides no training dynamics analysis. There are no training curves (loss or accuracy over steps), no plots of attention weight statistics (mean, variance, sparsity, positivity fraction) over the course of training, and no gradient analysis. The only temporal evidence is the final accuracy at the end of training (30 epochs for i21k, 300 epochs for i1k). The paper reports that both configurations use "a roughly similar number of steps of around 9e5" (Section 4), which implies comparable optimization difficulty, but this is a single number at the endpoint, not a trajectory comparison.

The $\alpha$ sweep in Figure 2 provides indirect evidence that the $L^{-1}$ scaling remains important throughout training (not just at initialization), because the experiments run to completion and the optimal $\alpha \approx 1$ is measured after full training. However, this does not reveal whether the mechanism that makes $\alpha \approx 1$ optimal at the endpoint is the same as the initialization-scale mechanism the paper hypothesizes — it could be that the scaling compensates for a training-dynamics effect unrelated to initialization.

Mitigation status. The paper does not address this limitation. It does not analyze training dynamics, does not provide training curves, and does not test whether the optimal $\alpha$ changes if measured at different points during training rather than only at convergence. The heuristic argument is presented as motivation, not as a claim about training dynamics, so the paper is not making claims it fails to support — but a practitioner deploying ReLU-attention in a new setting cannot predict from the paper's evidence whether the training will be stable or whether the $L^{-1}$ scaling will remain appropriate when training dynamics differ from the tested configurations.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the conversation around attention mechanisms from a mindset where softmax is assumed irreducible to one where it is understood as one point in a broader family of activation-plus-scaling choices. The conceptual shift is from "what can we build around softmax to make attention cheaper?" to "what properties of softmax actually matter, and can we satisfy them more cheaply?" This is a reframing rather than a paradigm shift — softmax is not dethroned, and the paper does not claim ReLU-attention is superior — but it changes what questions the field asks about attention design.

The diagnostic is more important than the solution. The paper's primary contribution is not ReLU-attention itself but the discovery that sequence-length scaling is the missing ingredient that prior work overlooked. Before this paper, the evidence pointed toward a negative conclusion: pointwise activations in attention degrade accuracy (Shen et al., 2023; Hua et al., 2022; Hron et al., 2020). The implicit interpretation was that softmax's particular functional form — its exponentiation, its normalization, its sum-to-one constraint — was functionally necessary. This paper demonstrates that interpretation was incorrect. The failure of prior pointwise replacements was not about softmax's functional form but about a scale mismatch: softmax produces attention weights of expected magnitude $O(L^{-1})$, naive pointwise replacements produce weights of $O(1)$, and the rest of the transformer was tuned for the $O(L^{-1})$ regime. The $L^{-1}$ scaling factor fixes the mismatch, and once it is fixed, the functional form of the activation matters surprisingly little — seven different nonlinearities all converge to similar accuracy at $\alpha \approx 1$ (Figure 2).

This is a simplicity result that should redirect research effort. The field has invested substantial work in complex solutions to the softmax bottleneck: gating mechanisms (Hua et al., 2022), linear attention reformulations (Katharopoulos et al., 2020; Lu et al., 2021), and scale-invariant normalization schemes (Li et al., 2022). This paper shows that for vision transformers at moderate scale, a one-line change — dividing by $L$ — suffices. This does not make those other lines of work obsolete — linear attention retains asymptotic complexity advantages for very long sequences, and gating may matter at scales this paper does not test — but it reframes the baseline. Any new mechanism that claims to improve on softmax must now compare against the simplest possible alternative (pointwise activation + $L^{-1}$ scaling), not just against softmax or against naive pointwise replacements. This raises the bar for complexity: if your mechanism is more complicated than ReLU(x)/L, it should demonstrate clear benefits over that simple baseline, not just over the straw man of unscaled pointwise attention.

Reconciling contradictory prior findings. The paper indirectly resolves a tension in the literature that the authors do not emphasize explicitly. On one hand, linear attention methods (Katharopoulos et al., 2020) showed that removing softmax entirely could work for certain tasks, but often with accuracy penalties on large-scale benchmarks. On the other hand, Hua et al. (2022) showed that squared ReLU with gating could match softmax, but added architectural complexity. The paper's unified picture is: accuracy degradation from removing softmax is largely explained by the loss of the $O(L^{-1})$ weight scaling, not by the loss of the exponential, the normalization sum, or the probability-distribution property. Different prior works were compensating for the scale loss in different ways — gating partially compensated, linear attention restructured the computation to avoid the problem — but none isolated the scale as the core issue. Once scale is restored via $L^{-1}$, the other interventions become unnecessary for achieving softmax-matched accuracy at the tested scales.

What research directions become more attractive. The paper makes scale-aware attention design a first-class research direction. Questions that were previously speculative become empirically tractable: Can the scaling factor be learned rather than hard-coded? Does the optimal scaling depend on the activation function? Does it change with sequence length in non-obvious ways? Does it interact with the number of attention heads? The $\alpha$ sweep methodology (Figure 2) provides a template for systematically investigating these questions — vary the scaling exponent, measure the effect, compare across conditions. Before this paper, the field lacked a clear demonstration that such sweeps would reveal anything interpretable; now there is a positive example to build on.

The paper also makes hardware-algorithm co-design more attractive for attention. If a pointwise activation with simple scaling can match softmax's representational quality, then hardware designers can optimize for the much simpler operation (a multiply, a max, and a divide, with no cross-position dependencies) rather than for the numerically complex softmax (exponentiation, max-over-sequence, sum-over-sequence, division). The paper does not demonstrate such hardware benefits, but by establishing functional equivalence, it gives hardware teams a credible target to optimize for.

What becomes less attractive. The paper makes gating as a general solution to softmax removal less attractive, at least for vision transformers at the tested scales. Figure 4 shows that gating provides minimal benefit at $\alpha \approx 1$ while adding ~9.3% computational overhead for the S/8 model. If the $L^{-1}$ scaling is sufficient, gating is an unnecessary cost. This does not mean gating is useless — it may matter at larger scales, in different domains, or when combined with learned scaling — but the default hypothesis should now be that gating is not required, and the burden of proof is on demonstrating settings where it provides benefits beyond what simple scaling achieves.

Similarly, the paper makes scale-invariant normalization that preserves the gather operation (Li et al., 2022) less attractive as a performance play. Those methods normalize attention weights to sum to one (or to have bounded norm) but retain the need to compute a normalization constant across the sequence, which is precisely the gather operation that prevents efficient parallelization. If $L^{-1}$ scaling achieves comparable accuracy without any cross-position communication, there is little reason to accept the parallelization penalty of gather-based normalization.

Follow-Up Research This Work Enables

1. Learned per-head scaling factors instead of fixed $L^{-1}$. The paper notes it is "unsure... if this term could be learned" (Section 5). A direct follow-up would replace the fixed $L^{-1}$ factor with a learned scalar parameter $\gamma_h$ per attention head $h$, initialized to $L^{-1}$ but allowed to evolve during training. This would test whether the optimal scaling factor is genuinely $L^{-1}$ or whether the model benefits from head-specific deviations. The experiment is simple: take the ReLU-attention configuration from Figure 1, add a learnable scalar multiply after the ReLU (so $\alpha_{ij} = \gamma_h \cdot \text{ReLU}(q_i^\top k_j / \sqrt{d})$), and compare final accuracy and the learned $\gamma_h$ values to the fixed $L^{-1}$ baseline. If $\gamma_h$ consistently converges near $L^{-1}$ across heads and model sizes, this strengthens the claim that $O(L^{-1})$ expected weight is the relevant invariant. If different heads learn substantially different scaling factors, it reveals that the optimal scale is head-specific (perhaps depending on the head's role in the network) and that a single fixed factor is a compromise. An intermediate design — learning a scalar per head that is regularized toward $L^{-1}$ — would test whether the fixed scaling is a useful inductive bias or an unnecessary constraint.

2. Attention weight statistics throughout training for ReLU-attention vs. softmax-attention. The paper's $O(L^{-1})$ argument operates at initialization, but the behavior throughout training is unexamined. A follow-up would instrument a training run to log, at regular intervals: (a) the mean, variance, and sparsity (fraction of zero weights) of the attention weights for both ReLU-attention and softmax-attention; (b) the fraction of query-key dot products that are positive (and thus survive the ReLU threshold); (c) the $L_2$ norm of the attention output vectors $o_i$; and (d) gradient norms for the query and key projection matrices. Comparing these statistics between ReLU-attention and softmax-attention across the S/16, B/16, and L/16 models would reveal whether training dynamics differ despite similar final accuracy. Specific hypotheses to test: Does ReLU-attention's sparsity increase during training (as the model learns to suppress irrelevant tokens) or remain near the ~50% expected at initialization? Does the absence of softmax's competitive normalization (where increasing one weight decreases others) lead to different attention-weight distributions? Does the attention output magnitude drift over training for ReLU-attention, or does the $L^{-1}$ scaling keep it stable? Answering these questions would build the mechanistic understanding that the paper's heuristic lacks, and could explain the failure of identity activation (which passes negative dot products through unchanged, potentially creating cancellation effects in the weighted sum that ReLU avoids via sparsity).

3. ReLU-attention in language modeling at sequence lengths where $L^{-1}$ becomes very small. The paper's longest tested sequence is 785 tokens (S/8). Modern language models routinely process 2,048–32,768 tokens, and long-context models reach 100,000+. At $L = 32,768$, the $L^{-1}$ factor is $3 \times 10^{-5}$, making the typical attention weight very small. A follow-up would train a GPT-style autoregressive transformer (with causal masking) using ReLU-attention with $L^{-1}$ scaling on a standard language modeling benchmark (e.g., C4 or The Pile) at multiple sequence lengths (512, 2,048, 8,192 tokens), measuring perplexity against a softmax baseline. The critical question is whether $\alpha = 1$ remains optimal, or whether longer sequences require a different scaling exponent (e.g., $\alpha = 1$ for short sequences, $\alpha < 1$ for long sequences to prevent weights from becoming too small). The causal mask interacts differently with ReLU than with softmax — in softmax, masked positions contribute nothing to the sum; in ReLU-attention, masked positions are simply excluded from the weighted sum, but the $L^{-1}$ factor still divides by the full sequence length rather than the number of unmasked positions. This mismatch could cause the effective attention weight scale to shrink as the model attends to early tokens (where most future positions are masked), and a follow-up could test whether normalizing by the number of unmasked positions ($\sum_{j \leq i} 1 = i$ for position $i$) rather than the full $L$ restores proper scaling in the causal setting.

4. ReLU-attention combined with FlashAttention-style kernel fusion to quantify actual throughput gains. The paper claims ReLU-attention "can be parallelized over the sequence length dimension with fewer gather operations" but provides no throughput or latency numbers. A follow-up would implement ReLU-attention within a FlashAttention-style fused kernel (Dao et al., 2022) — which already fuses the softmax with the matrix multiply to reduce memory I/O — and benchmark it against the standard FlashAttention softmax implementation at various batch sizes, sequence lengths (200 to 8,000 tokens), and head dimensions. The key comparison is wall-clock time per forward pass, not theoretical FLOP counts. FlashAttention's key optimization is computing softmax in blocks to avoid materializing the full $L \times L$ attention matrix in HBM; ReLU-attention removes the need for the running max and sum statistics that FlashAttention maintains across blocks for the softmax normalization. The hypothesis is that ReLU-attention within a fused kernel could reduce register pressure and simplify the block-wise computation, potentially yielding throughput gains above and beyond the elimination of exponentiation. Quantifying this gain — is it 5%, 20%, or 50%? — would convert the paper's "opportunity" into a concrete engineering case for (or against) adoption. If the gain is negligible (because exponentiation and gather costs are dwarfed by matrix multiplies in accelerators), the practical motivation for ReLU-attention weakens substantially. If the gain is large, it justifies investment in optimized ReLU-attention kernels.

5. Stress-test: ReLU-attention at scale where training instability matters. The paper's qk-layernorm ablation caveats that findings "may change at scale" (Figure 3). A follow-up would push ReLU-attention to model scales where training instability is a known challenge without qk-layernorm — for example, ViT models with ~1B+ parameters, or language models at the 1B–7B parameter range. The experiment would compare training stability (loss curves, gradient norms, occurrence of loss spikes) for three conditions: softmax with qk-layernorm (the stable baseline), ReLU-attention with qk-layernorm, and ReLU-attention without qk-layernorm. The finding that qk-layernorm is unnecessary for ReLU-attention at the tested scales (S-small to B-base on ImageNet) would be significantly strengthened if it holds at billion-parameter scale, because it would suggest ReLU-attention is inherently more stable than softmax attention — the $L^{-1}$ scaling might provide sufficient normalization without per-query-key LayerNorm. Conversely, if ReLU-attention without qk-layernorm becomes unstable at large scale while the softmax+qk-layernorm baseline remains stable, it reveals a limitation: the $L^{-1}$ scaling compensates for magnitude at initialization but does not prevent the scale of dot products from growing during training, whereas qk-layernorm actively constrains this growth. This negative result would be informative — it would clarify the boundary between what the $L^{-1}$ factor handles (initialization-scale matching) and what it does not (ongoing scale control during training).

6. ReLU-attention in domains where sparsity is beneficial: retrieval, long-document QA, video. ReLU-attention produces exactly zero weights for negative query-key dot products, unlike softmax which assigns exponentially small but non-zero weights. In domains where attention should be sparse — retrieving a few relevant documents from a large context, answering questions about specific video frames, or attending to relevant image regions in a cluttered scene — the induced sparsity might be a feature rather than a side effect. A follow-up would evaluate ReLU-attention against softmax-attention on a long-context retrieval task (e.g., answering questions about long documents where the answer appears in one or two passages) and on a video action recognition task where only a few frames are relevant to the action. The hypothesis is that ReLU-attention's hard zeroing of negative dot products acts as an implicit attention-pruning mechanism, reducing noise from irrelevant tokens and improving task performance beyond what the $L^{-1}$ scaling argument predicts. The experiment would measure not only accuracy but also the effective sparsity of the attention weights (what fraction are exactly zero after training) and the correlation between zero weights and token relevance (do the zeroed tokens correspond to distractor passages or irrelevant frames?). Even if accuracy is equivalent to softmax, demonstrated sparsity benefits could motivate ReLU-attention for memory-constrained deployment: zero weights need not be stored or communicated, potentially enabling sparse attention implementations that softmax (with its dense small-positive weights) cannot exploit.

Practical Applications and Downstream Use Cases

1. On-device vision transformers for mobile and edge deployment. Deploying vision transformers on mobile phones, AR/VR headsets, or embedded systems requires minimizing both latency and energy consumption. The softmax operation's exponentiation is implemented in hardware as a series of table lookups or polynomial approximations on most mobile accelerators, adding latency relative to simple arithmetic. The normalization sum creates a dependency between all sequence positions, preventing the attention computation from being split across processing elements without communication. ReLU-attention replaces both with a threshold, a multiply, and a division — operations that are natively fast on mobile GPUs and NPUs. While the paper does not measure inference latency, an edge deployment team could: (1) take a trained ViT model (e.g., MobileViT or a distilled ViT-S), (2) replace the attention softmax with ReLU-attention, (3) retrain or fine-tune on the target dataset (since the paper shows drop-in compatibility with existing training recipes), and (4) benchmark on-device inference time. The expected benefit is reduced per-inference latency from eliminating the exponentiation and gather, and reduced energy consumption from fewer compute operations. The risk — a small accuracy drop relative to softmax — may be acceptable for on-device applications where latency budgets are tight and absolute accuracy requirements are lower than for server-side models.

2. Large-scale batch inference on video understanding pipelines. Processing video with vision transformers requires applying attention across spatial patches and temporal frames, producing sequence lengths of thousands of tokens (e.g., 16 frames × 196 patches = 3,136 tokens). At these sequence lengths, the attention matrix is a primary memory bottleneck. ReLU-attention's induced sparsity — negative dot products become exactly zero — means that after the ReLU, a large fraction of the attention matrix need not be stored or used in the subsequent weighted sum. A video inference pipeline at scale (e.g., content moderation, video recommendation, or video search) processes millions of videos daily. If ReLU-attention can be implemented to exploit sparsity — only computing outputs for non-zero attention weights, and only storing non-zero weight indices — the memory footprint of the attention operation could shrink substantially. The paper does not measure sparsity rates after training, but at initialization approximately half the weights are zero. If training preserves or increases this sparsity, a video processing team could store only the non-zero attention weights (using sparse matrix formats) and compute the weighted sum over a subset of values, reducing both memory and compute. This matters more for batch inference where many videos are processed in parallel and memory pressure limits batch size — smaller per-sample memory means larger batches, higher throughput, and lower cost per video.

3. Training data generation and knowledge distillation pipelines where many forward passes are needed. In self-improvement, distillation, or synthetic data generation pipelines (e.g., generating rationales from a teacher model, scoring candidate outputs, or producing training data for a smaller student model), a transformer may be run for millions or billions of forward passes. A small per-forward-pass speedup from ReLU-attention would compound to substantial total compute savings. For example, if ReLU-attention eliminates 5% of the forward-pass latency relative to an optimized softmax implementation, and a pipeline requires $10^8$ forward passes, the savings would be the equivalent of 5 million forward passes. The paper's finding that ReLU-attention "matches the scaling trends" without hyperparameter modification (Figure 1) means that a team could substitute ReLU-attention into an existing model, verify that accuracy on the downstream task is preserved (as the paper does for ImageNet-1k and linear transfer), and immediately realize the latency savings without any pipeline re-engineering. The key uncertainty is whether the latency savings are realized in practice — which depends on whether the softmax is a meaningful fraction of total forward-pass time at the batch sizes and hardware used. A team considering this would first profile their model to determine what fraction of forward-pass time is spent in softmax operations, then evaluate whether ReLU-attention actually reduces that fraction on their hardware.

4. Hardware accelerator design targeting transformer workloads. Hardware teams designing next-generation ML accelerators must decide which operations to optimize in silicon. The paper provides evidence that the softmax operation — which currently receives dedicated hardware support in some accelerators (e.g., specialized exponential units, reduction trees for the sum) — may not be architecturally necessary for transformer performance. If ReLU-attention can match softmax across a range of scales and tasks (with the limitations acknowledged — vision-only, moderate scale), a hardware designer could choose to: (a) remove dedicated softmax hardware and use the silicon area for additional matrix multiply units, (b) design a fused "attention without gather" instruction that computes $o_i = \sum_j \text{ReLU}(q_i^\top k_j / \sqrt{d}) v_j / L$ in a single pass without cross-lane communication, or (c) optimize for sparse attention with ReLU-induced zeros as a first-class primitive. The hardware benefit is speculative without throughput measurements, but the paper's functional equivalence result is the necessary precondition: a hardware team would not optimize for an attention variant that degrades model quality. The paper makes ReLU-attention a credible enough alternative that hardware design studies can assume it as a target workload without worrying that the model accuracy will fall behind softmax-based transformers.