Year: 2021

Abstract:
The training process of neural networks usually optimizes weights and bias parameters of linear transformations, while nonlinear activation functions are prespecified and fixed. This work develops a systematic approach to constructing matrix activation functions whose entries are generalized from ReLU. The activation is based on matrix-vector multiplications using only scalar multiplications and comparisons. The proposed activation functions depend on parameters that are trained along with the weights and bias vectors. Neural networks based on this approach are simple and efficient and are shown to be robust in numerical experiments.

heading: 1 INTRODUCTION
text: In recent decades, deep neural networks (DNNs) have achieved significant successes in many fields such as computer vision and natural language processing (Voulodimos et al., 2018), (Otter et al., 2018). The DNN surrogate model is constructed using recursive composition of linear transformations and nonlinear activation functions. To achieve good performance, it is essential to choose activation functions suitable for specific applications. In practice, Rectified Linear Unit (ReLU) is one of the most popular activation functions for its simplicity and efficiency. A drawback of ReLU is the presence of vanishing gradient in the training process, known as dying ReLU problem (Lu et al., 2019). Several relatively new activation approaches are proposed to overcome this problem, e.g., the simple Leaky ReLU, and Piecewise Linear Unit (PLU) (Nicolae, 2018), Softplus (Glorot et al., 2011), Exponential Linear Unit (ELU) (Clevert et al., 2016), Scaled Exponential Linear Unit (SELU) (Klambauer et al., 2017), Gaussian Error Linear Unit (GELU) (Hendrycks & Gimpel, 2016).
Although the aforementioned activation functions are shown to be competitive in benchmark tests, they are still fixed nonlinear functions. In a DNN structure, it is often hard to determine a priori the optimal activation function for a specific application. In this paper, we shall generalize ReLU and introduce arbitrary trainable matrix activation functions. The effectiveness of the proposed approach is validated using function approximation examples and well-known benchmark datasets such as CIFAR-10 and CIFAR-100. There are a few classical works on adaptively tuning of parameters in the training process, e.g., the parametric ReLU (He et al., 2015b). However, our adaptive matrix activation functions are shown to be competitive and more robust in those experiments.

heading: 1.1 PRELIMINARIES
text: We consider the general learning process based on a given training set {(xn, fn)}Nn=1, where the inputs {xn}Nn=1 ⊂ Rd and outputs {fn}Nn=1 ⊂ RJ are implicitly related via an unknown target function f : Rd → RJ with fn = f(xn). The ReLU activation function is a piecewise linear function given by
σ(t) = max {t, 0} , for t ∈ R. (1) In the literature σ is acting component-wise on an input vector. In a DNN, let L be the number of layers and n` denote the number of neurons at the `-th layer for 0 ≤ ` ≤ L with n0 = d and nL = J . Let W = (W1,W2, . . . ,WL) ∈ ∏L `=1 Rn`×n`−1 denote the tuple of admissible weight
matrices and B = (b1, b2, . . . , bL) ∈ ∏L `=1 Rn` the tuple of admissible bias vectors. The ReLU
DNN approximation to f at the L-th layer is recursively defined as
ηL = ηL,W,B := hWL,bL ◦ σ ◦ · · · ◦ hW2,b2 ◦ σ ◦ hW1,b1 , (2)
where hW`,b`(x) = W`x + b` is an affine linear transformation at the `-th layer, and ◦ denotes the composition of functions. The traditional training process for such a DNN is to find optimal W∗ ∈ ∏L `=1 Rn`×n`−1 , B∗ ∈ ∏L `=1 Rn` , (and thus optimal ηL,W∗,B∗ ) such that
(W∗,B∗) = argmin W,B
1
N N∑ n=1 |fn − ηL,W,B(xn)|2 . (3)
In other words, ηL,W∗,B∗ best fits the data with respect to the ` 2 norm within the function class {ηL,W,B}. In practice, the sum of squares norm used in minimization could be replaced with other convenient norms.

heading: 2 TRAINABLE MATRIX ACTIVATION FUNCTION
text: Having a closer look at ReLU σ, we observe that the activation σ ◦ η`(x) = σ(η`(x)) could be realized as a matrix-vector multiplication σ ◦ η`(x) = D`(η`(x))η`(x) or equivalently σ ◦ η` = (D` ◦ η`)η`, where η` is column vector-valued in Rn` and D` is a diagonal matrix-valued function mapping from Rn` to Rn`×n` with entries from the discrete set {0, 1}. This is a simple but quite useful observation. There is no reason to restrict on {0, 1} and we thus look for a larger set of values over which the diagonal entries of D` are running or sampled. With slight abuse of notation, our new DNN approximation to f is calculated using the following recurrence relation
η1 = hW1,b1 , η`+1 = hW`+1,b`+1 ◦ [ (D` ◦ η`)η` ] , ` = 1, . . . , L− 1. (4)
Here each D` is diagonal and is of the form
D`(y) = diag(α`(y1), α`(y2), . . . , α`(yn`)), y ∈ Rn` , (5)
where α` is a nonlinear function to be determined. Since piecewise constant functions can approximate a continuous function within arbitrarily high accuracy, we choose α` to be the following step function
α`(s) =  t`,0, s ∈ (−∞, s1], t`,1, s ∈ (s1, s2],
... t`,m−1, s ∈ (sm−1, sm], t`,m, s ∈ (sm,∞),
(6)
where m is a positive integer and {t`,j}mj=0 and {sj}mj=1 are constants. For the time being, let D`(η`(x))η`(x) = σ`(η`(x)) with σ` being a scalar-valued function applied to η`(x) componentwise. In the following, we list several σ`s corresponding to special α` and visualize them in Figure 1.
m = 1, s1 = 0, t`,0 = 0, t`,1 = 1 =⇒ σ` is ReLU, m = 1, s1 = 0, t`,0 > 0, t`,1 = 1 =⇒ σ` is Leaky ReLU, m = 2, s1 = 0, s2 = 1, t`,0 = t`,2 = 0, t`,1 = 1 =⇒ σ` is discontinuous.
In our case, we will choose the grid points {sj}mj=1 a priori and train the step values ∪L`=1{t`,j}mj=0 of α`. In such a way, the resulting DNN may use different activation functions for different layers, and these activation functions are naturally adapted to the target function f and the target dataset. Since ReLU and Leaky ReLU are included by our DNN as special cases, the proposed DNN is clearly not worse than the traditional ones in practice. In the following, we call the neural network in equation 4, with the activation approach given in equation 5 and equation 6, a DNN based on the “Trainable Matrix Activation Function (TMAF)”.
REMARK
The activation functions used in TMAF neural network are not piecewise constants. Instead, TMAF activation is realized using matrix-vector multiplication, where entries of those matrices are trainable piecewise constants. There exist several trainable activation functions in the literature, e.g., Parametric ReLU (He et al., 2015b), Padé Activation Unit (PAU) (Molina et al., 2020), Piecewise Linear Unit (PWLU) (Zhou et al., 2021), Swish (Ramachandran et al., 2017), adaptive rational functions (Boullé et al., 2020), adaptive splines (Bohra et al., 2020), adaptive activation functions for physics informed neural networks (ada, 2020). We point out that TMAF is different from the aforementioned activation approaches. One might think TMAF is similar to PWLU because they are both defined piecewise. However, by construction TMAF includes discontinuous functions as special cases (see Figure 1) while PWLU always results in a continuous piecewise linear function with trainable slopes.
Starting from the diagonal activation D`, we can go one step further to construct more general activation matrices. First we note thatD` could be viewed as a nonlinear operator T` : [C(Rd)]n` → [C(Rd)]n` , where
[T`(g)](x) = D`(g(x))g(x), g = (g1, . . . , gn`) > ∈ [C(Rd)]n` , x ∈ Rd.
There seems to be no special reason, aside from keeping the computational cost as low as possible, to consider only diagonal operators D. A more ambitious approach is to consider a trainable nonlinear activation operator determined by more general matrices, for example, by a tri-diagonal operator
[T`(g)](x) =  α`(g1(x)) β`(g2(x)) 0 · · · 0 γ`(g1(x)) α`(g2(x)) β`(g3(x)) · · · 0 ... . . . . . . . . .
... 0 0 · · · α`(gn`−1(x)) β`(gn`(x)) 0 0 · · · γ`(gn`−1(x)) α`(gn`(x))  g(x), (7) for x ∈ Rd. The diagonal entry α` is given in equation 6 while the off-diagonal entries β`, γ` are piecewise constant functions defined in a fashion similar to α`. Theoretically speaking, even trainable full matrix activation is possible despite of potentially huge training cost. In summary, a DNN based on trainable nonlinear activation operators {T`}L`=1 reads
η1 = hW1,b1 , η`+1 = hW`+1,b`+1 ◦ T`(η`), ` = 1, . . . , L− 1. (8)
As we indicated above, TMAF can be used with any matrix and this will lead to a great flexibility in training and approximation. In fact, if a sparse matrix D is used, so that the matrix-vector multiplication is proportional to the number of rows in the matrix, then the resulting algorithm will be computationally efficient. We can choose p-diagonal matrices, as well as matrices with prescribed sparsity patterns. Clearly, based on the simple examples we have considered here, one may conclude that such possibilities, when investigated in depth, can increase the robustness and the accuracy of the TMAF-NN.
REMARK
Our observation also applies to an activation function σ other than ReLU. For example, we may rescale σ(x) to obtain σ(ωi,`x) using a set of constants {ωi,`}1≤i≤n`,1≤`≤L varying layer by layer and neuron by neuron. Then σ(ωi,`x) are used to form a matrix activation function and a TMAF DNN, where {ωi,`} are trained according to given data and are adapted to the target function.

heading: 3 NUMERICAL RESULTS
text: In this section, we demonstrate the feasibility and efficiency of TMAF by comparing it with traditional ReLU-type activation functions. Recall that neurons in the `-th layer will be activated by the matrix D`. In principle, all parameters in equation 6 are allowed to be trained. To ensure practical efficiency, we shall fix interval grid points {sj}mj=1 used in α` and only let function values {t`,j} in equation 6 be trained in the following. In each experiment, we use the same learning rates, stochastic optimization methods, and number NE of epochs (optimization iterations). In particular, the learning rate is 1e-4 from epoch 1 to NE2 and 1e-5 is used from epoch NE 2 + 1 to NE. The optimization method is ADAM (Kingma & Ba (2015)) based on mini-batches of size 500. Numerical experiments are performed in PyTorch (Paszke et al. (2019)). We provide two sets of numerical examples:
• Function approximations by TMAF networks and ReLU-type networks; • Classification problems for MNIST and CIFAR set solved by TMAF and ReLU networks.
For the first class of examples we use the `2-loss function as defined in equation 3. For the classification problems we consider the cross-entropy that is widely used as a loss function in classification models. The cross entropy is defined using a training set having p images, each withN pixels. Thus, the training dataset is equivalent to the vector set {zj}pj=1 ⊂ RN with each zj being an image. The j-th image belongs to a class cj ∈ {1, . . . ,M}. The neural network maps zj to xj ∈ RM ,
xj := ηL,W,B(zj) ∈ RM , zj ∈ RN , j = 1, . . . , p.
The cross entropy loss function of ηL,W,B then is defined by
E(W,B) = −1 p p∑ k=1 log ( exp(xck,k)∑M j=1 exp(xj,k) ) .

heading: 3.1 APPROXIMATION OF A SMOOTH FUNCTION
text: As our first example, we use neural networks to approximate
f(x1, · · · , xn) = sin(πx1 + · · ·+ πxn), xk ∈ [−2, 2], k = 1, . . . , n.
The training datasets consist of 20000 input-output data pairs where the input data are randomly sampled from the hypercube [−2, 2]n based on uniform distribution. The neural networks have single or double hidden layers. Each layer (except input and output layers) has 20 neurons. For TMAF D` in equation 5, the function α` uses intervals (−∞,−1.4), (−1.4,−0.92], (−0.92,−0.56], (−0.56,−0.26], (−0.26, 0], (0, 0.26], (0.26, 0.56], (0.56, 0.92], (0.92, 1.4], (1.4,∞) such that probability over each of the ten intervals is 0.1 with respect to Gaussian distribution. Moreover, we apply BatchNorm1d in PyTorch to the linear output of neural networks in each hidden layer. The approximation results are shown in Table 1 and Figures 2–3, where Para-ReLU stands for the parametric ReLU neural network. It is observed that TMAF is the most accurate activation approach in these examples.

heading: 3.2 APPROXIMATION OF AN OSCILLATORY FUNCTION
text: The next example is on approximating the following function having high, medium and low frequency components f(x) = sin(100πx) + cos(50πx) + sin(πx), (9) see Figure 4a for an illustration. In fact, the function in equation 9 is rather difficult to capture by traditional approximation methods as it is highly oscillatory. We consider ReLU, parametric ReLU,
and diagonal TMAF neural networks with 3 hidden layers and 400 neurons per layer (except the first and last layers). We also consider the tri-diagonal TMAF (denoted by Tri-diag TMAF, see equation 7) with 3 hidden layers and 300 neurons per layer. The training datasets are 20000 inputoutput data pairs where the input data are randomly sampled from the interval [−1, 1] based on uniform distribution.
The diagonal TMAF uses α = α` in equation 6 with intervals (−∞,−1], (−1+kh,−1+(k+1)h], (1,∞) for h = 0.02, 0 ≤ k ≤ 99. The tri-diagonal TMAF is given in equation 7, where α` is the same as the diagonal TMAF, while β` is piecewise constant with respect to intervals (−∞,−2.01+ h], { (−2.01 + kh,−2.01 + (k + 1)h] }99 k=0
, (−0.01,∞), and γ` is a piecewise constant based on (−∞, 0.01], { (0.01 + kh, 0.01 + (k + 1)h] }99 k=0
, (2.01,∞), respectively. Numerical results could be found in Figures 4b, 5, 6 and Table 2.
For this challenging problem, we note that the diagonal TMAF and tri-diagonal TMAF produce high-quality approximations (see Figures 5 and 4b) while ReLU and parametric ReLU are not able to approximate the highly oscillating function within reasonable accuracy, see Figure 4b and Table 2. Moreover, it is observed from Figure 6 that ReLU and parametric ReLU actually approximate the low frequency part of the target function. To capture the high frequency, ReLU-type neural networks are clearly required to use much more neurons, introducing significantly amount of weight and bias parameters. On the other hand, increasing the number of intervals in TMAF only lead to a few more training parameters.
3.3 CLASSIFICATION OF MNIST AND CIFAR DATASETS
We now test TMAF by classifying images in the MNIST, CIFAR-10 and CIFAR-100 dataset.
For the MNIST set, we implement double layer fully connected networks defined as in equation 2 and equation 4 with 10 neurons per layer (except at the first layer n0 = 764), and we use ReLU or
diagonal TMAF as described in Subsection 3.1. Numerical results are shown in Figures 7a, 7b and Table 3. We observe that performance of the TMAF and the ReLU networks are similar. Such a behavior clearly should be expected for a simple dataset such as MNIST.
For the CIFAR-10 dataset, we use the ResNet18 network structure provided by (He et al., 2015a). The activation functions are still ReLU and the diagonal TMAF used in Subsection 3.1. Numerical results are presented in Figures 8a, 8b and Table 3. Those parameters given in (Paszke et al., 2019) are already tuned well with respect to ReLU. Nevertheless, TMAF still produces smaller errors in the training process and returns better classification results in the evaluation stage, see Table 3.
For the CIFAR-100 dataset, we use the ResNet34 network structure provided by (He et al., 2015a) with the ReLU and TMAF activation functions in Subsection 3.1. Numerical results are presented in Figures 9a and 9b. In the training process, TMAF again outperforms the classical ReLU network.
It is possible to improve the performance of TMAF applied to those benchmark datasets. The key point is to select suitable intervals in α` to optimize the performance. A simple strategy is to let those intervals in equation 6 be varying and adjusted in the training (or the evaluation) process, which will be investigated in our future research.

Review summaries:
This paper proposes a new type of activation function, called Trainable Matrix Activation Functions (TMAF), to replace the existing activation functions in neural networks, such as ReLU. TMAF is realized using matrix-vector multiplications, where entries of the matrices are trainable piecewise constants. Empirical studies show that TMAF can approximate the sin/cos type of functions better than ReLU/PReLU in terms of approximation error. TMAF also outperforms ReLU in MNIST/CIFAR classifications tasks.

The authors propose a new approach to activation functions in DNNs. More precisely, the authors present the trainable matrix activation function (TMAF), which is an activation realized by a matrix-vector multiplication whose entries are trainable. The authors then present a diagonal and tri-diagonal operators. Here, these matrices represent a coefficient that multiplies the output of the layer of the neural network and thus can represent any general piecewise linear activation including relu and leaky relu.  These matrices are trainable and therefore add a small number of extra parameters to the network depending on the architecture of the layers.  Finally the authors demonstrate the benefits of TMAF in synthetic and real world datasets on a couple of architectures.

The paper introduces a trainable matrix of activation functions. The authors propose to replace activation with a custom learnable piecewise linear function. The results are evaluated on a custom sine function with different oscillations as well as CIFAR-10 and CIFAR-100.   

The paper describes an approach to train activation functions. The main idea is to rewrite ReLU as s(x)x (where s(x) is either 0 or 1 depending on the value of x), and then generalize this with a piecewise constant s(x). The paper also proposes a generalization where values at different neurons can be combined. Finally, this is evaluated on two artificial problems and two standard computer vision benchmarks.