Look Mom, No Softmax! Introducing Decoupled Polynomial Mapping in a Softmax-Free Attention-Based Transformer

The conventional attention-based transformer structure utilizes Softmax operations in attention, loss determination, and inference. While Softmax allows for the separation of high-probability signals from low-probability signals and is foundational to language model performance, applying an exponential function to the entire tensor introduces critical bottlenecks. Specifically, Softmax (a) does not efficiently introduce structural sparsity, often preserving background noise that makes convergence less efficient, and (b) is highly sensitive to steep gradients, which can cause saturation and halt convergence. Hua et al. (2022) introduced the concept of replacing Softmax with a ReLU-squared attention mechanism in a Gated Attention Unit (GAU), adapting the ReLU-squared activation popularized by So et al. (2021). Building on this prior work, this paper introduces Decoupled Polynomial Mapping (DPM) as a universal function to replace Softmax in self-attention transformers in self-attention, loss determination, and expression. Instead of exponential expansion, DPM uses a multi-term polynomial function with trained coefficients and a trained exponent variable, formulated as a * ReLU(x)² + b * ReLU(x)ᶻ , where a, b, and z are trainable. A 24.7 million parameter baseline transformer was trained using sandwich-style normalization and standard ReLU-squared attention over 10 runs of 100 epochs in token-level Wikitext-103 (Subword BPE-tokenized WikiText-103 using the GPT-2 tokenizer), reached an average low validation loss of 3.5774 (perplexity 35.78). When DPM replaced ReLU-squared in attention and Softmax in loss determination, the 24.7 million parameter network trained on the same dataset over 10 runs resulted in significantly improved validation losses, averaging 2.7481 (perplexity 15.61).