Self-Adjust Softmax

TL;DR

This paper introduces Self-Adjust Softmax (SA-Softmax), a modified softmax function designed to improve gradient flow in Transformer attention, with theoretical benefits and empirical validation across large-scale models and diverse tasks.

Contribution

The paper proposes SA-Softmax, a novel softmax variant that enhances gradient properties and can be easily integrated into Transformer models.

Findings

SA-Softmax improves gradient flow in attention mechanisms.

Models with SA-Softmax outperform those with vanilla softmax.

Empirical results on large models show consistent performance gains.

Abstract

The softmax function is crucial in Transformer attention, which normalizes each row of the attention scores with summation to one, achieving superior performances over other alternative functions. However, the softmax function can face a gradient vanishing issue when some elements of the attention scores approach extreme values, such as probabilities close to one or zero. In this paper, we propose Self-Adjust Softmax (SA-Softmax) to address this issue by modifying $softmax(x)$ to $x \cdot softmax(x)$ and its normalized variant $\frac{(x - min(x_{\min},0))}{max(0,x_{max})-min(x_{min},0)} \cdot softmax(x)$. We theoretically show that SA-Softmax provides enhanced gradient properties compared to the vanilla softmax function. Moreover, SA-Softmax 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 SA-Softmax 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.