Weight Conditioning for Smooth Optimization of Neural Networks
Hemanth Saratchandran Thomas X. Wang Simon Lucey




 










H. Saratchandran et al.




Australian Institute for Machine Learning, University of Adelaide Sorbonne Université, CNRS, ISIR 

hemanth.saratchandran@adelaide.edu.au








Weight Conditioning for Smooth Optimization of Neural Networks
    Hemanth Saratchandran1 Thomas X Wang2 Simon Lucey1
    September 9, 2024
==============================================================




§ ABSTRACT 
  In this article, we introduce a novel normalization technique for neural network weight matrices, which we term weight conditioning. This approach aims to narrow the gap between the smallest and largest singular values of the weight matrices, resulting in better-conditioned matrices. The inspiration for this technique partially derives from numerical linear algebra, where well-conditioned matrices are known to facilitate stronger convergence results for iterative solvers. We provide a theoretical foundation demonstrating that our normalization technique smoothens the loss landscape, thereby enhancing convergence of stochastic gradient descent algorithms. Empirically, we validate our normalization across various neural network architectures, including Convolutional Neural Networks (CNNs), Vision Transformers (ViT), Neural Radiance Fields (NeRF), and 3D shape modeling. Our findings indicate that our normalization method is not only competitive but also outperforms existing weight normalization techniques from the literature.
  





§ INTRODUCTION
 

Normalization techniques, including batch normalization <cit.>, weight standardization <cit.>, and weight normalization <cit.>, have become fundamental to the advancement of deep learning, playing a critical role in the development and performance optimization of deep learning models for vision applications <cit.>. By ensuring consistent scales of inputs and internal activations, these methods not only stabilize and accelerate the convergence process but also mitigate issues such as vanishing or exploding gradients. 
In this paper, we put forth a normalization method termed weight conditioning, designed to help in the optimization of neural network architectures, including both feedforward and convolutional layers, through strategic manipulation of their weight matrices. 
By multiplying a weight matrix of a neural architecture by a predetermined matrix conditioner, weight conditioning aims to minimize the condition number of these weight matrices, effectively narrowing the disparity between their smallest and largest singular values. Our findings demonstrate that such conditioning not only influences the Hessian matrix of the associated loss function when training such networks, leading to a lower condition number but also significantly enhances the efficiency of iterative optimization methods like gradient descent by fostering a smoother loss landscape.  
Our theoretical analysis elucidates the impact of weight matrices on the Hessian's condition number, revealing our central insight: optimizing the condition number of weight matrices directly facilitates Hessian conditioning, thereby expediting convergence in gradient-based optimization by allowing a larger learning rate to be used. Furthermore, weight conditioning can be used as a drop-in component along with the above normalization techniques, yielding better accuracy and training on a variety of problems.

We draw the reader's focus to fig. <ref>, which delineates the pivotal role of weight conditioning in enhancing batch normalization's effectiveness. Illustrated on the left, the figure compares training outcomes for a GoogleNet model on the CIFAR100 dataset under three distinct conditions: Batch Normalization (BN), Batch Normalization with Weight Standardization (BN + WS), and Batch Normalization with Weight Conditioning (BN + WC). Each variant was subjected to a training regimen of 40 epochs using Stochastic Gradient Descent (SGD) at an elevated learning rate of 2. Remarkably, the BN + WC configuration demonstrates a superior accuracy enhancement of nearly 15% over its counterparts, showcasing its robustness and superior adaptability to higher learning rates than typically documented in the literature. On the figure's right, we extend this comparative analysis to include ResNet18 and ResNet50 models, trained on CIFAR100, via SGD with a conventional learning rate of 1e-3 and over a span of 200 epochs. Consistently, BN + WC exhibits a pronounced performance advantage over the alternative normalization strategies, reinforcing its efficacy across diverse neural network frameworks.





To further demonstrate the versatility of weight conditioning, we rigorously tested its efficacy across a spectrum of machine learning architectures, including Convolutional Neural Networks (CNNs), Vision Transformers (ViTs), Neural Radiance Fields (NeRF), and 3D shape modeling. In each scenario, we juxtaposed weight conditioning against established normalization methods cited in current research, unveiling its potential to significantly boost the performance of these advanced deep learning frameworks. 

Our main contributions are:

    
  * We introduce a novel normalization strategy, termed weight conditioning, which strategically modifies the weight matrices within neural networks, facilitating faster convergence for gradient based optimizers. 
    
  * Through rigorous theoretical analysis, we validate the underlying principles of weight conditioning, offering a solid foundation for its implementation and understanding.
    
  * We present comprehensive empirical evidence showcasing the effectiveness of weight conditioning across a variety of machine learning models, highlighting its broad applicability and impact on model performance optimization.





§ RELATED WORK




  §.§.§ Normalization in deep learning:
 Normalization techniques have become pivotal in enhancing the training stability and performance of deep learning models. Batch normalization, introduced by Ioffe and Szegedy <cit.>, normalizes the inputs across the batch to reduce internal covariate shift, significantly improving the training speed and stability of neural networks. Layer normalization, proposed by Ba et al. <cit.>, extends this idea by normalizing inputs across features for each sample, proving particularly effective in recurrent neural network architectures <cit.> and transformer networks <cit.>. Weight normalization, by Salimans and Kingma <cit.>, decouples the magnitude of the weights from their direction, facilitating a smoother optimization landscape. Lastly, weight standardization, introduced by Qiao et al. <cit.>, standardizes the weights in convolutional layers, further aiding in the optimization process, especially when combined with batch normalization. Together, these techniques address various challenges in training deep learning models, underscoring the continuous evolution of strategies to improve model convergence and performance.









§ NOTATION

 Our main theoretical contributions will be in the context of feedfoward layers. Therefore, we fix notation for this here.
Let F denote a depth L neural network with layer widths 
{n_1,…,n_L}. We let X ∈^N× n_0 denote the training data, with n_0 being the dimension of the input. The output at layer k will be denoted by F_k and is defined by

    F_k =  
        
            F_L-1W_L + b_L,     k = L 
    ϕ(F_k-1W_k + b_k),     k ∈ [L-1] 
    
            X,     k = 0

where the weights W_k ∈^n_k-1× n_k and the biases 
b_k ∈^n_k and ϕ is an activation applied component wise.
The notation
[m] is defined by [m] = {1,…,m}.
We will also fix a loss function ℒ for minimizing the weights of F. In the experiments this will always be the MSE loss or the Binary Cross Entropy (BCE) loss. Note that ℒ depends on F.





§ MOTIVATION

In this section, we give some brief motivation for the theoretical framework we develop in the next section.


  
A simple model:Consider the quadratic objective function given by

    ℒ(θ) := 1/2θ^TAθ - b^Tθ

where A is a symmetric n × n matrix of full rank and b is an n × 1 vector. The objective function ℒ is used when solving the eqn. Ax = b. One can see that the solution of this equation if given by x = A^-1b which is precisely the minimum θ^* of ℒ. Thus minimizing the objective function ℒ with a gradient descent algorithm is one way to find a solution to the matrix equation Ax = b. We want to consider the gradient descent algorithm on ℒ and understand how the convergence of such an algorithm depends on characteristics of the matrix A.

Observe that 

    ∇ℒ(θ) = Aθ - b  and  
        H(ℒ)(θ) = A

where H(ℒ)(θ) denotes the Hessian of ℒ at θ. If we consider the gradient descent update for this objective function with a learning rate of η, eqn. (<ref>) implies

    θ^t+1 = θ^t - η∇ℒ(θ^t) = 
        θ^t - η (Aθ^t - b)

Taking the singular value decomposition (SVD) of A we can write 

    A = Udiag(σ_1,⋯ ,σ_n)V^T

where U and V are unitary matrices and σ_1 ≥⋯≥σ_n are the singular values of A. The importance of the SVD comes from the fact that we can view the gradient descent update (<ref>) in terms of the basis defined by V^T. Namely, we can perform a change of coordinates and define

    x^t = V^T(θ^t - θ^*)

The gradient update for the ith-coordinate of x^t, denoted x_i^t becomes 

    x^(t+1)_i = x^t_i - ησ_ix_i^t = (1- ησ_i)x_i^t
        = (1-ησ_i)^t+1x_0.

If we write V = [v_1,⋯ ,v_n]  with each v_i ∈^n× 1 we then have

    θ^t - θ^* = Vx^t = ∑_i=1^nx_i^0(1-ησ_i)^t+1
        v_i.

Eqn. (<ref>) shows that the rate at which gradient descent moves depends on the quantities 
1-ησ_i. This implies that in the direction v_i, gradient descent moves at a rate of (1-ησ_i)^t from which it follows that the closer 1-ησ_i is to zero the faster the convergence. In particular, provided η is small enough, the directions corresponding to the larger singular values will converge fastest.
Furthermore, eqn. (<ref>) gives a condition on how the singular values of A affect the choice of learning rate η. We see that in order to guarantee convergence we need 

    | 1 - ησ_i| < 1  for all  1 ≤ i ≤ n

which implies

    0 < ησ_i < 2  for all  1 ≤ i ≤ n.

This means we must have η < 2/σ_i for each i and this will be satisfied if η < 2/σ_1 since σ_1 is the largest singular value. Furthermore, we see that the progress of gradient descent in the ith direction v_i is bounded by

    ησ_i < 2σ_i/σ_1 < 
        2σ_1/σ_n.

Since σ_1 ≥σ_n we thus see that gradient descent will converge faster provided 
the quantity σ_1/σ_n is as small as possible.
This motivates the following definition.

    Let A be a n × m matrix of full rank. The condition number of A is defined by 
    
    κ(A) := σ_1(A)/σ_k(A)

   where σ_1(A) ≥⋯≥σ_k(A) > 0 and 
   k = min{m, n}.

Note that because we are assuming A to be full rank, the condition number is well defined as all the singular values are positive.






  
Preconditioning: Preconditioning involves the application of a matrix, known as the preconditioner P, to another matrix A, resulting in the product PA, with the aim of achieving κ(PA) ≤κ(A) <cit.>. This process, typically referred to as left preconditioning due to the multiplication of A from the left by P, is an effective method to reduce the condition number of A. Besides left preconditioning, there is also right preconditioning, which considers the product AP, and double preconditioning, which employs two preconditioners, P_1 and P_2, to form P_1AP_2. Diagonal matrices are frequently chosen as preconditioners because their application involves scaling the rows or columns of A, thus minimally adding to the computational cost of the problem. Examples of preconditioners are:

    
  1. Jacobi Preconditioner: Given a square matrix A the Jacobi preconditioner D consists of the inverse of the diagonal elements of A, A → diag(A)^-1A <cit.>.

    
  2. Row Equilibration: Given a n × m matrix A, row equilibration is a diagonal n× n matrix with the inverse of the 2-norm of each row of A on the diagonal, 
    A → (|| A_i:||_2)^-1A, where A_i: denotes the ith-row of A <cit.>.

    
  3. Column Equilibration: Given a n × m matrix A, column equilibration is a diagonal m× m matrix with the inverse of the 2-norm of each column of A on the diagonal, 
    A → A(|| A_:i||_2)^-1, where A_:i denotes the ith-column of A.
    
  4. Row-Column Equilibration: This is a double sided equilibration given by row equilibration on the left and column equilibration on the right. 

The interested reader can consult <cit.> for more on preconditioner. In Sec. <ref> we will explain how row equilibration helps reduce the condition number. 

If we precondition A yielding PA and consider the new objective function 

    ℒ_P(θ) = θ^TPAθ - b^TP^Tθ


Then provided we have that κ(PA) ≤κ(A) the above discussion shows that gradient descent on the new objective function
ℒ_P will converge faster. For this problem the preconditioning can also be thought of in terms of the matrix eqn. Ax = b. We seek to multiply the system by P yielding the new system 
PAx = Pb and provided κ(PA) ≤κ(A), this new system will be easier to solve using a gradient descent.




  
General objective functions: 
The above discussion focused on a very simple objective function given by eqn. (<ref>). In general, objective functions are rarely this simple. However, given a general objective function ℒ, about a point θ_0 we can approximate ℒ by a second order Taylor series yielding 

    ℒ(θ) ≈1/2(θ - θ_0)^TH(θ_0)(θ - θ_0) + 
        (∇ℒ(θ_0))(θ - θ_0) + 
        ℒ(θ_0)

where H, the Hessian matrix of ℒ, is a symmetric square matrix that captures the curvature of ℒ. This expansion offers insights into the local behavior of gradient descent around θ_0, particularly illustrating that if H is full rank, the convergence speed of the descent algorithm is influenced by the condition number of H. However, for many neural network architectures, which typically have a vast number of parameters, directly computing H is impractical due to its 𝒪(n^2) computational complexity, where n represents the parameter count. To address this challenge, the subsequent section will introduce the concept of weight conditioning, a technique aimed at reducing the condition number of H without the necessity of direct computation.






§ THEORETICAL METHODOLOGY


In this section, we explain our approach to conditioning the weight matrices within neural networks, aiming to enhance the convergence efficiency of gradient descent algorithms. Fixing a loss function ℒ we saw in the previous sec. <ref> that we can approximate ℒ locally by a second order quadratic approximation via the Taylor series about any point θ_0 as

    ℒ(θ) ≈1/2(θ - θ_0)^TH(θ_0)(θ - θ_0) + 
        (∇ℒ(θ_0))(θ - θ_0) + 
        ℒ(θ_0)

where H is the Hessian matrix associated to ℒ. Note that H is a square symmetric matrix. The discussion in sec. <ref> made two key observations:

      
  * The convergence rate of gradient descent on ℒ is significantly influenced by the condition number of H, provided H is full rank.
    
  * Inspired by the first point, the convergence rate can be accelerated by diminishing the condition number of H.

It was mentioned at the end of sec. <ref> that in general we cannot expect to have direct access to the Hessian H as this would require a computational cost of 𝒪(n^2), where n is the number of parameters, which can be extremely large for many deep learning models. In this section we introduce weight conditioning, which is a method of conditioning the weights of a neural networks that leads to a method of bringing down the condition number of H without directly accessing it.





 §.§ Weight conditioning for feed forward networks




We fix a neural network F(x;θ) with L layers, as defined in sec. <ref>. Given a collection of preconditioner matrices P = {P_1,⋯, P_l} where 
P_k ∈^n_k-1× n_k we define a preconditioned network 
F^pre(x;θ) as follows. The layer maps of F^pre(x;θ) will be defined by:

    F^pre_k =
        ϕ_k((P                                     _kW_k)^T(F^pre_k-1)(x) + b_k),     k = [1, L] 
    
        x,     k = 0.

We thus see that the layer-wise weights of the network F^pre are the weights W_k of the network F preconditioned by the preconditioner matrix P_k. Fig. <ref> provides a visual depiction of the preconditioned network F^pre(x;θ), illustrating how the preconditioner matrices are applied to each layer's weights to form the updated network configuration. 


    Given a feed forward neural network F(x;θ) we call the process of going from F(x;θ) to F^pre(x;θ) weight conditioning. 

It's important to observe that both networks, F and F^pre, maintain an identical count of parameters, activations, and layers. The sole distinction between them lies in the configuration of their weight matrices.

Our objective is to rigorously demonstrate that an optimal choice for weight conditioning a network is to use row equilibration.
Recall from sec. <ref>, that row equilibration preconditioner for a weight matrix W_k ∈ℝ^n_k-1× n_k is defined as a diagonal matrix E_k ∈ℝ^n_k-1× n_k-1. Each diagonal element of E_k is determined by the inverse of the 2-norm of the i-th row vector of W_k, i.e. 
|| (W_k)_i:||_2^-1. 



All our statements in this section hold for the column equilibrated preconditioner and the row-column equilibrated preconditioner. Therefore, from here on in we will drop the usage of the word row and simply call our preconditioner an equilibrated preconditioner.





    
    
  




There are two main reasons we are choosing to use equilibration as the preferred form to condition the weights of the neural network F. The first reason is that it is a diagonal preconditioner, therefore requiring low computational cost to compute. Though more importantly, as the following theorem of Van Der Sluis <cit.> shows, it is the optimum preconditioner amongst all diagonal preconditioners to reduce the condition number.



    Let A be a n× m matrix, P an arbitrary diagonal n × n matrix and E the row equilibrated matrix built from A. Then κ(EA) ≤κ(PA).


For the fixed L layer neural network F, let F^eq denote the equilibrated network with weights E_kW_k, where E_k is the equilibrated preconditioner corresponding to the weight matrix W_k. We have the following proposition.


κ(E_kW_k) ≤κ(W_k) for any 1 ≤ k ≤ L. In other words the weight matrices of the network F^eq have at least better condition number than those of the network F


    The matrix W_k can be written as the product 
    I_n_k-1W_k where 
    I_n_k-1∈^n_k-1× n_k-1 is the n_k-1× n_k-1 identity matrix. Applying thm. <ref> we have 
    
    κ(E_kW_K) ≤κ(I_n_k-1W_K) = 
        κ(W_k).









For the following theorem we fix a loss function ℒ see Sec. <ref>. Given two neural networks F and G we obtain two loss functions 
ℒ_F and ℒ_G, each depending on the weights of the respective networks. We will denote the Hessian of these two loss functions at a point θ by H_F(θ) and H_G(θ) respectively. 




Let F(x;θ) be a fixed L layer feed forward neural network. 
Let F^eq(x;θ) denote the equilibrated network obtained by equilibrating the weight matrices of F. Then 

    κ(H_F^eq(θ)) ≤κ(H_F(θ))

for all parameters θ at which both H_F^eq and 
H_F have full rank.


The proof of Thm. <ref> is given in Supp. material Sec. 1.

Thm. <ref> shows that weight conditioning by
equilibrating the weights of the network F thereby forming F^eq leads to a better conditioned Hessian of the loss landscape. From the discussion in sec. <ref> we see that this implies that locally around points where the Hessian has full rank, a gradient descent algorithm will move faster. 

Weight conditioning can also be applied to a convolutional layer. Please see Supp. material Sec. 1 for a detailed analysis on how this is done.















§ EXPERIMENTS




 §.§ Convolutional Neural Networks (CNNs)



CNNs are pivotal for vision-related tasks, thriving in image classification challenges due to their specialized architecture that efficiently learns spatial feature hierarchies. Among the notable CNN architectures in the literature, we will specifically explore two significant ones for our experiment: the Inception <cit.> architecture and DenseNet <cit.>, both of which continue to be highly relevant and influential in the field.





  §.§.§ Experimental setup:
 
We will assess four normalization strategies on two CNN architectures. Our study compares BN, BN with weight standardization (BN + WS), BN with weight normalization (BN + W), and BN with equilibrated weight conditioning (BN + E). Training employs SGD with a 1e-3 learning rate on CIFAR10 and CIFAR100, using a batch size of 128 across 200 epochs. For more details on the training regiment, see Supp. material Sec. 2. For results on ImageNet1k Supp. material Sec. 2.






  §.§.§ Inception:
 
In our experiment, we employed a modified Inception architecture, as proposed in <cit.>, featuring eight inception blocks. Weight conditioning was effectively applied to just the initial convolutional layer and the final linear layer. For detailed insights into the architecture and normalization applications, see Sec. 2 of the Supp. material.

Results on the CIFAR10 dataset, depicted in Fig. <ref>, highlight that the Inception model with BN + E exhibits the quickest loss reduction and accelerated convergence in Top-1% accuracy among all four normalization strategies. Similarly, Fig. <ref> illustrates that, on the CIFAR100 dataset, BN + E achieves the lowest training loss and highest Top-1% accuracy, outperforming other methods in convergence speed. Across both datasets, as summarized in Tab. <ref>, BN + E consistently delivers superior Top-1% and Top-5% accuracies, affirming its effectiveness.






















  §.§.§ DenseNet:
 For this experiment we tested four normalizations on the DenseNet architecture <cit.>. We found that it sufficed to apply equilibrated weight conditioning to the first convolutional layer and the last feedforward layer as in the case for Inception. An in-depth description of the architecture employed in our study, and how we applied each normalization is given in Sec. 2 of the Supp. material.

Fig. <ref> and Fig. <ref> display CIFAR10 and CIFAR100 dataset results, respectively. For CIFAR10, BN + E demonstrates quicker convergence and higher Top-1% accuracy as shown in the train loss curves and accuracy figures. CIFAR100 results indicate BN + E outperforms other normalizations in both convergence speed and Top-1% accuracy. Tab. <ref> summarizes the final Top-1% and Top-5% accuracies for all normalizations, with BN + E leading in performance across both datasets.
















 §.§ Vision Transformers (ViTs)


Vision Transformers (ViTs) <cit.> have emerged as innovative architectures in computer vision, showcasing exceptional capabilities across a broad spectrum of tasks. In general, we found that most work in the literature applied layer normalization to vision transformers which was much more robust for training than batch normalization. We found that if we removed layer normalization training was impeded significantly leading to gradients not being able to be backpropgated. Therefore, for this experiment we will consider 3 normalization scenarios, layer normalization (LN), layer normalization with weight normalization (LN + W) and layer normalization with equilibrated weight conditioning (LN + E).





































  §.§.§ ViT-Base (ViT-B):
 The ViT-B architecture <cit.>, with its 86 million parameters, exemplifies a highly overparameterized neural network. We investigated three variations of ViT-B, each modified with a different normalization: LN, LN + W, and LN + E; the implementation details are provided in Sec. 2 of the Supp. material. These models were trained on the ImageNet1k dataset using a batch size of 1024 and optimized with AdamW.


Tab. <ref> shows the final Top-1% and Top-5% accuracies for the ViT-B architecture with the above three different normalizations. The table shows that LN + E outperformed the other normalization schemes.












  §.§.§ ViT-Small (ViT-S):
 For experiments on smaller transformer networks on the CIFAR100 dataset, please see Sup. material sec. 2.





    
    











 §.§ Neural Radiance Fields (NeRF)






Neural Radiance Fields (NeRF) <cit.> have emerged as a pioneering approach in 3D modeling, using Multi-Layer Perceptrons (MLPs) to reconstruct 3D objects and scenes from multi-view 2D images. We utilized the standard NeRF model from <cit.>, noting that unlike transformers or CNNs, NeRF architectures are relatively shallow, typically comprising 8 hidden layers, where common normalization techniques can hinder training. We explored the performance of a standard NeRF against an equilibrated weight conditioned NeRF (E-NeRF). For an in-depth look at the NeRF setup and our application of weight conditioning, refer to Sec. 2 in the Supp. material. Both models were trained on the LLFF dataset <cit.>.

Fig. <ref> presents outcomes for the Fern and Room instances from LLFF, demonstrating that E-NeRF outperforms the vanilla NeRF by an average of 0.5-1 dB. Tab. <ref> gives the test PSNR averaged over three unseen scenes over the whole LLFF dataset. E-NeRF on average performs better over the whole dataset.

















 §.§ Further Experiments


Further experiments can be found in Supp. material Sec. 3:
Applications to 3D shape modelling, cost analysis and ablations.
   





§ LIMITATIONS


Weight conditioning entails the application of a preconditioner to the weight matrices of a neural network during the forward pass, which does extend the training duration per iteration of a gradient optimizer compared to other normalization methods. We leave it as future work to develop methods to bring down this cost.









§ CONCLUSION

In this work, we introduced weight conditioning to improve neural network training by conditioning weight matrices, thereby enhancing optimizer convergence. We developed a theoretical framework showing that weight conditioning reduces the Hessian's condition number, improving loss function optimization. Through empirical evaluations on diverse deep learning models, we validated our approach, confirming that equilibrated weight conditioning consistently aligns with our theoretical insights.




















§ ACKNOWLEDGEMENTS


Thomas X Wang acknowledges financial support provided by DL4CLIM (ANR-
19-CHIA-0018-01), DEEPNUM (ANR-21-CE23-0017-02), PHLUSIM (ANR-23-CE23-
0025-02), and PEPR Sharp (ANR-23-PEIA-0008”, “ANR”, “FRANCE 2030”). 

Part of this work was performed using HPC resources from GENCI–IDRIS (Grant 2023-AD011013332R1).






splncs04

