N EURAL NETWORKS WITH TRAINABLE MATRIX ACTIVATION FUNCTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
The training process of neural networks usually optimizes weights and bias parameters of linear transformations, while nonlinear activation functions are prespeciﬁed and ﬁxed.
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 efﬁcient and are shown to be robust in numerical experiments
I NTRODUCTION
In recent decades, deep neural networks (DNNs) have achieved signiﬁcant successes in many ﬁelds 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 speciﬁc applications.
In practice, Rectiﬁed Linear Unit (ReLU) is one of the most popular activation functions for its simplicity and efﬁciency.
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 ﬁxed nonlinear functions.
In a DNN structure, it is often hard to determine a priori the optimal activation function for a speciﬁc 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.
1. P RELIMINARIES
We consider the general learning process based on a given training set { ( x n , f n ) } Nn =1 , where the inputs { x n } Nn =1 ⊂ R d and outputs { f n } Nn =1 ⊂ R J are implicitly related via an unknown target function f : R d → R J with f n = f ( x n ) .
The ReLU activation function is a piecewise linear function given by
[Equation]
In the literature σ is acting component-wise on an input vector.
In a DNN, let L be the number of layers and n (cid:96) denote the number of neurons at the (cid:96) -th layer for 0 ≤ (cid:96) ≤ L with n 0 = d and n L = J .
Let W =
( W 1 , W 2 , . . .
, W L ) ∈ (cid:81) L(cid:96)
=1 R n (cid:96) × n (cid:96) − 1 denote the tuple of admissible weight matrices and B = (
b 1 , b 2 , . . .
, b L ) ∈ (cid:81) L(cid:96)
=1 R n (cid:96) the tuple of admissible bias vectors.
The ReLU
DNN approximation to f at the L -th layer is recursively deﬁned as
[Equation]
where h W (cid:96) ,b (cid:96) ( x ) = W (cid:96) x
+ b (cid:96) is an afﬁne linear transformation at the (cid:96) -th layer, and ◦ denotes the composition of functions.
The traditional training process for such a DNN is to ﬁnd optimal W ∗ ∈ (cid:81)
[Equation] =1 R n (cid:96) × n (cid:96) [Equation]
1 , B [Equation] ∈ (cid:81) L(cid:96) =
[Equation]
n [Equation] , (and thus optimal η L, W ∗
[Equation]
B ∗ ) such that [Equation] In other words, η L, W ∗ , B ∗ best ﬁts the data with respect to the (cid:96) 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.
T RAINABLE MATRIX ACTIVATION FUNCTION
Having a closer look at ReLU σ, we observe that the activation σ ◦ η (cid:96) ( x ) =
σ ( η (cid:96) ( x )) could be realized as a matrix-vector multiplication σ ◦ η (cid:96) ( x ) =
D (cid:96) ( η (cid:96) ( x ))
η (cid:96) ( x ) or equivalently σ ◦ η (cid:96) =
( D (cid:96) ◦
η
(cid:96) )
η (cid:96) , where η (cid:96) is column vector-valued in R n (cid:96) and D (cid:96) is a diagonal matrix-valued function mapping from R n (cid:96) to R n (cid:96) × n (cid:96) 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 (cid:96) are running or sampled.
With slight abuse of notation, our new DNN approximation to f is calculated using the following recurrence relation
[Equation]
Here each D (cid:96) is diagonal and is of the form
[Equation]
where α (cid:96) is a nonlinear function to be determined.
Since piecewise constant functions can approximate a continuous function within arbitrarily high accuracy, we choose α (cid:96) to be the following step function
[Equation]
where m is a positive integer and { t (cid:96),j } mj =0 and { s j } mj =1 are constants. For the time being, let D
[Equation]
x ))
[Equation]
x ) = [-1 , t 
-]σ [Equation] η [Equation]
x )) with σ (cid:96) being a scalar-valued function applied to η (cid:96) ( x ) componentwise. In the following, we list several σ (cid:96) s corresponding to special α (cid:96) and visualize them in Figure 1. [Equation]
In our case, we will choose the grid points { s j } m j =1 a priori and train the step values ∪ L (cid:96) =
1 { t (cid:96),j } m j = of α (cid:96) .
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)”. [Figure]
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), Pad´e Activation Unit (PAU)
(Molina et al., 2020), Piecewise
Linear Unit (PWLU)
(Zhou et al., 2021), Swish (Ramachandran et al., 2017), adaptive rational functions (Boull´e 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 deﬁned 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 (cid:96) , we can go one step further to construct more general activation matrices.
First we note that D (cid:96) could be viewed as a nonlinear operator T (cid:96) :
[ C ( R d )] n (cid:96) → [ C ( R d )] n (cid:96) , where
[Equation]
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 [Equation] g )]
[Equation]
x ) =  
[Equation]

 g
[Equation] [Equation] for x ∈ R d . The diagonal entry α (cid:96) is given in equation 6 while the off-diagonal entries β (cid:96) , γ (cid:96) are piecewise constant functions deﬁned in a fashion similar to α (cid:96) .
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 (cid:96) } L(cid:96) =1 reads
[Equation]
As we indicated above, TMAF can be used with any matrix and this will lead to a great ﬂexibility 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 efﬁcient. 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.
R EMARK
Our observation also applies to an activation function σ other than ReLU.
For example, we may rescale σ ( x ) to obtain σ ( ω i,(cid:96) x ) using a set of constants { ω i,(cid:96) } 1 ≤ i ≤ n (cid:96) , 1 ≤ (cid:96) ≤ L varying layer by layer and neuron by neuron.
Then σ ( ω i,(cid:96) x ) are used to form a matrix activation function and a TMAF DNN, where { ω i,(cid:96) } are trained according to given data and are adapted to the target function.
N UMERICAL RESULTS
In this section, we demonstrate the feasibility and efﬁciency of TMAF by comparing it with traditional ReLU-type activation functions.
Recall that neurons in the (cid:96) -th layer will be activated by the matrix D (cid:96) .
In principle, all parameters in equation 6 are allowed to be trained.
To ensure practical efﬁciency, we shall ﬁx interval grid points { s j } mj =1 used in α (cid:96) and only let function values { t (cid:96),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
-]NE2 + 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.
We provide two sets of numerical examples:
• Function approximations by TMAF networks and ReLU-type networks; • Classiﬁcation problems for MNIST and CIFAR set solved by TMAF and ReLU networks.
For the ﬁrst class of examples we use the (cid:96) 2 -loss function as deﬁned in equation 3.
For the classiﬁcation problems we consider the cross-entropy that is widely used as a loss function in classiﬁcation models.
The cross entropy is deﬁned using a training set having p images, each with N pixels.
Thus, the training dataset is equivalent to the vector set { z j } pj =1 ⊂ R N with each z j being an image.
The j -th image belongs to a class c j ∈ { 1 , . . .
, M } .
The neural network maps z j to x j ∈ R M , x j :=
η L, W , B ( z j ) ∈ R M , z j ∈ R N , j = 1 , . . . p.
The cross entropy loss function of η L, W , B then is deﬁned by
[Equation]
3. A PPROXIMATION OF A SMOOTH FUNCTION
As our ﬁrst example, we use neural networks to approximate [Equation]
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 (cid:96) in equation 5, the function α (cid:96) uses intervals ( −∞ , − 1 . 4) , ( − 1 . 4 , − 0 . 92] , ( − 0 .
92 , − 0 .
56] , ( − 0 .
56 , − 0 .
26] , ( − 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.
3. A PPROXIMATION OF AN OSCILLATORY FUNCTION
The next example is on approximating the following function having high, medium and low frequency components
[Equation]
see Figure 4a for an illustration.
In fact, the function in equation 9 is rather difﬁcult to capture by traditional approximation methods as it is highly oscillatory.
We consider ReLU, parametric ReLU, [Table] [Figure] and diagonal TMAF neural networks with 3 hidden layers and 400 neurons per layer (except the ﬁrst 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 α
= α (cid:96) 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 α (cid:96) is the same as the diagonal TMAF, while β (cid:96) is piecewise constant with respect to intervals ( −∞ , − 2 .
01 + h ] , (cid:8) ( − 2 .
01 + kh, − 2 .
+ ( k + 1) h ] (cid:9) 99 k =0 , ( − 0 .
01 , ∞ ) , and γ (cid:96) is a piecewise constant based on ( −∞ , 0 . 01] , (cid:8) (0 .
+ kh, 0 .
+ ( k + 1) h ] (cid:9) 99 k =0 , (2 . 01 , ∞ ) , respectively.
Numerical results could be found in Figures 4b, 5, 6 and Table 2.
[Figure]
[Table] 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 signiﬁcantly 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.
[Figure]
[Figure]
3. C LASSIFICATION 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 deﬁned as in equationand equation 4 with 10 neurons per layer (except at the ﬁrst layer n 0 = 764 ), and we use ReLU or [Figure] 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 classiﬁcation 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 α (cid:96) 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.
[Figure]