As technologies emerge that skip or optimize 'softmax,' a chronic computational bottleneck in AI transformer models, a new AI era capable of processing larger volumes of information faster is opening.
Imagine you are trying to find specific information in a massive library. What if the librarian had to open and sort every single book before giving you an answer? As the number of books increases, the time it takes for the librarian to answer would grow exponentially, beyond imagination. The problem currently faced by AI models—especially the ‘Transformer’ (the core AI structure that grasps relationships between words in a sentence)—is quite similar.
Recently, in the field of artificial intelligence, innovative attempts are being actively made to skip or optimize the ‘softmax’ operation, which has been taken for granted, to solve this chronic ‘bottleneck.’
Why It Matters
The Transformer is the foundation for almost all modern AI. Within this structure, the softmax function is like a ‘toll’ that must be paid every time information is processed. Softmax plays an essential role in helping AI make final choices or judgments by converting various data values into a probability distribution Source: Softmax function - Wikipedia.
The issue is that the amount of information AI handles has exploded recently. As more data needs to be processed, softmax operations become the primary culprit for consuming massive amounts of memory and slowing down overall computational speed. If this operation can be bypassed or optimized, AI can process longer contexts much faster while using far less energy. This signals a hopeful future where we can comfortably use smarter, more responsive AI assistants in our daily lives.
The Explainer
To understand softmax, let’s use a simple analogy. Think of the process of picking your ‘favorite product’ at a shopping mall. The process of comparing the price, quality, and design of numerous products and then assigning a probability to how much each product appeals to you is essentially a softmax operation. You score all the choices and sum them up to convert them into a 100% probability distribution.
But what if an AI model had to read an entire book’s worth of content that is very long and complex? Calculating the precise probability by comparing every single word one by one is an extremely arduous task.
Researchers have recently found a few clever ways around this.
- Static Contraction Technique: Just like paving a road in advance, this method pre-blocks unnecessary calculation paths during the operation and replaces them with efficient formulas. This significantly reduces the risk of AI models experiencing Out-Of-Memory (OOM) issues in long contexts Source: GitHub - PJHkorea/jax-softmax-bypass, Source: ShowHN: Bypassing Transformer Softmax via Static Contraction.
- Softmax-free: A model called ‘SOFT’ uses a relatively simple mathematical function called a ‘Gaussian kernel’ instead of existing complex operations. It is like choosing a more intuitive shortcut instead of punching into a complex scientific calculator every time for probability calculations Source: SOFT: Softmax-free Transformer with Linear Complexity, Source: [2110.11945] SOFT: Softmax-free Transformer with Linear Complexity.
-
Forgetting Mechanism: ‘ForgettingTransformer(FoX)’ utilizes a ‘forget gate’ that appropriately forgets unnecessary information. Just as we naturally memorize important information and let the rest go, it selectively adjusts the attention scores, making context processing much easier Source: ForgettingTransformer: Softmax Attention with a Forget Gate, [Source: ForgettingTransformer: Softmax Attention with… Papers with Code](https://paperswithcode.co/paper/2503.02130), Source: GitHub - zhixuan-lin/forgetting-transformer.
Where We Stand
| Currently, these technologies are rapidly developing in various forms, from laboratory-stage Proof of Concepts (PoC) to academic proposals. However, it is premature to say that softmax has been completely replaced. While it is clearly proven to be more efficient than standard existing methods, there remain aspects that require additional verification before being applied to all general-purpose AI models. Nevertheless, attempts such as ‘UniAttn,’ which drastically reduce computational costs while minimizing performance degradation, are consistently achieving results, raising expectations [Source: UniAttn: Reducing Inference Costs via Softmax… | Papers with Code](https://paperswithcode.co/paper/2502.00439). |
What’s Next
The future AI technology race will shift from merely creating ‘bigger models’ to ‘who can create more efficient models.’ These softmax bypassing techniques will be essential keys to running more powerful AI on devices with physically limited computational resources, like the smartphones we use. If these current studies come to fruition, we will encounter smarter AI in our daily lives that perfectly remembers longer conversation histories and responds at a faster speed.
References
- GitHub - PJHkorea/jax-softmax-bypass: https://github.com/PJHkorea/jax-softmax-bypass
- Vertex-Softmax: Tight Transformer Verification via Exact Softmax Optimization∗: https://arxiv.org/pdf/2605.10974
- SimA: Simple Softmax-free Attention for Vision Transformers: https://openaccess.thecvf.com/content/WACV2024/papers/Koohpayegani_SimA_Simple_Softmax-Free_Attention_for_Vision_Transformers_WACV_2024_paper.pdf
- SOFT: Softmax-free Transformer with Linear Complexity: https://proceedings.neurips.cc/paper/2021/file/b1d10e7bafa4421218a51b1e1f1b0ba2-Paper.pdf
- [2110.11945] SOFT: Softmax-free Transformer with Linear Complexity: https://arxiv.org/abs/2110.11945
-
Differential Transformer Hacker News: https://news.ycombinator.com/item?id=41776324 - Softmax function - Wikipedia: https://en.wikipedia.org/wiki/Softmax_function
- ForgettingTransformer: Softmax Attention with a Forget Gate: https://arxiv.org/abs/2503.02130
- ShowHN: Bypassing Transformer Softmax via Static Contraction: https://news.ycombinator.com/item?id=49666335
-
ForgettingTransformer: Softmax Attention with… Papers with Code: https://paperswithcode.co/paper/2503.02130 -
In-Context Learning with Transformers: Softmax… OpenReview: https://openreview.net/forum?id=lfxIASyLxB - Transformers are RNNs: Fast Autoregressive… - YouTube: https://www.youtube.com/watch?v=hAooAOFRsYc
-
UniAttn: Reducing Inference Costs via Softmax… Papers with Code: https://paperswithcode.co/paper/2502.00439 - GitHub - zhixuan-lin/forgetting-transformer: https://github.com/zhixuan-lin/forgetting-transformer
- Softmax Function in Deep Learning: https://lzwjava.com/notes/2025-06-03-softmax-en
- Compress data
- Convert numbers into a probability distribution
- Delete data
- Gaussian kernel function
- Linear function
- Log function
- ForgettingTransformer(FoX)
- Softmax-free Transformer
- UniAttn