To update or not to update? Neurons at equilibrium in deep models
Abstract
Recent advances in deep learning optimization showed that, with some a-posteriori information on fully-trained models, it is possible to match the same performance by simply training a subset of their parameters.
Such a discovery has a broad impact from theory to applications, driving the research towards methods to identify the minimum subset of parameters to train without look-ahead information exploitation.
However, the methods proposed do not match the state-of-the-art performance and rely on unstructured sparsely connected models.
In this work we shift our focus from the single parameters to the behavior of the whole neuron, exploiting the concept of neuronal equilibrium (NEq).
When a neuron is in a configuration at equilibrium (meaning that it has learned a specific input-output relationship), we can halt its update; on the contrary, when a neuron is at non-equilibrium, we let its state evolve towards an equilibrium state, updating its parameters.
The proposed approach has been tested on different state-of-theart learning strategies and tasks, validating NEq and observing that the neuronal equilibrium depends on the specific learning setup.
Introduction
In recent years, deep learning has become a staple solution to different tasks, such as computer vision, bio-informatics, speech recognition, and many more.
Unfortunately, modern neural network architectures are becoming more and more “expensive”: performing simple inferences requires a lot of computational resources and training the models even more.
Such cost poses several challenges for the research community: the training of a network model is associated with large carbon footprints and the commercialization of AI research (especially for edge devices) is hindered by the resource requirements of the models Strubell et al.
For several years now many works in literature have shown that is possible to shrink both the size and resource requirements, mainly via quantization Yang et al.
(2020); Jin et al.
and pruning Wang et al.
; Tang et al.
; Tanaka et al.
In particular, pruning techniques can drastically reduce the number of operations needed to perform inference, without affecting the overall performance of the model.
Pruning targets the reduction of parameters in the model, which has advantages after training; unfortunately, it is unable to reduce the cost of the training cycle, and on the contrary, it requires, in general, iterations of few-shot pruning, followed by fine-tuning Lee et al.
; Tartaglione et al.
A recent work shows the existence of sub-graphs in the whole deep learning model which, when trained in isolation, match the performance of the whole model Frankle and Carbin (2019).
This opens the road towards a whole field of research, where many approaches are proposed to find these parameters a-priori: this shows a potential environmental impact, considering that a smaller model is trained, leading to a lower energy consumption Strubell et al.
Indeed, when training a neural network, the back-propagation procedure and the weights update lead to the larger part of FLOPs (compared to forward propagation) Plaut et al.
Baydin et al.
Researchers started to experiment pruning early in the training or even before the training starts in the hope of training a restricted number of parameters.
Unfortunately, training a sparse network with standard optimizers leads to subpar results Evci et al.
or the final result does not differ much from magnitude pruning at the end of the training Frankle et al.
The causes for such behaviors are still a matter of debate among the community 2019; 2021.
In this work, we shift the focus from the single parameter to the whole neuron, and we propose
NEq, an approach to evaluate whether a given neuron is at equilibrium for the learning dynamics.
If the neuron is in such a state, its parameters have already reached a target configuration and do not require a further update.
Unlike many other recent approaches, NEq disables entire neurons (hence, in a structured way), does not require prior knowledge of the specific task (for example by first training a model to convergence), and automatically self-adapts to the specific learning policy deployed.
Unlike pruning techniques, NEq does not remove the neurons’ contribution to the output; instead, it only prevents unnecessary updates to their weights: as a result, we reduce the number of operations performed by the back-propagation algorithm and the optimizer.
The rest of the paper is organized as follows.
In Sec.
2 we discuss the related literature; Sec.presents the concept of neuronal equilibrium and how to evaluate it; Sec. 4 presents the experimental validation inclusive of an ablation study and Sec. 5 draws the conclusions.
Related works
It is broadly acknowledged that the typically-deployed deep learning models on the state-of-the-art scenarios are over-parametrized Mhaskar and Poggio (2016); Brutzkus et al.
This ignites two lines of research: reducing the size of these models (with pruning algorithms) or saving computational resources at training time.
While the first has been broadly explored, the latter suffered a stalemate until a recent work suggested its feasibility Frankle and Carbin (2019).
Pruning strategies.
Attempts to reduce the number of parameters from learned models date back to 1989 when Mozer and Smolensky proposed skeletonization , a technique to identify less relevant neurons in a trained model and to remove them 1989.
This was accomplished thanks to the evaluation of an error function penalty from a pre-trained model.
In the same years, LeCun et al.
also proposed a work where the information from the second-order derivative of the error function is leveraged to rank the parameters of the trained model on a saliency-like basis LeCun et al.
In the last decade, thanks to the broad availability of computational resources, pruning approaches gainedback popularity, with approaches including the exploitation of dropout Molchanov et al.
Gomez et al.
(2019), sensitivity-based approaches Lee et al.
; Tartaglione et al.
(2021, 2022), relaxation of ℓ 0 regularization Louizos et al.
Sanh et al.
(2020) and optimization of auxiliary parameters Xiao et al.
Despite leading to very compact deep models, the demanded training complexity is frequently very high, as many of the most well-known approaches rely on iterative finetuning (or few-shot pruning) approaches.
Hence, the training complexity for these techniques is larger than training a vanilla model (which in many cases is used as initialization).
Much computational complexity could be saved if structured pruning is applied before training the model itself (namely, entire neurons/channels are pruned).
Some works on pruning have suggested the possibility of re-allocating previously pruned parameters Tresp et al.
or filters He et al.
(2018) to learn new functions and minimize the loss, towards higher generalization capability Blalock et al.
The lottery ticket hypothesis.
In a recent paper, Frankle and Carbin provided empirical evidence that, at initialization, there exists a sub-graph of the original deep learning model such that, when trained in isolation, it can match the performance of the complete model Frankle and Carbin (2019).
Indeed, the authors claim that the parameters in the sub-graph have “won at the lottery of initialization”.
From a very practical perspective, this means that all the parameters not in the selected sub-graph are, de facto, pruned from the model, not requiring any gradient computation.
This empirical evidence potentially opens the road to the development of algorithms for pruning at initialization. Unfortunately, at such point, two obstacles are yet to be tackled. From one side, Frankle and Carbin (2019) uses an iterative procedure to identify the sub-graph: it simply proves its existence but does not provide a method to [Figure]
a t e d
[Figure]
nd no r m a li
ze
d c on [Figure] t e n a t
e
d a nd [Figure]
li ze
[Figure] find it directly at initialization.
On another side, it focuses on un-structured sub-graphs, meaning that parameters are treated in isolation: is it possible to find structured sub-graphs (or in other words, removing entire nodes and not just arcs)?
Beyond the lottery ticket hypothesis.
Many works of the last two years have received inspiration from the lottery ticket hypothesis, in the quest of determining the early lottery tickets.
This is, however, a tough task to solve: Frankle et al.
(2020) shows that there is a region, at the very early stages of learning, where the lottery tickets identified with iterative pruning are “not stable” (meaning that tickets extracted at different moments of this early stages are essentially different).
This suggests that, in the very first epochs, the neural network evolves in very different states, making the problem of a-priori identifying winning tickets hard Tartaglione (2022).
This is also endorsed by other works, including Morcos et al. ; Malach et al.
On the other hand, other approaches reduce the overall complexity of the iterative training by drawing early-bird tickets You et al.
(meaning that they learn the lottery tickets when the model has not yet reached full convergence), even reducing the training data Zhang et al.
or moving the first steps towards structurally-sparse winning tickets, yet still at iterative fashion, applying similar concept as Frankle and Carbin to entire neurons and channels Chen et al.
With NEq we learn the important lessons related to the lottery ticket hypothesis, targeting the reduction of computational complexity at training time, without exploiting knowledge on pre-training models, or rewinding.
Hence, we do not target the achievement of sparse architectures, but we aim at determining when a whole neuron requires to be updated or when the computation of the gradient for its parameters is not necessary.
As such, the comparison with the other presented approaches will be in general unfair, as they require a much greater computational complexity for training because of un-structured sparsity (which introduces an overhead in the representation of the tensors) and the iterative strategies.
Furthermore, we will observe the possibility, along the training process, that some neurons, already kept in a “frozen” state, might unfreeze, requiring additional update steps.
Although resource re-allocation has been exploited before Tresp et al.
; He et al.
, our unfreezing is different as it involves learning of a specific target function by the neuron, and not learning new ones through their re-allocation.
In the next section, the notion of neuronal equilibrium will be presented, as well as the strategy to determine which neurons will require gradient computation.
Neurons at equilibrium
In this section, we will treat the problem of determining when a given neuron, along with the learning dynamics, finds itself at equilibrium.
Towards this end, we define y t i,ξ as the output of the i -th neuron when the input ξ is fed to the whole model trained after t epochs.
Given a set of inputs ξ ∈ Ξ val (where Ξ val is the validation set), it is possible to compare each n -th element y ti,n,ξ with y t − 1 i,n,ξ , for the same model’s input: what changes are the parameters of the model.
Fig.
1 provides an overview of the nomenclature used: in the rest of the section we will see how to determine when a neuron is at equilibrium.
3.1 Neuronal equilibrium
In this section, we are interested in evaluating when the relationship between the input of the model and the output of the i -th neuron is modified.
When this happens, the neuron is at non-equilibrium , meaning that its learned function, in the whole picture (or in other words, taking into account the evolution of the neurons in the previous layers as well), is still “evolving”.
We are interested in identifying the scenarios where the neuron is at equilibrium at the net of the interactions with the other neurons.
To assess it, let us define the cosine similarity between all the outputs of the i -th neuron at time t and at time t − 1 for the whole validation set Ξ val as
[Equation]
Here we can determine that, when ϕ i = 1 , the i -th neuron produces the same (eventually scaled) output between the evaluation at time t and at time t − 1 for the same input ξ of the model.
We say the i -th neuron reaches equilibrium when we have
[Equation]
where k ∈ [ − 1; +1] is some constant value.
We can have the following scenarios:
• k = 1 .
In this case, the two outputs are perfectly correlated, meaning that the relationship bounding the input of the whole model ξ and the output of the specific i -th neuron is
maintained.
• k ∈ (0; 1) .
The outputs correlate, but we are in presence of an oscillatory behavior (in the sense that the cosine similarity varies by a constant value between consecutive evaluations). This effect can be caused by stochastic effects like high learning rate/regularization, small batch size, or a combination of them.
• k ∈ [ − 1; 0] .
Also in this case we are in the presence of oscillatory behavior, but the outputs are anti-correlated or de-correlated.
Here follows the evaluation framework to determine the arrival to an equilibrium state for the i -th neuron.
3.2 Neuron dynamics evaluation
To assess the convergence to equilibrium for (2), it is important to evaluate the variation of the similarities ϕ ti over time.
Towards this end, let us introduce the variation of similarities
[Equation]
According to the analysis in Sec. 3.1, in this case, we say we reach equilibrium when ∆ ϕ ti → 0 .
Hence, it is useful to keep track of the recent evolution over the similarity scores in the model: towards this end, we can introduce the velocity of the similarity variations:
[Equation]
where µ eq is the momentum coefficient.
We can rewrite (4) making the similarity scores explicit, obtaining
[Equation]
where ( · ) m indicates power of m .
If we assume ϕ t ∈ [0; 1] ∀ t (which is the case of ReLU-activated neurons), in order to prevent (5) from exploding, we need to set µ eq ∈ [0; 0 . 5] .
We can extend this setup to any layer if we assume that the neurons in the trained model will not reach equilibrium with anti-correlated outputs.
3.3 Selection of trainable neurons at non-equilibrium
To evaluate when a given neuron has reached equilibrium, exploiting (2), we can say that the i -th neuron is at equilibrium when it can satisfy
[Equation]
It is important to notice that, once (6) is satisfied for a certain t , in case something changes in the learning dynamics (for example, the learning rate is re-scaled), there may exist some t ′ >
t such that the constraint is not satisfied anymore.
When this happens, it means that the neuron is driven towards new states , and is no anymore at equilibrium.
Hence, it requires to be updated again.
3.4 Overall training scheme
[Figure]
END
End of training?
[Figure]
The overall training scheme is summarized in Fig.
The model is trained for one epoch, after which neurons at equilibrium are identified.
We split this into two phases: in the first (neuron dynamics evaluation), the velocity of the similarities is evaluated according to (4), while in the second (trainable neurons at non-equilibrium) the hidden neurons at non-equilibrium, which will be trained for the next epoch, are identified according to (6).
For the first epoch, all the neurons are considered at non-equilibrium by default.
The evaluation of neurons at equilibrium is agnostic to the general training strategy, which can include arbitrary re-scaling for the learning rate/hyper-parameters or the most common optimizers.
In the next section, we will test this procedure on very different architectures, tasks, and learning strategies.
Experiments
In this section, we report the experiments supporting the approach as presented in Sec. 3.4.
First, we will perform an ablation study, analyzing single contributions for the introduced hyper-parameters and providing an overview of neuronal equilibrium along the training process; then, we will test the proposed technique on state-of-the-art network architectures, datasets, and learning policies.
All experiments were performed using 8 NVIDIA A40 GPUs and the source code uses PyTorch 1.10.
4.1 Ablation study
We performed our ablation study training a ResNet-32 He et al.
[-U pd a t e d N e u r o n s × 
-]
(2015) model on CIFAR10 Krizhevsky et al..
Unless differently specified, following the hyper-parameters setup of Zagoruyko
[Figure]
[Figure]
[Figure]
and Komodakis (2016), the model is trained using SGD as an optimizer, with a starting learning rate η = 0 . 1 and momentum µ opt = 0 .
9 and a weight decay of 5 × 10 − 4 for 250 epochs.
The learning rate is decayed by a factor of 10 after 100 and 150 epochs, using the formulation as in (4), with ∥ Ξ val ∥ 0 = 50 , µ eq = 0 .
5 , and ε = 0 .
4.1.1 SGD vs Adam
To show that our technique automatically self-adapts to the training policy, we compare the evolution of the FLOPs required for a back-propagation step and the number of updated neurons of two different training of the ResNet-32: one using the SGD optimizer with µ opt = 0 . 9 , and the other using the
Adam optimizer.
For Adam we leave the hyper-parameters to their default values ( η = 0 .
001 , β 1 = 0 .
9 , and β 2 = 0 .
and use the same weight decay as for SGD ( 5 × 10 − 4 ).
Fig.
3 shows the trends for the two training procedures.
We can see that in the first phase of the train, where η is high, the amount of the trained neurons (and the FLOPs required for the backward pass) is higher.
This is related to the general lack of equilibrium in the neurons of the network: at high learning rates, the configuration of the neurons’ parameters is subjected to high stochastic noise.
As the train progresses, and the network progresses toward its final configuration, fewer and fewer neurons need to be updated.
Noticeably, Adam drives the neurons towards equilibrium in a faster way, as expected; however, in simple tasks like the considered one, converges to lower accuracy scores ( 92 .
96% for SGD and 92.01% for Adam).
Furthermore, at the first learning rate decay (epoch 100 ), for SGD the number of updated neurons first decreases and then increases, such a phenomenon is not present in the Adam case.
This is explained by the different working principles of the two optimizers: SGD explores the solution space looking for large minima, searching for configurations that prevent equilibrium in high learning rate regimes. [Figure]
(a) Distribution of ϕ . [Figure] (c) Distribution of v ∆ ϕ with µ eq = 0 . [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
4.1.2 Distribution of ϕ & choice of µ eq
Looking at different values for µ eq in Fig.
4, we observe for all the values a convergence to similar accuracy.
Despite without warranty from the theory, we tested a very large value for the momentum coefficient ( 0 . 9 ): the convergence of v t ∆ ϕ
i shows that the neurons are in general in a very correlated case of equilibrium, with very high values for k in (2), which is also empirically observed in Fig.
5a.
However, including a very large value for µ eq maintains the memory of very old variations, producing a sub-optimal reduction in terms of FLOPs reduction.
We find that a good compromise, supported by the findings as in Sec. 3.2, is to set µ eq to 0 .
Fig.
5 reports the distribution for the velocities for ϕ , ∆ ϕ , and v ∆ ϕ , observing that the average converges to a specific k for each of the three learning rates used.
4.1.3 Impact of the validation set size and ε
Tab. 1b provides an empirical evaluation of the impact on the performance and on the FLOPs varying the validation set size.
We indeed observe not a significant impact on the performance of the model varying it.
Interestingly, the approach produces extremely good results even for extremely low cardinality for the validation set (down to even a single image): this can be explained by the presence of convolutional layers (the only fully-connected layer is the output layer, excluded by default) which even with little images produce high-dimensionality output in every neuron (Fig. 1) and by the homogeneity of the considered dataset.
Investigating the impact of ε , instead, we find for very high values of ε a drop in performance, identifying a good compromise for classification tasks to 0 . [Table]
4.2 Main experiments
In this section [Table] we show the results of the proposed method.
For our experiments, we focused on different state-of-the-art architectures trained on standard classification and semantic segmentation datasets.
All the learning policies used are borrowed from other works and are un-optimized to test the adaptability of NEq.
ResNet-32 trained on CIFAR-10 .
The training spans 250 epochs, using SGD as optimizer with momentum µ opt = 0 .
9 , weight decay 5 × 10 − 4 and initial learning rate η = 0 .
1 , reduced by a factor of 10 after 100 and 150 epochs.
To evaluate the neuronal equilibrium we used a ∥ Ξ val ∥ 0 ofimages, µ eq = 0 . 5 , and ε = 0 .
We used a batch size of 100 images during training.
ResNet-18 trained on ImageNet-1K Krizhevsky et al.
This model was trained with SGD as optimizer for 90 epochs, with η = 0 . 1 , reduced by a factor of 10 every 30 epochs, µ opt = 0 .
9 and weight decay 10 − 4 using a batch size of 128 .
We used a ∥ Ξ val ∥ 0 of 1.2k images, µ eq = 0 .
5 , and ε = 0 .
Swin Transformer Liu et al.
(Swin-B) trained on ImageNet-1K .
To test our technique on more modern models and training policies, we used the Swin-B architecture.
Here we trained the model starting from a pre-trained checkpoint trained on ImageNet-21K, following the official GitHub repository 2 released under the MIT License.
We used a ∥ Ξ val ∥ 0 of 1.2k images, µ eq = 0 .
5 , and ε = 0 .
DeepLabv3 Chen et al.
(2017) trained on COCO Lin et al.
Other than classification tasks of varying complexity, we tested our procedure on a semantic segmentation problem.
For this experiment, we used DeepLabv3 with a ResNet-50 backbone and the COCO dataset.
To train the network, we followed the state-of-the-art procedure defined in PyTorch 3 .
We evaluated the neuronal equilibrium using a ∥ Ξ val ∥ 0 of 320 images, µ eq = 0 .
5 , and ε = 0 .
4.3 Discussion
The results (average over five different runs) are reported in Tab. 2.
For each experiment, we compare our technique with a “stochastic” approach.
Namely, we randomly halt, at every epoch and with probability p , the update of a given neuron.
We test on three different probabilities: 0 . 2 , 0 . 5 , and a probability that is as close as possible to the average over the one achieved by NEq - indicated with “*”.
To evaluate the effectiveness of the proposed procedure we focus on the average computational complexity of the back-propagation for a single update iteration (expressed in FLOPs) and the network generalization capabilities at the end of the training.
In all the considered scenarios, it is possible to observe a reduction of FLOPs with very marginal or no performance drop for NEq.
When compared to the stochastic approach, with fixed probabilities, the amount of saved computation is similar in all the considered scenarios, but the loss in performance varies, depending on the specific architecture/dataset.
On the contrary, NEq remains consistent in performance, self-adapting to the specific setup and saving the largest FLOPs for the given performance.
Furthermore, testing the stochastic approach with the same FLOPs saving (hence, even letting that information leak in favor of the stochastic approach), the performance loss is lower.
Limitations.
The current approach analyzes the behavior of an entire neuron.
However, empirical experiments show that there could be further improvements considering ensembles of neurons.
For example, Fig. 5 shows the average value for the similarities close to a constant but many neurons are still away from the convergence value, meaning that these neurons, at isolation, are still not at equilibrium: is the scenario changing when considering the dynamics of groups of neurons?
Furthermore, to validate the adaptability of NEq to the most popular training schemes, no optimization of the hyper-parameters for the training procedure has been performed (as it is out of scope for our evaluation).
However, higher savings in computational complexity are possible by tuning the training strategy as well.
In such a direction, prospectively, it will be of interest to design more efficient learning strategies which keep into account the concept of neuronal equilibrium.
Conclusions
The work by Frankle and Carbin (2019) showed the existence of sub-graphs in deep models which, when trained in isolation, can match the original performance of the whole model.
Finding these sub-graphs is, however, a complex task, as in the first stages of the learning the model itself is at non-equilibrium.
Identifying these with a dynamic strategy, without requiring a posterior over the whole training process, is a crucial task to be solved, towards computational resources saving.
Differently from the vast majority of the literature which focuses on the identification of sub-graphs without any concrete computational saving (as they rely on iterative or roll-back algorithms), we have introduced the knowledge of neuronal equilibrium, looking for entire structures of the deep model at equilibrium, not requiring further optimization and gradient computation, which self-adapts to very specific experimental setups on very different learning scenarios.
This work opens the doors toward a deeper understanding of the deep neural network’s learning dynamics and to the development of new training strategies exploiting this knowledge