Neural networks with trainable matrix activation functions

1 INTRODUCTION . 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 . 1.1 PRELIMINARIES . 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 . 2 TRAINABLE MATRIX ACTIVATION FUNCTION . 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 .

Summary: 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.