# A Study on ReLU and Softmax in Transformer

arXiv: https://arxiv.org/abs/2302.06461 · HTML: https://mlsanity.com/paper/2302.06461/

## Key facts

- Cited by 53 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 33 references that are in the map.
- First posted on arXiv 13 Feb 2023.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)
- [Knowledge Neurons in Pretrained Transformers](https://mlsanity.com/paper/2104.08696/) (2021)
- [Sequence Level Training with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06732/) (2015)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Sparse Sequence-to-Sequence Models](https://mlsanity.com/paper/1905.05702/) (2019)
- [An Actor-Critic Algorithm for Sequence Prediction](https://mlsanity.com/paper/1607.07086/) (2016)
- [Large Memory Layers with Product Keys](https://mlsanity.com/paper/1907.05242/) (2019)
- [Achieving Human Parity on Automatic Chinese to English News Translation](https://mlsanity.com/paper/1803.05567/) (2018)
- [Augmenting Self-attention with Persistent Memory](https://mlsanity.com/paper/1907.01470/) (2019)
- [Selective Attention for Context-aware Neural Machine Translation](https://mlsanity.com/paper/1903.08788/) (2019)
- [Sparse Attention with Linear Units](https://mlsanity.com/paper/2104.07012/) (2021)
- [Non-Autoregressive Neural Machine Translation with Enhanced Decoder Input](https://mlsanity.com/paper/1812.09664/) (2018)
- [Kformer: Knowledge Injection in Transformer Feed-Forward Layers](https://mlsanity.com/paper/2201.05742/) (2022)

## Built on by (most cited first, 30 of 53)

- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models](https://mlsanity.com/paper/2310.04564/) (2023)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [Beyond the Limits: A Survey of Techniques to Extend the Context Length in Large Language Models](https://mlsanity.com/paper/2402.02244/) (2024)
- [How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations](https://mlsanity.com/paper/2310.10616/) (2023)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [ProSparse: Introducing and Enhancing Intrinsic Activation Sparsity within Large Language Models](https://mlsanity.com/paper/2402.13516/) (2024)
- [Replacing softmax with ReLU in Vision Transformers](https://mlsanity.com/paper/2309.08586/) (2023)
- [Improving fine-grained understanding in image-text pre-training](https://mlsanity.com/paper/2401.09865/) (2024)
- [What can a Single Attention Layer Learn? A Study Through the Random Features Lens](https://mlsanity.com/paper/2307.11353/) (2023)
- [Frequency Decoupling for Motion Magnification via Multi-Level Isomorphic Architecture](https://mlsanity.com/paper/2403.07347/) (2024)
- [Approximating Two-Layer Feedforward Networks for Efficient Transformers](https://mlsanity.com/paper/2310.10837/) (2023)
- [On Provable Length and Compositional Generalization](https://mlsanity.com/paper/2402.04875/) (2024)
- [HSR-Enhanced Sparse Attention Acceleration](https://mlsanity.com/paper/2410.10165/) (2024)
- [Ultra-Sparse Memory Network](https://mlsanity.com/paper/2411.12364/) (2024)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)
- [Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems](https://mlsanity.com/paper/2310.12956/) (2023)
- [More Expressive Attention with Negative Weights](https://mlsanity.com/paper/2411.07176/) (2024)
- [ZETA: Leveraging Z-order Curves for Efficient Top-k Attention](https://mlsanity.com/paper/2501.14577/) (2025)
- [Rethinking Attention: Polynomial Alternatives to Softmax in Transformers](https://mlsanity.com/paper/2410.18613/) (2024)
- [Unification of Symmetries Inside Neural Networks: Transformer, Feedforward and Neural ODE](https://mlsanity.com/paper/2402.02362/) (2024)
- [Learning Neural Networks with Sparse Activations](https://mlsanity.com/paper/2406.17989/) (2024)
- [Quantum Doubly Stochastic Transformers](https://mlsanity.com/paper/2504.16275/) (2025)
- [Mitigating Sycophancy in Decoder-Only Transformer Architectures: Synthetic Data Intervention](https://mlsanity.com/paper/2411.10156/) (2024)
- [Threshold Differential Attention for Sink-Free, Ultra-Sparse, and Non-Dispersive Language Modeling](https://mlsanity.com/paper/2601.12145/) (2026)
- [ZeroS: Zero-Sum Linear Attention for Efficient Transformers](https://mlsanity.com/paper/2602.05230/) (2026)
- [A Non-negative VAE:the Generalized Gamma Belief Network](https://mlsanity.com/paper/2408.03388/) (2024)
- [PERK: Long-Context Reasoning as Test-Time Learning](https://mlsanity.com/paper/2507.06415/) (2025)
- [On Space Folds of ReLU Neural Networks](https://mlsanity.com/paper/2502.09954/) (2025)
- [Self-Adjust Softmax](https://mlsanity.com/paper/2502.18277/) (2025)

## Abstract (arXiv)

The Transformer architecture consists of self-attention and feed-forward networks (FFNs) which can be viewed as key-value memories according to previous works. However, FFN and traditional memory utilize different activation functions (i.e., ReLU and Softmax respectively), which makes them not equivalent. In this paper, we first rebuild the connections between FFN and key-value memory by conducting extensive studies on ReLU and Softmax, and find they are equivalent when adding an additional layer normalization module on Softmax. In addition, ReLU outperforms Softmax on both FFN and key-value memory when the number of value slots is large. We analyze the reasons and then explore this good property of ReLU on the self-attention network where the original Softmax activation performs poorly on long input sequences. We then propose a full ReLU architecture named ReLUFormer which performs better than the baseline Transformer on long sequence tasks such as document translation. This paper sheds light on the following points: 1) Softmax and ReLU use different normalization methods over elements which lead to different variances of results, and ReLU is good at dealing with a large number of key-value slots; 2) FFN and key-value memory are equivalent, and thus the Transformer can be viewed as a memory network where FFNs and self-attention networks are both key-value memories.
