@inproceedings{zheng-etal-2025-self,
    title = "Self-Adjust Softmax",
    author = "Zheng, Chuanyang  and
      Gao, Yihang  and
      Chen, Guoxuan  and
      Shi, Han  and
      Xiong, Jing  and
      Ren, Xiaozhe  and
      Huang, Chao  and
      Li, Zhenguo  and
      Li, Yu",
    editor = "Christodoulopoulos, Christos  and
      Chakraborty, Tanmoy  and
      Rose, Carolyn  and
      Peng, Violet",
    booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2025",
    address = "Suzhou, China",
    publisher = "Association for Computational Linguistics",
    url = "https://preview.aclanthology.org/ingest-emnlp/2025.emnlp-main.397/",
    pages = "7827--7847",
    ISBN = "979-8-89176-332-6",
    abstract = "The softmax function is crucial in Transformer attention, which normalizes each row of the attention scores with summation to one. **Usually, tokens with larger attention scores are important for the final prediction.However, the softmax function can face a gradient vanishing issue for such important tokens (e.g., probabilities close to one), leading to optimization difficulties for the important tokens so that the performance may not be better.**In this paper, we propose Self-Adjust Softmax (SA-Softmax) to address this issue by modifying $softmax(z)$ to $z \cdot softmax(z)$ and its normalized variant $\frac{(z - min(z_{\min},0))}{max(0,z_{max})-min(z_{min},0)} \cdot softmax(z)$.We theoretically show that SA-Softmax provides enhanced gradient properties compared to the vanilla softmax function.Moreover, Attention can be seamlessly integrated into existing Transformer models to their attention mechanisms with minor adjustments.We conducted experiments to evaluate the empirical performance of Transformer models using compared to the vanilla softmax function. These experiments, involving models with up to 2.7 billion parameters, are conducted across diverse datasets, language tasks, and positional encoding methods."
}