Enhancing the trainability and expressivity of deep MLPs with globally orthogonal initialization
Abstract
Multilayer Perceptrons (MLPs) deﬁnes a fundamental model class that forms the backbone of many modern deep learning architectures.
Despite their universality guarantees, practical training via stochastic gradient descent often struggles to attain theoretical error bounds due to issues including (but not limited to) frequency bias, vanishing gradients, and stiff gradient ﬂows.
In this work we postulate that many of such issues ﬁnd origins in the initialization of the network’s parameters.
While the initialization schemes proposed by Glorot et al.
and He et al.
have become the de-facto choices among practitioners, their goal to preserve the variance of forward- and backward-propagated signals is mainly achieved by assumptions on linearity, while the presence of nonlinear activation functions may partially destroy these efforts.
Here, we revisit the initialization of MLPs from a dynamical systems viewpoint to explore why and how under these classical scheme, the MLP could still fail even at the beginning.
Drawing inspiration from classical numerical methods for differential equations that leverage orthogonal feature representations, we propose a novel initialization scheme that promotes orthogonality in the features of the last hidden layer, ultimately leading to more diverse and localized features.
Our results demonstrate that network initialization alone can be sufﬁcient in mitigating frequency bias and yields competitive results for high-frequency function approximation and image regression tasks, without any additional modiﬁcations to the network architecture or activation functions.
Introduction
It is now well understood that frequency bias [1, 2, 3] can prevent MLPs from approximating multiscale target functions, such as the ones that routinely arise when modeling physical and engineering systems governed by differential equations [4, 5].
While many novel architectures such as residual networks [6], Fourier feature networks [7, 8], and SIREN networks [9] have been proposed to enhance the trainability and expressivity of MLPs, less attention has been placed on the importance of network initialization, especially in the context of frequency bias.
Previous studies from Glorot et al.
He
[Figure]
et al.
and Sitzmann et al.
[10, 9, 11] have focused on studying how the input signal variance can be preserved from one layer to another, and deduce initialization schemes for the network weights to achieve so.
More recently
[Figure]
Cyr et al.
proposed the Box initialization scheme [12] for the robust initialization of MLPs with piece-wise linear activation functions
[Figure]
which preserves diversity in the last layer features, while Singh et al [13] use an additional regularization term to promote diversity in the ﬁrst layer features.
Inspired by these studies, we propose to model the initialization of deep
MLPs as a discrete dynamical system, study its stability, and deduce a novel initialization scheme that prevents vanishing information ﬂow and collapsed outputs
[Figure]
Variance ﬂow as a dynamical system
Consider a linear layer Y = XW , where X is the input vector, Y is the output vector, and W is the weight matrix.
W assume that the entries of W are independent and identically distributed (i.i.d) with zero mean and standard deviation c/ √ d in , where d in is the input dimension, while any bias
[Figure]
parameters are initialized to zero.
In the following text, c will be denoted as a “variance offset” or the “scale” of the initialization of W .
Then the output variance of the linear layer then takes the form
Var ( X i · W j,i ) =
c 2 E ( X 2 i ) /d in , while the mean is E ( X i · W j,i )
= 0 .
Moreover, the inner product
Y k = (cid:104) X, W · ,k (cid:105) =
(cid:80) d in i =
1 X
i W i,k converges in distribution to a normal distribution with zero mean and variance c
2 E ( X 2 i ) as d in → ∞ , by the Lindeberg–Lévy Central Limit Theorem [14], if E ( X 2 i ) is ﬁnite.
Consequently, we can use a normal distribution to approximate pre-activation distribution, and estimate the corresponding post-activation second moment as
[Equation]
where f i is the activation function, and s i is the square root of the second moment of the post activation, of the i th layer.
The system described by Eq.
2.1 can be simulated numerically for common choices of activation functions, as well as back-propagation (see Appendix B).
Fig.
1 conﬁrms the known behavior of vanishing second moment for hyperbolic tangent and ReLU activation functions, when c
= 1 as the conventional Glorot scheme [10] initializes, eventually leading to collapsed outputs especially for deeper MLPs (see Fig. 1(c)).
It also conﬁrms the preservation of the second moment in
ReLU networks initialized by the He scheme [11] where c = √ 2 .
To further illustrate the effect of the variance offset c , we analyze the ﬁxed points of Eq.
(2.1) for hyperbolic tangent networks, since the symmetry of tanh activation allows us to replace the second moment in Eq.
(2.1) with variance, and also replace s i with standard deviation.
As illustrated in
Fig. 2(a), our numerical results suggest that a bifurcation occurs when varying the variance offset c .
We numerically locate the new ﬁxed point and estimate its stability, and Fig. 2(b) shows that the larger variance offset c leads to two ﬁxed points, and the only stable one is the positive ﬁxed point, to which the variance of the hidden layer output with slightly larger variance offset converges (see Fig. 1(c)).
Clearly, an intervention that controls the onset of this bifurcation could lead to a drastic change to the network’s behavior and help prevent vanishing information ﬂows that leads to collapsed outputs.
In the next section we propose an initialization procedure that automatically tunes the layer-wise variance offset and bias, ensuring a diverse set of features in the network’s last layer prior to the training.
Globally Orthogonal (GO) initialization of MLPs
In the previous section we show that MLP initializations that do not carefully consider the effect of non-linear activations, will likely suffer from collapsing outputs.
To mitigate this adverse effect here we propose to use orthogonality in the last layer features as a measure to quantify the expressivity of an MLP representation.
This is motivated by the fact that orthogonality can quantify dissimilar behavior in the feature functions.
The concept of orthogonality is also widely adopted for representing functions in low-dimensional spaces among many classical numerical methods for solving differential equations by projecting the target function onto an orthogonal basis (e.g. spectral methods [15]).
Given a MLP f w , with m hidden layers, and weights w , we use g w to denote the post-activation outputs of the last hidden layer with d features.
Letting g jw denote the j th feature, the cosine of the angle θ i,j between g iw , g jw is deﬁned as
[Equation]
In practice, given input data X = ( x 1 , ..., x N ) , the inner product (cid:104)· , ·(cid:105) is evaluated with respect to the empirical data distribution p ( x ) = 1 N (cid:80) Ni =1 δ ( x − x i ) , with Dirac delta function δ .
Please note that we normalize the inner product of the features with the length of the features on multiple data points to avoid the scenario when the orthogonality loss on features with sharp angle that should otherwise be large remains to be small because of their lengths.
Cosine of angles between the last layer features functions are then transformed to a loss objective, by which the initialization variances and the biases of the each layer except the last layer can be calibrated.
Speciﬁcally, let the entries of the weight matrix of each layer be initialized with standard deviation c j / (cid:113)
d jin , j = 1 , ..., m , and let the bias of each layer be initialized as b j 1 , j = 1 , ..., m , we proceed by optimizing { c , b } by minimizing the following regularized loss
[Equation]
Intuitively, this orthogonality objective aims to reduce the correlation between the gradients of the parameters at different locations and therefore drastically changes the expressiveness of an MLP at initialization, as shown in Fig.
As opposed to previous studies that focus on the parameters orthogonality [16, 17], we directly optimize the orthogonality of the features in the function space, and therefore name it as global orthogonality objective.
In the next section, we will empirical analyze the capability of GO initialization to overcome the frequency bias from the perspective of the Neural
Tangent Kernel.
Empirical NTK analysis: The NTK for input locations x, y can be expressed as
[Equation]
where ∇ w denotes the gradient with respect to the parameter w .
Several studies [18, 19, 20] have shown that the NTK converges to a deterministic kernel as the width of the network increases to inﬁnity.
Under this setting, the MLP approximation can be asymptotically viewed as a kernel regression deﬁned by the NTK.
Furthermore, Geifman et al.
[21] show that MLPs with ReLU activations relate closely to the Laplace kernel.
We therefore believe that the behavior of the NTK at the initialization could be a representative characteristic.
A visualization of the resulting NTK matrix of a 512 × 8 MLP reveals the correlation between the predicted MLP outputs corresponding to different input locations when responding to a small perturbation in the network’s parameters.
We densely sample the nearby points at different input locations in the domain to approximate the kernel functions at each input point.
The kernel functions are then normalized to have unit area.
Fig.
[Figure]
shows the normalized kernel functions evaluated at an ordered grid of points for a wide and deep neural network initialized using Glorot initialization, and the proposed GO initialization, respectively.
Figure 4 empirically shows that MLPs initialized with the proposed GO scheme have a limiting kernel function whose distance parameter is smaller, as the regularizer coefﬁcient decreases.
Smaller distance parameters suggest less correlation between the prediction and the data output whose input location is further away, thus preventing frequency bias arising from kernels with an overly large distance parameter that tend to average the oscillations brought by the high frequency components and only retain low frequency components in the prediction.
Such tendency is also observed in previous works on the NTK of physics-informed neural networks (PINNs)
[5, 8] in the context of solving multi-scale differential equations.
Results
Synthetic data: In this example we demonstrate the effectiveness of GO initialization, on a 256 ×
MLP, in mitigating frequency bias and enabling the approximation of synthetic high frequency functions in two dimensions.
Instead of increasing the frequency of the target function, we increase the size of the training data-set exponentially in each training session, until the training and testing errors converge.
The training and testing errors associated with data-sets of various sizes are reported in Fig.
We in addition also test the orthonormal weights initialization, as proposed in[16, 17], where the square weight matrices are the orthonormal matrices returned from the QR factorization of the originally initialized weight matrices, normalized to the original Frobenious norm.
The MLPs with Glorot initialization fail to accommodate the ﬁtting of increasingly complex target functions, and eventually collapse with a training and testing relative L 2 errors of O (1) .
While the orthonormal initialization helps mitigate the collapsing problem, the large empirical conﬁdence intervals still
[Figure]
indicates that the MLP will still suffer from that problem with considerable probability.
On the contrary, the MLPs initialized with the GO scheme are able to capture broadband frequencies, and the empirical conﬁdence interval of the errors also suggests that the MLP initialized with the GO scheme is less sensitive to the initial points, and is therefore more robust.
However, we should acknowledge that the realization of the theoretical error bound and complexity of a given architecture could be heavily inﬂuenced by the choice of training hyperparameters, and our result in Fig.
5 shows that the initialization could be one of these crucial factors.
Image regression:
In this example, we will use a high deﬁnition RGB colored image.
The sample image contains rich details of a city landscape.
A continuous representation of the image is recovered by ﬁtting 25% of the pixels with total variation regularization.
As shown in Fig.
6, an MLP with
GO initialization is able to improve the SNR from 16 .
8 dB to 20 .
dB from the Glorot scheme, and recovers visually sharper details on light, shadow and texture that are on the contrary heavily distorted and smudged together in the representation trained from Glorot scheme.
A comparison of SNRs for different sampling ratios and different competing approaches can be found in Appendix C.5.
Summary
By deriving a simpliﬁed discrete dynamical system to quantify information ﬂow through MLPs at initialization, we identify and analyze vanishing output pathologies of conventional initialization schemes that hamper the trainability of deep fully-connected architectures.
Motivated by these observations, we propose a novel initialization procedure that is designed to promote global orthogonality in the last layer’s features, and demonstrate that such initialization leads to rich and localized feature representations.
Our numerical results conﬁrm that the proposed GO initialization can greatly enhance the trainability of deep MLPs and effectively mitigate frequency bias, without requiring any elaborate changes to the network architecture or activation functions.
Our ﬁndings motivate further investigation to elucidate the effect of orthogonal features as a mechanism to promote expressiveness of MLPs.
In future work we plan to extend our analysis to consider a wider range of activation functions, as well as loss objectives for different problems.
For example, for approximating solution of partial differential equations, the orthogonality objective may be augmented by the orthogonality of a set features that also shows diverse behavior after being transformed by a differential operator acting on the network’s outputs.
Acknowledgements
This work received support from DOE grant DE-SC0019116, AFOSR grant FA9550-20-1- and
DOE-ARPA grant DE-AR
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[N/A] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] The experiments remain in toy level and the computational cost is much below average.
We only include the computation cost for a few representative ones.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Appendix A Computational cost
GO initialization is typically carried out using a relatively small number of stochastic gradient descent iterations.
In all experiments reported in this work we employed 5 , 000 iterations to roughly calibrate the network’s weights and biases at initialization.
This typically amounts to a fraction of the total cost required for training the network to ﬁt observed data.
For example, applying GO initialization to an MLP with 4 hidden layers and 32 neurons per layer using a batch size of 256 takes roughlyseconds, while for a larger MLP with 16 hidden layers and 256 neurons per layer, the initialization takes roughly 1 minute for a larger batch size of 2 , 048 (all timings reported on a signle NVIDIA RTX
A6000 GPU).
We therefore consider the computational allocated to applying the GO initialization as negligible, compared to the cost of training.
Furthermore, the initialization cost could be further reduced by ﬁxing one draw of the parameters instead of generating different random draws of weights and biases at each iteration of optimizing the orthogonality objective.
B Back-propagation as a dynamical system
We analyze the backward pass of an MLP [22] as a dynamical system to show that the increased variance offset discussed in section 2 will not lead to vanishing gradients in MLPs activated with hyperbolic tangent activations.
To this end, a linear layer followed by a hyperbolic tangent activation can be expressed as Y = tanh( XW ) , where X ∈ R N × m 0 is the input of the hidden layer, and
W ∈ R m 0 × m 1 is the weight matrix.
Let dY =
∂l∂Y , i.e., the back propagated gradient with respect to the objective l .
We again assume that the entries of dY is distributed i.i.d.
Then the gradient of the weight matrix is
[Equation]
and the further back-propagated gradient is
[Equation]
To study the impact of the variance offset c on the back-propagated gradients dX , we choose to trace the square root of second moment of T X = tanh (cid:48) ( dX ) .
Let s TX =
(cid:112) E ( T 2 X ) ,
T Y = tanh (cid:48)
( T X W (cid:124) ) , and s TY
=
(cid:112) E ( T 2 Y ) .
Then, similarly, as m
= m 1 → ∞ , with the offset c , the distribution of dX is approximately normal with zero mean and standard deviation cs TX .
Therefore, we can write a discrete dynamical system for back-propagated gradients as
[Equation]
Again, considering the difﬁculty in obtaining the analytic form of the integration, we resort to the law of large numbers to estimate the expectation.
Furthermore, given that tanh (cid:48) ( x ) = 1 − tanh( x )which is an even function with monotonically decreasing component on R + , we have that
[Equation]
for all ﬁxed x , if σ 0 dX > σ 1 dX ≥ 0 .
Then since tanh (cid:48) ( x ) is bounded (and therefore integrable), the right hand side of (B.3) decreases monotonically for positive s TY .
And since (cid:112)
E (tanh (cid:48) (0 N (0 , 1)) 2 ) = 1 , the only ﬁxed point of the system is positive.
Fig.
7 shows the evolution of (B.3) at different variance offset value c s.
As before, we use a bisection method to roughly identify the ﬁxed points given different variance offsets, and observe their stability.
As shown in Fig. 7(b), our numerical results suggest that for a ﬁxed variance offset at reasonable range, the only positive ﬁxed point for the back-propagated gradient is stable.
Moreover, although the second moment shrinks for larger variance offset, it does not vanish.
0.0 0.5 1.0 1.5 2.0 s TX 0.
s TY offset = 1 offset =offset =0.0 0.5 1.0 1.5 2.0 s TX    d s TY / d s T X offset = 1 offset =offset =(a)  20 Offset 0.
F i x e d p o i n t  20 Offset 0.
D e r
i v a t
i v e a t t
h e f i x e d p o i n t
[Figure]
C Experimental setup
All numerical experiments presented in this work were performed in Python using the JAX library [23].
Unless otherwise speciﬁed, the optimizer is Adam with default settings [24].
The input and output of the data are normalized to zero mean and unit variance.
In the following we provide the necessary details for reproducing the main results and conclusions presented in the manuscript.
C.1 Variance ﬂow as a dynamical system
For Fig.
1(a,b), Eq.
(2.1) and its derivative are approximated with 4096 draws from the unit normal distribution.
For Fig.
1(c), the neural network consists of 32 layers, each with 1024 neurons, and a hyperbolic tangent activation function.
Moreover, the function and its derivative shown in Fig.
2 are evaluated the same as that of Fig. 1.
The ﬁxed point is obtained by bisection method.
C.2 Globally Orthogonal (GO) initialization of MLPs
For the visualization of the function behavior at initialization depicted in Fig.
3, we choose 2 16 data points selected uniformly from the square region [ − 1 , 1] ×
[ − 1 , 1] and apply normalization for both dimension.
The width of the MLP network is set to be 256 .
The visualization of the outputs is on a 200 × 200 grid.
The GO initialization scheme was applied using the Adam optimizer with a ﬁxed learning rate of 0 .
The orthogonality loss objective is regularized with λ = 10 − 4 , and 10 ,stochastic gradient descent steps were performed.
The results presented in Fig.4 are obtained using an MLP network with 16 layers, each withneurons activated with the hyperbolic tangent function.
The GO initialization is trained on 2048 of uniformly distributed points with zero mean and unit variance.
The regularizer λ = 10 − 3 and the learning rate is ﬁxed at 0 .
The GO initialization is trained for 10 , 000 iterations.
The NTK is calculated on 256 evenly spaced points with zero mean and unit variance.
[Figure]
C.3 Synthetic benchmark
The results presented in Fig.5 correspond to a regression task using target function with the following form
[Equation]
where M = 18 is the maximum frequency.
Fig.
8 shows the target function that is used for data generation.
Here we have considered an MLP network with 16 layers, each with 256 neurons activated by the hyperbolic tangent.
For the GO initialization, 2 16 data points are generated uniformly randomly in the square [ − 1 , 1] ×
[ − 1 , 1] .
They are then normalized in each dimension.
The GO initialization is trained with λ = 10 − 4 and a learning rate of 0 .
For each iteration, 2 , 048 data points are selected uniformly from the pre-drawn points.
The initialization is trained for 20 , 000 iterations.
Once the model is initialized with proceed with solving the regression task, and 2048 data points are chosen per batch uniformly and randomly from the data-set with given size.
The MLP is then trained using the Adam optimizer with a learning rate of 5 e − 4 and a decay rate of 0.99 everyiterations, for a total of ∼ 100 , 000 iterations.
We have performed a total of 20 independent trials with parameters initialized from both schemes (i.e. Glorot and GO), and use the ensemble to estimate the 80% probability interval.
The testing error is computed on 200 × 200 regular grid, normalized by the training data normalization.
C.4 Image regression benchmark
The benchmark presented in Fig.6 involves a real RGB image with a resolution of 960 × 960 × 3 .
We use 25% of the pixels as the training data-set, and the signal-to-noise ratio (SNR) is evaluated on the entire image including the training data.
The MLP network we employed consists of 8 layers with a width of 512 , and a hyperbolic tangent activation function.
The GO initialization is trained using the same exact settings is in the synthetic example discussed above, using a batch size of 512 .
Once the model is initialized we proceed to solve the regression task, using the MSE loss objective and a total variation ( TV ) regularization on the Frobenius norm of the MLP Jacobian matrix
[Equation]
where X 1 is a batch of normalized indices, Y 1 is a batch of corresponding RGB value vectors, X 2 is a batch of points for evaluating the total variation regularization, λ r is the regularization coefﬁcient, and N 1 , N 2 are the corresponding batch sizes.
In this case X 1 is drawn uniformly randomly from the sampled normalized indices, X 2 is sampled uniformly randomly in the normalized domain,
N 1 = N 2 = 512 , and λ r = 0 .
The number of training iterations is 400 , 000 .
The optimizer learning rate starts from 10 − 4 , and decays with a rate of 0 .
99 per 1 , 000 steps.
[Table] [Table]
with scheme is able to marginally improve the performance of SIREN when data is sufﬁcient.
C.5 Comparison against SIREN
For a comparison against state-of-the-art methods for image regression, we report the results obtained by an MLP with GO initialization versus a scaled SIREN model [9].
The weight matrices of the
SIREN network are initialized in accordance with [9], without using any heuristics on the ﬁrst layer variance offset, and the architecture stays the same, except for the periodic activation function used in the SIREN case.
Table 1 shows the SNRs for different conﬁgurations.
Interestingly, the proposed
GO initialization can achieve comparable image reconstructions to the SIREN network, suggesting that the proposed GO initialization alone can be sufﬁcient in mitigating frequency bias, without any additional modiﬁcations to the network architecture or activation functions.