AR E LU: A TTENTION - BASED R ECTIFIED L INEAR U NIT
Anonymous authors
Paper under double-blind review
A BSTRACT
Element-wise activation functions play a critical role in deep neural networks via affecting the expressivity power and the learning dynamics.
Learning-based activation functions have recently gained increasing attention and success.
We propose a new perspective of learnable activation function through formulating them with element-wise attention mechanism .
In each network layer, we devise an attention module which learns an element-wise, sign-based attention map for the pre-activation feature map.
The attention map scales an element based on its sign.
Adding the attention module with a rectiﬁed linear unit (ReLU) results in an ampliﬁcation of positive elements and a suppression of negative ones, both with learned, data-adaptive parameters.
We coin the resulting activation function
Attention-based Rectiﬁed Linear Unit (AReLU).
The attention module essentially learns an element-wise residue of the activated part of the input, as ReLU can be viewed as an identity transformation.
This makes the network training more resistant to gradient vanishing.
The learned attentive activation leads to well-focused activation of relevant regions of a feature map.
Through extensive evaluations, we show that AReLU signiﬁcantly boosts the performance of most mainstream network architectures with only two extra learnable parameters per layer introduced.
Notably, AReLU facilitates fast network training under small learning rates, which makes it especially suited in the case of transfer learning and meta learning.
I NTRODUCTION
Activation functions, introducing nonlinearities to artiﬁcial neural networks, is essential to networks’ expressivity power and learning dynamics.
Designing activation functions that facilitate fast training of accurate deep neural networks is an active area of research (Maas et al., 2013; Goodfellow et al., 2013; Xu et al., 2015a;
Clevert et al., 2015; Hendrycks & Gimpel, 2016; Klambauer et al., 2017;
Barron, 2017; Ramachandran et al., 2017).
Aside from the large body of hand-designed functions, learning-based approaches recently gain more attention and success (Agostinelli et al., 2014;
He et al., 2015; Manessi & Rozza, 2018; Molina et al., 2019; Goyal et al., 2019).
The existing learnable activation functions are motivated either by relaxing/parameterizing a non-learnable activation function (e.g. Rectiﬁed Linear Units (ReLU) (Nair & Hinton, 2010)) with learnable parameters
(He et al., 2015), or by seeking for a data-driven combination of a pool of pre-deﬁned activation functions (Manessi & Rozza, 2018).
Existing learning-based methods make activation functions data-adaptive through introducing degrees of freedom and/or enlarging the hypothesis space explored.
In this work, we propose a new perspective of learnable activation functions through formulating them with element-wise attention mechanism .
A straightforward motivation of this is a straightforward observation that both activation functions and element-wise attention functions are applied as a network module of element-wise multiplication .
More intriguingly, learning element-wise activation functions in a neural network can intuitively be viewed as task-oriented attention mechanism (Chorowski et al., 2015; Xu et al., 2015b), i.e., learning where (which element in the input feature map) to attend (activate) given an end task to fulﬁll .
This motivates an arguably more interpretable formulation of attentive activation functions .
Attention mechanism has been a cornerstone in deep learning.
It directs the network to learn which part of the input is more relevant or contributes more to the output.
There have been many variants of attention modules with plentiful successful applications.
In natural language processing, vector-wise attention is developed to model the long-range dependencies in a sequence of word vectors (Luong et al., 2015; Vaswani et al., 2017).
Many computer vision tasks utilize pixel-wise or channel-wise attention modules for more expressive and invariant representation learning (Xu et al., 2015b; Chen et al., 2017).
Element-wise attention (Bochkovskiy et al., 2020) is the most ﬁne-grained where each element of a feature volume can receive different amount of attention.
Consequently, it attains high expressivity with neuron-level degrees of freedom.
Inspired by that, we devise for each layer of a network an element-wise attention module which learns a sign-based attention map for the pre-activation feature map.
The attention map scales an element based on its sign.
Through adding the attention and a ReLU module, we obtain Attention-based
Rectiﬁed Linear Unit (AReLU) which ampliﬁes positive elements and suppresses negative ones, both with learned, data-adaptive parameters.
The attention module essentially learns an element-wise residue for the activated elements with respect to the ReLU since the latter can be viewed as an identity transformation.
This helps ameliorate the gradient vanishing issue effectively.
Through extensive experiments on several public benchmarks, we show that AReLU signiﬁcantly boosts the performance of most mainstream network architectures with only two extra learnable parameters per layer introduced.
Moreover, AReLU enables fast learning under small learning rates, making it especially suited for transfer learning.
We also demonstrate with feature map visualization that the learned attentive activation achieves well-focused, task-oriented activation of relevant regions.
R ELATED W ORK
Non-learnable activation functions Sigmoid is a non-linear, saturated activation function used mostly in the output layers of a deep learning model.
However, it suffers from the exploding/vanishing gradient problem.
As a remedy, the rectiﬁed linear unit (ReLU) (Nair & Hinton, 2010) has been the most widely used activation function for deep learning models with the state-of-the-art performance in many applications.
Many variants of ReLU have been proposed to further improve its performance on different tasks LReLU (Maas et al., 2013), ReLU6 (Krizhevsky & Hinton, 2010), RReLU (Xu et al., 2015a).
Besides that, some speciﬁed activation functions also have been designed for different usages, such as CELU (Barron, 2017), ELU (Clevert et al., 2015), GELU (Hendrycks & Gimpel, 2016), Maxout (Goodfellow et al., 2013), SELU (Klambauer et al., 2017), (Softplus) (Glorot et al., 2011), Swish (Ramachandran et al., 2017).
Learnable activation functions Recently, learnable activation functions have drawn more attentions.
PReLU (He et al., 2015), as a variants of ReLU, improves model ﬁtting with little extra computational cost and overﬁtting risk.
Recently, PAU (Molina et al., 2019) is proposed to not only approximate common activation functions but also learn new ones while providing compact representations with few learnable parameters.
Several other learnable activation functions such as
APL (Agostinelli et al., 2014), Comb (Manessi & Rozza, 2018), SLAF (Goyal et al., 2019) also achieve promising performance under different tasks.
Attention Mechanism Vector-Wise Attention Mechanism (VWAM) has been widely applied in
Natural Language Processing (NLP) tasks (Xu et al., 2015c;
Luong et al., 2015; Bahdanau et al., 2014; Vaswani et al., 2017; Ahmed et al., 2017).
VWAM learns which vector among a sequence of word vectors is the most relevant to the task in hand.
Channel-Wise Attention Mechanism (CWAM) can be regarded as an extension of VWAM from NLP to Vision tasks (Tang et al., 2019b; 2020;
Kim et al., 2019).
It learns to assign each channel an attentional value.
Pixel-Wise Attention Mechanism (PWAM) is also widely used in vision (Tang et al., 2019c;a).
Element-Wise Attention Mechanism (EWAM) assigns different values to each element without any spatial/channel constraint.
The recently proposed YOLOv4 (Bochkovskiy et al., 2020) is the ﬁrst work that introduces EWAM implemented by a convolutional layer and sigmoid function.
It achieves the state-of-the-art performance on object detection.
We introduce a new kind of EWAM for learnable activation function.
M ETHOD
We start by describing attention mechanism and then introduce element-wise sign-based attention mechanism based on which AReLU is deﬁned.
The optimization of AReLU then follows.
3. A TTENTION M ECHANISM
Let us denote V
= { v i } ∈ R D 1 v × D 2 v ×··· a tensor representing input data or feature volume.
Function Φ , parameterized by Θ = { θ i } , is used to compute an attention map S
= { s i } ∈ R D
θ (1) v × D θ (2) v ×···
[Figure]
over a subspace of V (let θ ( · ) denote a correspondence function for the indices of dimension):
[Equation]
Φ can be implemented by a neural network with Θ being its learnable parameters.
We can modulate the input V with the attention map S using a function Ψ , obtaining the output U
[Equation]
Ψ is an element-wise multiplication.
In order to perform element-wise multiplication, one needs to ﬁrst extend S to the full dimension of V .
We next review various attention mechanisms with attention map at different granularities.
Figure 1(left) gives an illustration of various attention mechanisms.
Vector-wise Attention Mechanism In NLP, attention maps are usually computed over different word vectors.
In this case, V = { v i } ∈ R N × D represents a sequence of N feature vectors with dimension
D .
S =
{ s i } ∈ R N is a sequence of attention values for the corresponding vectors.
Channel-wise Attention Mechanism In computer vision, a feature volume V = { v i } ∈ R W × H × C has a spatial dimension of W × H and a channel dimension of C .
S =
{ s i } ∈ R C is an attention map over the C channels.
All elements in each channel share the same attention value.
Spatial-wise Attention Mechanism Considering again V
= { v i } ∈ R W × H × C with a spatial dimension of W × H .
S =
{ s i } ∈ R W × H is an attention map over the spatial dimension.
All channels of a given spatial location share the same attention value.
Element-wise Attention Mechanism Given a feature volume V = { v i } ∈ R W × H × C containing
W × H × C elements, we compute an attention map over the whole volume (all elements), i.e.,
S =
{ s i } ∈ R W × H × C , so that each element has an independent attention value.
3. EL EMENT - WISE S IGN - BASED A TTENTION (ELSA)
We propose, ELSA, a new kind of element-wise attention mechanism which is used to deﬁne our attention-based activation.
Considering a feature volume V = { v i } ∈ R W × H × C , we compute an element-wise attention map S
=
{ s i } ∈ R W × H × C :
[Equation]
where Θ = { α, β } ∈ R 2 is learnable parameters.
C ( · ) clamps the input variable into [0 . 01 , 0 . 99] .
σ is the sigmoid function.
The modulation function Ψ is deﬁned as:
[Equation]
In ELSA, positive and negative elements receive different amount of attention determined by the two parameters α and β , respectively.
Therefore, it can also be regarded as sign-wise attention mechanism.
With only two learnable parameters, ELSA is light-weight and easy to learn.
3. AR E LU: A TTENTION - BASED R ECTIFIED L INEAR U NITS
We represent the function Φ in ELSA with a network layer with learnable parameters α and β :
[Equation]
where X = { x i } is the input of the current layer.
In constructing an activation function with ELSA, we combine it with the standard Rectiﬁed Linear Units
[Equation]
Adding them together leads to a learnable activation function:
[Equation]
This combination ampliﬁes positive elements and suppresses negative ones based on the learned scaling parameters β and α , respectively.
Thus, ELSA learns an element-wise residue for the activated elements w.r.t.
ReLU which is an identity transformation, which helps ameliorate gradient vanishing.
3. T HE O PTIMIZATION OF AR E LU
AReLU can be trained using back-propagation jointly with all other network layers.
The update formulation of α and β can be derived with the chain rule.
Speciﬁcally, the gradient of α is:
[Equation]
[Equation]
where E is the error function to be minimized.
The term ∂ E ∂ F ( x i ,α,β ) is the gradient propagated from the deeper layer.
The gradient of the activation of X with respect to α is given by:
[Equation]
Here, the derivative of the clamp function C ( · ) is handled simply by detaching the gradient backpropagation when α < 0 .
01 or α > 0 . 99 .
The gradient of the activation of X with respect to β is:
[Equation]
The gradient of the activation with respect to input x i by:
[Equation]
It can be found that AReLU ampliﬁes the gradients propagated from the downstream when the input is activated since 1 + σ ( β ) > 1 ; it suppresses the gradients otherwise.
On the contrary, there is no such ampliﬁcation effect in the standard ReLU and its variants (e.g., PReLu (He et al., 2015)) — only suppression is available.
The ability to amplify the gradients over the activated input helps avoiding gradient vanishing, and thus speeds up the training convergence of the model (see Figure 3).
Moreover, the ampliﬁcation factor is learned to dynamically adapt to the input and is conﬁned with the sigmoid function.
This makes the activation more data-adaptive and stable (see Figure 1(right) for a visual comparison of post-activation feature maps by AReLU and ReLU).
The suppression part is similar to PReLu which learns the suppression factor for ameliorating zero gradients.
AReLU introduces a very small number of extra parameters which is 2 L for an L -layer network.
The computational complexity due to AReLU is negligible for both forward and backward propagation.
Note that the gradients of α and β depend on the entire feature volume X .
This means that ELSA can be regarded as a global attention mechanism: Although the attention map is computed in an element-wise manner, the parameters are learned globally accounting for the impact of the full feature volume.
This makes our AReLU more data-adaptive and hence the whole network more expressive.
[Figure]
E XPERIMENTS
We ﬁrst study the robustness of AReLU in terms of parameter initialization.
We then evaluate convergence of network training with different activation functions on two standard classiﬁcation benchmarks (MNIST (LeCun et al., 1998)) and CIFAR100 (Krizhevsky et al., 2009).
We compare
AReLU with 18 different activation functions including 13 non-learnable ones and 5 learnable ones; see the list in Table 1.
The number of learnable parameters for each learnable activation function are also given in the table.
In the end, we also demonstrate the advantages of AReLU in transfer learning.
Please refer to supplemental material for more results and experiments details.
4. I NITIALIZATION OF L EARNABLE P ARAMETER α AND β
For evaluation purpose, we design a neural network (MNIST-Conv) with three convolutional layers each followed by a max-pooling layer and an AReLU, and ﬁnally a fully connected layer followed by a softmax layer.
Details of this network can be found in the supplemental material.
The experiment on parameter initialization is conducted with MNIST-Conv over the MNIST dataset.
As shown in
Figure 2(a), AReLU is insensitive to the initialization of α and β .
Different initial values result in close convergence rate and classiﬁcation accuracy.
Generally, a large initial value of β can speed up the convergence.
Figure 2(b) shows the learning procedure of the two parameters and (c) plots the learned ﬁnal AReLU’s for the three convolutional layers.
In the following experiments, we initialize α = 0 .
9 and β = 2 . 0 by default.
4. C ONVERGENCE ON MNIST
On the MNIST dataset, we evaluate MNIST-Conv implemented with different activation functions and trained with the ADAM or SGD optimizer.
The activation function is placed after each maxpooling layers.
We compare AReLU with both learnable and non-learnable activation functions under different learning rates of 1 × 10 − 2 , 1 × 10 − 3 , 1 × 10 − 4 , and 1 × 10 − 5 .
To compare the convergence speed of different activation functions, we report the accuracy after the ﬁrst epoch, again taking the mean over ﬁve times training; see Table 1.
In the table, we report the improvement of AReLU over the best among other non-learnable and learnable methods. In Figure 3, we plot the mean accuracy over increasing number of training epochs. As shown in Table 1, AReLU outperforms most existing non-learnable and learnable activation functions in terms of convergence speed and ﬁnal classiﬁcation accuracy on MNIST. A note-worthy phenomenon is that AReLU can achieve a more effective training with a small learning rate (see the signiﬁcant improvement when the learning rate is 1 × 10 − 4 or 1 × 10 − 5 ) than the alternatives.
This can also be observed from Figure 3.
Generally, smaller learning rates would cause lower learning efﬁciency since the vanishing gradient issue is intensiﬁed in such case. AReLU can overcome this difﬁculty thanks to its gradient ampliﬁcation effect. Efﬁcient learning with a small learning rate is very useful in transfer learning where a pre-trained model is usually ﬁne-tuned on a new domain/dataset [Table]
A cc u r a c y
Epoch
A cc u r a c y
Epoch
CELU
ELU
GELU
LReLU
RReLU
Maxout
ReLU ReLU
SELU
Sigmoid
Softplus
Swish
SLAF
Tanh
APL
Combine
PAU
PReLU
AReLU with a small learning rate which is difﬁcult for most existing deep networks. Section 4.4 will demonstrate this application of AReLU.
4. C ONVERGENCE ON CIFAR
In order to better demonstrate the effect of ELSA, we regard the ReLU, without ELSA, as our baseline. For plot clarity, we choose to compare only with those most representative competitive activation functions including PAU, SELU, ReLU, LReLU (LReLU), and PReLU.
More results can be found in the supplemental material.
We evaluate the performance of AReLU with ﬁve different mainstream network architectures on CIFAR100.
We use the SGD optimizer and follow the training conﬁguration in (Pereyra et al., 2017): The learning rate is 0 . 1 , the batch size is 64 , the weight decay is 5 × 10 − 4 , and the momentum is 0 . A cc u r a c y
Epoch
MobileNet-v        0
A cc u r a c y
Epoch
PreActNet0
A cc u r a c y
Epoch
ResNet0
A cc u r a c y
Epoch
SeResNet0
A cc u r a c y
[Figure] [Figure]
The results are plotted in Figure 4.
Learnable activation functions generally have a faster convergence compared to non-learnable ones.
AReLU achieves a faster convergence speed for all the ﬁve network architectures.
It is worth to note that though PAU can achieve a faster convergence at the beginning in some networks such as SeResNet-18, it tends to overﬁt later with a fast saturation of accuracy.
AReLU avoids such overﬁtting with smaller number of parameters than PAU (2 vs 10).
We also conduct a qualitative analysis of AReLU by visualizing the learned feature maps using
Grad-CAM Selvaraju et al.
(2017) using testing images of CIFAR100.
Grad-CAM is a recently proposed network visualization method which utilizes gradients to depict the importance of the spatial locations in a feature map.
Since gradients are computed with respect to a speciﬁc image class, Grad-CAM visualization can be regarded as a task-oriented attention map.
In Figure 5, we visualize the ﬁrst-layer feature map of ResNet-18.
As shown in the ﬁgure, the feature maps learned with AReLU leads to semantically more meaningful activation of regions with respect to the target class.
This is due to the data-adaptive, attentive ability of AReLU.
4. P ERFORMANCE IN T RANSFER L EARNING
We evaluate transfer learning of MNIST-Conv with different activation functions between two datasets:
MNIST and SVHN 1 .
The data preprocessing for adapting the two datasets follows (Shin et al., 2017).
We train three models and test them on SVHN: 1) one is trained directly on SVHN without any pretraining, 2) one trained on MNIST but not ﬁnetuned on SVHN, and 3) one pretrained on MNIST and ﬁnetuned on SVHN.
In pretraining, we train MNIST-Conv using SGD with a learning rate of 0 .
01 for 20 epochs which is sufﬁcient for all model variants to converge.
In ﬁnetuning, we train the model on SVHN with a learning rate of 1 × 10 − 5 , using SGD optimizer for 100 epochs.
The testing results on SVHN are reported in Table 2 where we compare AReLU with several competitive alternatives.
Without pretraining, it is hard to obtain a good accuracy on the difﬁcult task of SVHN.
Nevertheless, MNIST-Conv with AReLU performs the best among all alternatives; some activation functions even failed in learning.
In the setting of transfer learning (pretrain + ﬁnetune),
AReLU outperforms all other activation functions for different amount of pretraining, thanks to it high learning efﬁciency with small learning rates. [Table] [Table]
4. P ERFORMANCE IN M ETA L EARNING
We evaluate the meta learning performance of MNIST-Conv with the various activation functions based on the MAML framework Finn et al.
MAML is a fairly general optimization-based algorithm compatible with any model that learns through gradient descent.
It aims to obtain metalearning parameters from similar tasks and adapt the parameters to novel tasks with the same distribution using a few gradient updates.
In MAML, model parameters are explicitly trained such that a small number of gradient updates over a small amount of training data from the novel task could lead to good generalization performance on that task.
We expect that the fast convergence of AReLU would help MAML to adapt a model to a novel task more efﬁciently and with better generalization.
We set the fast adaption steps as 5 and use 32 tasks for each steps.
We train the model for 100 iterations with a learning rate of 0 .
We report in Table 3 the ﬁnal test accuracy for different activation functions on a 5-ways-1-shots task and a 5-ways-5-shots task, respectively.
The results show that AReLU shows clear advantage compared to the alternative activation functions.
One noteworthy phenomenon is the performance of PAU (Molina et al., 2019):
It performs well in other evaluations but not on meta learning which is probably due to its overﬁtting-prone nature.
4. T HE GENERALIZED EFFECT OF ELSA
In this experiemnt, we show that ELSA (Element-wise Sign-based Attention) can serve as a general module which can be plugged in to any existing activation function and obtain a performance boost for most cases.
We deﬁne a new activation F (cid:48) the same as Eq.
(7), but replace the ReLU function R with speciﬁed activation function.
We keep the same experiment settings as Sec. 4.2. As shown in Tab: 4, after plugging with a ELSA module, we can obtain a performance boost for most cases compared with Table 1, indicating the well generalized effect of ELSA.  C ONCLUSION
We have presented AReLU, a new learnable activation function formulated with element-wise signbased attention mechanism.
Networks implemented with AReLU can better mitigate the gradient vanishing issue and converge faster with small learning rates.
This makes it especially useful in transfer learning where a pretrained model needs to be ﬁnetuned in the target domain with a small learning rate.
AReLU can signiﬁcantly boost the performance of most mainstream network architectures with only two extra learnable parameters per layer introduced.
In the future, we would like to investigate the application/extension of AReLU to more diverse tasks such as object detection, language translation and even structural feature learning with graph neural networks.
[Table]