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