Transformers with Learnable Activation Functions (2023)

Abstract: Activation functions can have a significant impact on reducing the topological complexity of input data and therefore, improving a model’s performance. However, the choice of activation functions is seldom discussed or explored in Transformer-based language models. As a common practice, commonly used activation functions like Gaussian Error Linear Unit (GELU) are chosen beforehand and then remain fixed from pre-training to fine-tuning. In this paper, we investigate the impact of activation functions on Transformer-based models by utilizing rational activation functions (RAFs). In contrast to fixed activation functions (FAF), RAFs are capable of learning the optimal activation functions from data. Our experiments show that the RAF-based Transformer model (RAFT) achieves a better performance than its FAF-based counterpart (). For instance, we find that RAFT outperforms on the GLUE benchmark by 5.71 points when using only 100 training examples and by 2.05 points on SQuAD with all available data. Analyzing the shapes of the learned RAFs further unveils that they vary across different layers and different tasks; opening a promising way to better analyze and understand large, pre-trained language models.

## introduction activation functions introduce non-linearity and increase neural networks' representational capacity, and therefore, play an essential role in designing deep learning models @xcite @xcite . @xcite explain the importance of activation functions by proposing to consider data as a topology with its own shape. they empirically show that activation functions accelerate the data topology transformation through different layers of a neural network to simplify its complexity and make it linearly separable in the output space. their experiments show that choosing the right activation function can have a significant impact on the overall performance. while any activation function can be used with transformers @xcite , their choice is made before pre-training and remains fixed afterwards. hence, the inductive bias an activation function imposes on the model cannot be adjusted during pre-training or fine-tuning. as many transformer-based models are pre-trained on a large amount of data, and changing the activation function for or during fine-tuning may negatively impact the performance foot_1 . moreover, the simple case of finding the optimal combination of k different activation functions in n different feedforward layers results in k n possible combinations and becomes intractable; e.g., 531,441 experiments for a 12-layer bert model and three different activation functions. as a result, most transformerbased pre-trained models adopt the gelu activation function that has been initially used for the bert model @xcite . to overcome the limitation of using a potentially suboptimal activation function that remains fixed during training, we propose to use a learnable activation function, namely, the rational activation function @xcite ). the raf is a universal function approximator that can approximate any existing activation function. the advantage of using rafs over fixed activation functions (faf) such as relu or gelu, is that the model can learn the optimal activation function from the data during (pre)training without the need to consider the choice of activation function as an additional dimension during hyperparameter tuning. 3 to evaluate the effectiveness of rafs, we pre-train two encoder-only transformers using raf and gelu respectively, within an academic budget. in our experiments, we find that: ## related work activation functions. there exists various predefined activation functions such as sigmoid, hyperbolic tangent (tanh), rectified linear unit (relu, fukushima 1969), and gaussian error linear unit (gelu, @xcite . there are also approaches that leverage automatic search to obtain optimal combinations of several base activation functions in a predefined search space @xcite @xcite @xcite . for instance, ramachandran et al. (2018) discovered the swish activation function by using this method. @xcite show that further extending the search space using evolutionary algorithms can also lead to an improvement. finally, several search-based works investigate how to train a combination of a set of activation functions to better adapt to specific tasks and architectures (manessi and rozza, ments such as pre-training large language models. 2018; @xcite . one substantial drawback of these search-based methods is that they are computationally expensive. especially for pre-trained language models where pre-training is costly, it is infeasible to perform a hyperparameter search for selecting the best activation function (even more so their combination). in contrast, the flexibility of rational activation functions (rafs) allows them to be trained along with the model parameters in an endto-end fashion @xcite . therefore, they can learn the optimized activation function from data during training. rafs have been successfully used in deep reinforcement learning for improving plasticity @xcite , cell detection models in biology @xcite , and adapter architectures @xcite . model act. funct. bert @xcite ) gelu gpt-1 @xcite gelu roberta @xcite gelu xlnet @xcite gelu albert @xcite gelu gpt-2 * @xcite gelu megatron-lm @xcite gelu electra + @xcite gelu t5 @xcite relu t5v1.1 @xcite geglu deberta + @xcite gelu bart @xcite ., 2020) gelu gpt-3 * (brown et al., 2020) gelu jurassic * (lieber et al., 2021) gelu gopher * (rae et al., 2021) gelu megatron-turing nlg * (smith et al., 2022) gelu chinchilla * (hoffmann et al., 2022) gelu canine + @xcite gelu lambda @xcite geglu opt @xcite relu ## raft: raf-based transformers we adopt the bert architecture @xcite where all activation functions in feedforward layers activation(w 1 x)w 2 are replaced with rational activation functions (illustrated in appendix a). the equation of rational activation function f (x) is as below: where a and b are learnable parameters, and m and n are degrees of f (x), which decide the complexity and fitting ability of rational functions. following @xcite , we use the safe pau formulation that further stabilizes training. selecting m and n. similar to taylor series, the higher the degrees m and n are, the more precise is the approximation of rational functions. however, indefinitely increasing the degrees also means adding more complexity and increasing training time. the challenge is to find suitable degrees that leads to rational functions with a strong fitting ability while keeping their complexity as low as possible. as this is still an open question, we set the search space of m and n to {4, 5}, and evaluate their ability to approximate the gelu function in the range of @xcite . our results show that using @xmath0, are adopted in this work (cf. figure 5 , appendix b ). this matches the findings in previous work @xcite @xcite as well. so overall, each rational activation function adds nine parameters, resulting in a total of 108 additional parameters in a 12-layer transformer model (less than 0.000098% of its original parameters). the weights of f (x) can further be initialized to approximate any existing activation functions. in our experiments, we initialize it with weights that approximate gelu. ## pre-training to evaluate the viability of raft, we pre-train two comparable transformer models from scratchone using the common fixed gelu activation function (faft), and another one using rafs (raft). ## analysis impact of raf initialization. to investigate how initialization affects the performance of raft, we train raft models initialized with gelu, relu, and the identity function. other hyperparameters are the same as those in section 4. ## conclusion and future work in this work, we propose to utilize rational activation functions (raf) in transformers to directly learn optimal activation functions from data during pre-training and fine-tuning. to evaluate the effectiveness of rational activation functions, we pre-trained a transformer-based language model, namely, raft. raft achieves a lower validation perplexity than faft during pre-training. our experimental results show that raft performs better than faft in general language understanding tasks and reading comprehension tasks across different data size scenarios. we further visualize and analyze rational activation functions across different layers and tasks after pre-training and fine-tuning and find that they can substantially vary across different layers and tasks. this provides us a new way to analyze and better understand transformerbased language models. for instance, we can investigate whether layers with similar rational activation functions encode similar linguistic properties. we further find that some layers exhibit a close to zero throughput of the rational activation function which indicates that the corresponding feedforward layer does not contribute too much to a model's prediction. we consider these as our future work.