Unit-level surprise in neural networks
Abstract
To adapt to changes in real-world data distributions, neural networks must update their parameters.
We argue that unit-level surprise should be useful for: (i) determining which few parameters should update to adapt quickly; and (ii) learning a modularization such that few modules need be adapted to transfer.
We empirically validate (i) in simple settings and reﬂect on the challenges and opportunities of realizing both (i) and (ii) in more general settings
Introduction
Neural networks have achieved remarkable successes on problems with independent and identicallydistributed (IID) data [23].
However, real-world data is not IID—environmental conditions shift [8], new tasks are encountered [27], and agents alter their behaviour [6].
Ideally, we would like our networks to transfer or adapt quickly to such out-of-distribution (OOD) data [30].
Fast adaptation to OOD or shifted data is often achieved by updating only a few parameters—the ﬁnal layer of an ImageNet-pretrained network [13, 42], the batch-normalization layers [24, 34, 38], or speciﬁc mechanisms or modules [31, 15].
For a given problem or (expected) shift type , a human usually decides: (i) which few parameters should be updated; and (ii) what architecture or modularization will allow the network to be adapted by updating only a few parameters (e.g. Rebufﬁ et al.
[34] use residual adapters for domain adaptation).
However, we often do not know what type of shift will occur and thus how to decide (i) and (ii).
In this work we argue that unit-level surprise can help identify what has changed, and this in turn can be used to: (i) infer which few parameters should be updated to adapt quickly by transferring past knowledge; and (ii) learn an appropriate modularization such that very few modules need be adapted to transfer.
For (i), we propose the use of unit-level surprise— “This looks the same as before, no need for me to update!”
— along with modulatory update-in-progress signals— “This is being dealt with by someone else, hold tight!” .
These two additional pieces of information at the unit-level allow the network to update only the appropriate parameters to facilitate fast adaptation.
This idea is depicted in Figure 1 and further motivated in Appendix A. For (ii), we propose to use the number of surprised units as a proxy score for the number of parameters that need to be adapted to a given shift, which can then be optimized to learn a modularization that changes sparsely and thus can be adapted quickly.
As a ﬁrst step, we empirically validate the usefulness of unit-level surprise by focusing on use case (i) above.
Through experiments on shifted EMNIST [7] datasets, we show that unit-level surprise can be used to create shift-dependent ﬁne-tuning strategies that often align with our intuition about which few parameters should update.
We then critique these results, identifying several challenges and opportunities for using unit-level surprise in more general settings where its full potential could be unlocked.
Despite such challenges, the “beauty” of unit-level surprise—its alignment with our intuitions about how to handle OOD data, its seemingly strong signal for learning how to modularize knowledge, and its relations to neuroscience concepts like metaplasticity [1, 2, 41] and the free-energy principle [12]—convinces us that it should be better for realizing (i) and (ii) in more general settings. Shift [Figure]
[Figure]
Initial validation
In this section we provide an initial empirical validation of the usefulness of unit-level surprise by showing how it can be used to determine which few parameters should be adapted to a given shift.
Experimental setup.
We train a simple 5-layer network with 3 convolutional (conv) and 2 fullyconnected (FC) layers on 37 of the 47 classes in EMNIST [7].
We then adapt this network to new data from 1 of 10 distribution shifts (see Appendix B.1) for which we expect it to be optimal to update different parts of the network.
In particular, we use 7 low-level shifts where we expect only the early layers to need to adapt (e.g. new background) and 3 high-level shifts where we expect only the later layers to need to adapt (e.g. held-out/unseen classes).
Full experimental details are in Appendix B.2, and code is available at https://github.com/cianeastwood/unit-level-surprise .
Calculating unit-level surprise.
For each unit in a network we store the 1D distribution of its activations under the training data, P ( A ) , and compare this with the distribution of activations under a shifted data distribution, Q ( A ) .
More speciﬁcally, we parameterize P ( A ) and Q ( A ) as softly-binned histograms [40] and infer the parameters of these distributions from the training and shifted data respectively.
We then calculate the (Bayesian) surprise of a unit as s ( A )
=
D KL ( Q ( A ) ||
P ( A ))
We discuss this quantity in Appendix B.3 and how to calculate it from bin counts in Appendix B.4.
[Figure]
Surprise patterns.
Together, a network’s surprise values yield shift-dependent patterns.
These patterns, shown in
Fig. 2, help us to understand the level of abstraction at which a shift occurs.
For example, when the colour of the
EMNIST characters is changed under the crystals shift (see Fig. 4 in Appendix B.1), we intuitively expect the low-level ﬁlters (e.g. black-and-white edge detectors) to be surprised. Fig. 2a shows this to indeed be the case with units in the ﬁrst layer exhibiting high surprise.
Naturally, the abnormal activations which surprised these units propagate through the rest of the network causing the succeeding units to also be surprised.
In contrast, when the network is presented with unseen classes, we intuitively expect the low-level features to be unsurprised (similar distribution of black-and-white edges) but the later layers to be surprised (new combinations of low-level features).
Fig. 2b shows this to indeed be the case with units in the later layers surprised while those in earlier layers are not.
Creating an update rule.
Once units have been equipped with the ability to calculate their surprise, how can we best
[Table]
[Table]
make use of this?
We investigate whether or not surprise can give us some indication of the speciﬁc few units to update in order to adapt quickly to new data.
In particular, we focus on few-shot learning as updating (the parameters of) fewer units should improve sample-efﬁciency.
We consider three possible adaptation strategies or shift responses :
SGD.
Ignore surprise patterns and use stochastic gradient descent (SGD) to update all parameters. FlexTune [35] .
Ignore surprise patterns and select the best individual layer to update using a test-domain validation set.
This serves as an upper-bound on single-layer performance Update rule based on unit-level surprise
Use a rule that determines whether or not a unit should update based on surprise values.
We design an update rule with the goal of leveraging unit-level surprise to preserve structure by preventing unnecessary updates.
This rule can be summarised as “update only if: (i) you are surprised; and (ii) your parents are not” , with the key insight being that surprised parents (units in the preceding layer that are connected to it) indicate that the surprise-causing shift may be resolved in earlier layers.
Here, (i) ensures that only affected units update— “This looks the same as before, no need for me to update!”
—while (ii) acts as a modulatory update-in-progress signal from a unit’s parents— “This is being dealt with by someone else, hold tight!” .
We formally deﬁne this update rule in Appendix B.5.
Results.
Table 1 gives the 5-shot results for training individual network layers and using each of ourresponses.
We see that low-level shifts are best dealt with by ﬁne-tuning the ﬁrst layer and high-level shifts the last.
This aligns with our intuition about which units should update and also with the surprise patterns in Figure 2.
Without a priori knowledge of the type of shift that will occur (choosing
Conv1 for low-level shifts, FC2 for high-level) or exhaustively searching over all possibilities using an oracle-like test-domain validation set (FlexTune), we achieve the best average performance over low- and high-level shifts with our surprise-based update rule—where units only update if they are sufﬁciently surprised and their parents are not.
These results show that it is not always optimal to ﬁne-tune only the later layers of a network (as also found in [35]) and demonstrate that unit-level surprise can indeed help determine which few parameters should update to adapt quickly (i.e. with few samples) to new data.
Results for individual shifts are given in Appendix E.
We also evaluate each of these adaptation strategies or responses as a function of the number of shifted samples that are available.
As shown in Table 2, all strategies perform well using large amounts of data (2000 samples-per-class).
However, as the number of samples-per-class drops, it becomes more important to update few parameters, and thus to choose the right strategy.
With 5samples-per-class, simple SGD (training all layers) is outperformed by both other strategies which seek to maintain structure by selecting few units to update.
While FlexTune consistently outperforms the update rule (by a small margin), it requires training 5 different models on labelled test-domain data (one per layer) in order to select the best layer to update.
Which units does SGD update?
To visualize which parameters are updated we calculate “how far” each unit’s input parameters are moved and use this to compare SGD and our update rule.
Speciﬁcally, we calculate the Euclidean distance between a parameter’s value before and after adaptation, then average these distances layerwise to see which layers are being updated by each adaptation strategy.
Figures 3a & 3b show that SGD moves all layers of the network similarly and makes no distinction between different types of shift. This conﬁrms that a simple gradient signal is not sufﬁcient to select C1 C2 C3 FC1 FC2 Layer 0. A v g .
D i s t .
M o v e d (a) Low, SGD C1 C2 C3 FC1 FC2 Layer 0. A v g . D i s t .
M o v e d (b) High, SGD
C1 C2 C3 FC1 FC2 Layer 0.
A v g .
D i s t .
M o v e d (c) Low, Upd.
Rule
C1 C2 C3 FC1 FC2 Layer 0.
A v g .
D i s t .
M o v e d (d) High, Upd.
Rule which few parameters should update to a given shift, explaining the results in Table 1.
In contrast,
Figures 3c & 3d show that our update rule exhibits shift-speciﬁc behaviour which aligns with the
While we achieved some promising results in carefully-designed settings, it remains unclear how unit-level surprise should be used in more realistic settings to determine which few parameters to update.
Below we discuss the main challenges discovered during the initial idea validation above.
Networks are not sufﬁciently modular to show sparse surprise patterns.
For the various shifts we considered, networks did not exhibit sparse or modular surprise patterns—low-level shifts surprised almost every unit in the network (Fig. 2a), while high-level shifts surprised all units in both FC layers (Fig. 2b).
This then affects the updates of our update rule, with single, entire layers being updated rather than a modular selection of units (Fig. 3).
We posit that such patterns of surprise occur because, using standard training techniques, simple neural networks do not form modules or mechanisms that change independently (and thus can be surprised independently).
Moreover, we ﬁnd ourselves asking if it is ever possible to learn such modules or mechanisms that change sparsely (i.e. align with the shifts in data distribution) without explicitly optimizing for it at training time over multiple shifts.
Parents and children may need to update.
As discussed above, our update rule tends to update single layers, and this can lead to sub-optimal solutions (e.g. updating only FC1 for high-level shifts).
We attribute this to the fact that surprised parent units tend to stay surprised, continually blocking their children from updating.
One solution would be to update P ( A ) with samples from Q ( A ) as training progresses, for units that are being adapted.
Intuitively, this would allow parent units to gradually become less surprised by the new data, eventually freeing their children to also update.
Surprise may not be sufﬁcient.
As shown in Table 1 and Figure 3, the update rule does not always select the “optimal” units to update (updating FC1 rather than FC2 for high-level shifts).
This raises some questions—is surprise alone sufﬁcient to determine which units should update?
What other unit-level information may be required?
Here we list some examples: (a) gradient magnitude —may help us to update FC2 rather than FC1 to adapt to high-level shifts; (b) task importance [21, 43]—may help in continual learning settings to best preserve performance on past tasks when adapting to new ones; (c) parameter uncertainty —may help speed-up learning by only updating the parameters with high uncertainty [3], akin to Bayesian optimization.
Additionally, one can imagine shifts that are best resolved at a mid- or high-level but cause changes to the low-level feature distributions—e.g. if we occlude part of a digit, the low-level features (edges) should still be valid, we just need to update how they combine to form the class label.
However, the changed low-level feature distributions (some edges are occluded) will surprise units in the early layers, causing them to update.
Another similar challenge is that of changing feature frequency.
In particular, if we take the view of units as feature detectors, where a speciﬁc unit is sensitive to a speciﬁc feature, then the activation distributions of units will change if the features they detect occur with a different frequency (i.e. appear more or less often).
In this situation, a unit can show high surprise
but we may not wish to update its parameters if we still wish to detect this feature.
One potential solution to these problems is to explore alternative measures that distinguish between being more or less surprised than expected (see Appendix C).
Lack of interpretability makes validation difﬁcult.
Finally, it is difﬁcult to come up with experimental setups where we know what the “right” units to update are.
This makes the design and validation of update rules quite challenging.
While we may hope for interpretable edges-parts-wholes feature hierarchies (see Figure 9 in Appendix E.1), this is often not the case.
In fact, it can take as little as a change in random seed to change the semantic meaning of features, and with it, the optimal units to update.
Work on unit-level interpretability [14, 29] may eventually help, but it is not currently at a level to decide which units should update for a given shift.
Future work
As discussed above, creating a general update rule using unit-level surprise is challenging as: (i) neural networks are not very modular by default; and (ii) surprise may not be sufﬁcient to determine “who” should update.
Below we discuss two potential avenues to overcome these two challenges.
Learning modular structure with unit-level surprise.
Recent work in causal discovery assumes that the ground-truth data-generative process consists of independent mechanisms or modules, with many modules expected to behave similarly across different tasks and environments [5, 31, 32, 36, 37].
Thus, if an appropriate modular representation of the world has been learned, very few modules should need to be adapted in order to transfer [5, 15, 32, 37].
Bengio et al.
[5] exploit this through a meta-learning objective that uses adaptation speed to optimize the way in which knowledge is represented or modularized.
Here, adaptation speed is a proxy score for the number of modules and parameters that need be adapted and ultimately for how well the learned modularization ﬁts the underlying causal dependencies.
We believe surprise could provide an alternative proxy score (e.g. the number 2 of surprised units or groups of units) that is easier to meta-optimize (no inner-loop steps) and arguably a more direct measure of the number of modules and parameters that need be adapted.
Learning optimizers to better utilize unit-level surprise.
As discussed in the previous section, unit-level surprise may not be sufﬁcient for some shifts—further information may be required to determine “who” should update.
This makes it difﬁcult to handcraft a general update rule as one must specify how surprise should be used with this other information.
In fact, even without additional information, there is still an enormous space to be explored surrounding how best to use unit-level surprise, e.g. setting the learning rate based on the surprise by using soft thresholding.
One solution is to learn an optimizer [4, 33], or the parameters of an update rule [25], across multiple shifts.
However, learned optimizers are notoriously difﬁcult to train [28] and it can be difﬁcult to beat the heavily-tuned few-shot benchmarks against which they are often compared [16].
Conclusion
In this work we provided a preliminary validation of the usefulness of unit-level surprise in neural networks.
In particular, we showed that it can be useful for analysing where dataset shifts are noticed in a network and for devising shift-dependent ﬁne-tuning strategies.
We also discussed some of the “beauty” that makes us believe that unit-level surprise should be useful in more general settings—it often aligns with our intuitions about how to handle OOD data, it seems a strong signal for learning how to modularize knowledge such that only a few modules are surprised (and thus need be adapted), and it has several neuroscience relations such as metaplasticity and the free-energy principle.
However, in formally specifying how to use unit-level surprise in simple settings, we encountered several challenges that made us wary of tackling the more complicated settings in which its full potential could be realized. Acknowledgements
We thank Matthias Hennig, Tim Hospadales, Amos Storkey, and Sébastian Arnold for helpful discussions and comments.
CE acknowledges support from The National University of Ireland via his Travelling Studentship in the Sciences.
IM is supported by the Engineering and Physical Sciences
Research Council (EPSRC).
A Example to further motivate surprise
Imagine you are given a network which has been trained to classify different Ford car models outdoors (in green pastures), and asked to classify pictures of these same models indoors (in a showroom).
Intuitively, most of the learned function is still valid so only a small number of parameters need to be tweaked, requiring only a small number of indoor examples to do so.
Perhaps tweaking parameters in the ﬁrst layer would be sufﬁcient, re-extracting the same features from e.g. the darker pixels of indoor images as were previously extracted from the outdoor images.
Similarly, if you were asked to classify pictures of Fiat cars, you may expect that only a small number of parameters need to be adapted (likely in later layers this time), perhaps adjusting the (conceptual) wheel and mirror detectors for the smaller wheels and more rounded mirrors of Fiats.
However, as we discuss in
Section 2, minimizing a standard loss function (e.g. cross-entropy) with stochastic gradient descent (SGD) does not result in such intuitive updates—all of the network’s parameters are updated and learnt structure is unnecessarily destroyed.
These gradients tell us which parameters can reduce the error, but not which parameters should reduce the error in order to maximally-transfer knowledge and thus speed-up learning.
Doing so requires additional information—such as unit-level surprise.
B Implementation details B.1 Data
Our networks are trained using the 47-class EMNIST dataset [7] (“identity” shift in Figure 4).
We use 2000 samples per class from the training split with the remaining 400 forming a validation set, we report results on the separate test set also containing 400 examples per class.
We adapt to 10 new data distributions—7 low-level shifts from EMNIST-DA [9] and 3 high-level label shifts.
From the 14 EMNIST-DA shifts depicted in Figure 4, we chose 7 shifts that adversely affect accuracy (without adaptation) and intuitively affect the early convolutional ﬁlters: crystals, fog, gaussian blur, grass, impulse noise, sky and stripe.
To create the 3 label shifts, we train our networks on only the ﬁrst 37 classes of EMNIST and then choose 5 of the 10 unseen classes three times from the range [38 , 47] to arrive at: H1: [38 , 39 , 40 , 41 , 42] , H2: [43 , 44 , 45 , 46 , 47] , and H3:
To evaluate sample efﬁciency we run experiments with varying amounts of data, we experiment with using 2 , 5 , 10 , 20 and 50 samples per class as well as using all the data ( 2000 samples per class).
The full results of these experiments are given in Appendix E.
B.2 Experimental setup
Table 3 provides the architectural details of the simple 5-layer convolutional neural network (CNN) that we use.
During pre-training we use dropout between the layers, for adaptation we do not as it unnecessarily complicates the propagation of surprise and makes little difference to the ﬁnal results.
During pre-training and adaptation we use a batch size of 256 .
We pre-train for 150 epochs with a learning rate of 0 . 01 .
Due to using small amount of data, during adaptation we train with early stopping for a maximum of 100 epochs using a patience of 10 . We use a learning rate of 0 . 1 for all experiments except for when ﬁne-tuning all layers simultaneously which requires a learning rate of 0 .
01 to prevent divergence.
We optimize using stochastic gradient descent with momentum set to 0 .
For experiments using our update rule the thresholds α and β in Equation 3 are set to 0 .
Experiments are run over 3 seeds from which we report a mean and one standard deviation.
[Table]
B.3 Measuring unit-level surprise
A single unit in a feed-forward neural network outputs an activation a = g ( w T h + b ) , where h is the hidden unit activations of the previous layer, w the learned weight vector, b the learned bias and g some non-linearity.
During training a unit can store a distribution P ( A ) which captures the distribution that its activation can take.
A unit is surprised by new data if the activation distribution changes, i.e. P ( A ) (cid:54) =
Q ( A ) , where Q ( A ) is the distribution of the unit’s activations under the new data.
We quantify surprise using the KL-divergence from P ( A ) to Q ( A ) , i.e. s ( A ) =
D KL ( Q ( A ) ||
P ( A ))
The surprisal (or information content) of an event X = x , with X ∼ P ( X ) , is given by log(1 /P
( x )) .
Intuitively, this quantity represents how “surprised” we are to see X
= x , with unlikely events having high surprisal.
Surprise itself is a somewhat overloaded term but can be used to describe the entropy, H ( X )
= − (cid:80) P ( x ) log P ( x ) , that is the expected surprisal.
If we now receive a sample of a different random variable Y =
y , Y ∼ Q ( Y ) , but we have assumed as a prior that we are receiving samples from P ( X ) , the amount of additional surprisal we receive on account of our assumption is log(1 /P
( X = y ))
− log(1 /Q
( Y = y )) .
The expected value of this quantity over Q is the KLDivergence D KL ( Q || P )
[22], which is the expected surprisal of receiving samples from Q when we have assumed the distribution to be P .
We can also interpret D KL ( Q || P )
=
H ( Q, P ) − H ( Q ) as the expected extra message-length per datum that must be communicated if a code that is optimal for
P is used to communicate Q , compared to a code that is optimal for Q .
We have seen this quantity referred to as Bayesian surprise, information gain, asymmetric surprise or simply surprise [10, 19].
Throughout this work we refer to this quantity as surprise for simplicity.
B.4 Calculating surprise from bin counts
After pre-training we parameterize activation distributions with softly-binned histograms.
To calculate
P ( A ) we run one further forward-pass of the network over the training data and bin the activations using the same procedure as in [9], with 10 bins, which outputs 10 normalized bin counts π p 1 , . . .
π pfor each unit.
π pi represents the probability a falls into bin i and (cid:80) 10 i =1 π pi = 1 .
The blue curves in Fig. 5 depict examples of such distributions. These distributions can be considered as representing the “normal” activation values of a unit, i.e. the values it expects to take on.
We then receive some data from a new distribution, possibly the same as the pre-training data distribution, which is fed into the network.
During adaptation we can parameterize Q ( A ) in the same way as P ( A ) , using a batch of this new data (Fig. 5–orange curves) to calculate normalized bin counts π q 1 , . . .
π q 10 which change as the network learns.
The surprise for a unit can then be calculated as
[Equation]
B.5 A surprise-based update rule
Let p i denote the surprise of “parent” unit i in layer l , c j the surprise of “child” unit j in layer l + 1 , and w ij the weight that connects parent unit i with child unit j . This setup is depicted in Fig.
6 below. [Figure] [Equation] [Figure]
[Equation] [Figure] For child unit j , we calculate its aggregate parent surprise ¯ p j as a weighted average of its parent surprises.
More speciﬁcally, we calculate
[Equation]
where the normalized weight value [Equation] | ensures comparable scaling across units in a layer. We then use this to create an update rule where the input weights of child j are updated if and only if child j is surprised (i.e. c j is above some threshold α ) but its parents are not (i.e. ¯ p is below some threshold β ).
In particular, we create the following update rule:
[Equation]
[Equation] ) where I [ c j > α ] is an indicator function that is 1 when c j > α and 0 otherwise, I [¯ p j < β ] is similarly deﬁned, η the learning rate, and L the loss function (cross-entropy in our case).
I [ c j > α ] ensures that only surprised units update, and can be compared with metaplasticity 4 in the brain.
I [¯ p j < β ] prevents/blocks simultaneous changes in later units who may also surprised by their input, and can be compared with neuromodulation 5 in the brain.
C Alternative surprise measures
Imagine a unit or feature detector with a bi-modal activation distribution where the modes roughly represent on (detected) and off (not detected).
Perhaps such a unit being off more often in the new data (as when part of an image is occluded, discussed in Section 3) is not a good signal to update.
To differentiate this situation from e.g. the unit being on more often in the new data, we can deﬁne a new measure which we call the surprise increase (SI):
[Equation]
(4) where H ( P ) is the entropy of P and H ( Q, P ) is the cross-entropy of P relative to Q .
Unlike the KL-divergence, SI can be negative.
SI is negative (i.e. a surprise decrease) when an event that is already quite likely under P becomes even more likely under Q —as shown in Figure 7a.
This could be an interesting alternative surprise measurement as it can distinguish between surprise increases and decreases.
[Figure]
D Does batch normalization solve the problem?
Batch normalization (BN, 18) standardizes activation distributions across batches, bringing P ( A ) and
Q ( A ) closer together.
This naturally raises the question as to whether or not BN solves the problem of differing unit-level distributions, thus removing the need for unit-level surprise.
We investigate this below.
Do the same surprise patterns exist?
Figure 8 shows the ideal situation for BN, where the BN statistics are (re)calculated using the new data distribution before we calculate surprise.
Compared to Figure 2, the magnitude of the surprise is indeed lessened as BN standardizes P ( A ) and Q ( A ) , but it does not make units unsurprised.
Moreover, we still see the same patterns of surprise in the early layers for low-level shifts and in the later layers for high-level shifts.
Is an adaptative strategy still best?
As shown in Table 4, it is indeed still best to employ an adaptive update strategy to train speciﬁc layers for speciﬁc shifts, e.g. use FlexTune [35] to select
Conv1 & BN1 for low-level shifts and FC2 for high-level shifts.
This further conﬁrms that BN alone does not solve the problem of changes in activation distribution, as selective adaptation strategies are still superior to updating all layers with SGD.
In Table 4 we also show that, if we only use the BN statistics without any training (AdaBN, 24), or only update the BN parameters and statistics on the new data (labelled “BN params”), performance is poor compared to the other strategies.
[Figure]
[Table]
E Further Results E.1 Max-activating patches
Figure 9 shows the maximum-activating image patches on the EMNIST-DA [9] grass shift for selected units in each layer.
Note that we do not have a perfect edges-parts-wholes hierarchy—the receptive ﬁeld of Conv2 seems too large as it almost sees the entire image (can be a “whole” rather than a part).
[Figure]
E.2 Per-shift Results
What follows is tables of full results for each shift at each number of shots, these are provided for completeness with no further analysis.
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]