TASKNORM: Rethinking Batch Normalization for Meta-Learning John Bronskill *1 Jonathan Gordon *1 James Requeima 12 Sebastian Nowozin 3 Richard E. Turner 13 Abstract Modern meta-learning approaches for image clas- sification rely on increasingly deep networks to achieve state-of-the-art performance, making batch normalization an essential component of meta-learning pipelines. However, the hierarchi- cal nature of the meta-learning setting presents several challenges that can render conventional batch normalization ineffective, giving rise to the need to rethink normalization in this setting. We evaluate a range of approaches to batch normal- ization for meta-learning scenarios, and develop a novel approach that we call TASKNORM. Ex- periments on fourteen datasets demonstrate that the choice of batch normalization has a dramatic effect on both classification accuracy and train- ing time for both gradient based- and gradient- free meta-learning approaches. Importantly, TAS- KNORM is found to consistently improve perfor- mance. Finally, we provide a set of best practices for normalization that will allow fair comparison of meta-learning algorithms. 1. Introduction Meta-learning, or learning to learn (Thrun & Pratt, 2012; Schmidhuber, 1987), is an appealing approach for design- ing learning systems. It enables practitioners to construct models and training procedures that explicitly target de- sirable charateristics such as sample-efficiency and out-of- distribution generalization. Meta-learning systems have been demonstrated to excel at complex learning tasks such as few-shot learning (Snell et al., 2017; Finn et al., 2017) and continual learning (Nagabandi et al., 2019; Requeima et al., 2019a; Jerfel et al., 2019). Recent approaches to meta-learning rely on increasingly deep neural network based architectures to achieve state-of- * Equal contribution 1 University of Cambridge 2 Invenia Labs 3 Microsoft Research. Correspondence to: John Bronskill <[email protected]>. Proceedings of the 37 th International Conference on Machine Learning, Vienna, Austria, PMLR 119, 2020. Copyright 2020 by the author(s). the-art performance in a range of benchmark tasks (Finn et al., 2017; Mishra et al., 2018; Triantafillou et al., 2020; Requeima et al., 2019a). When constructing very deep net- works, a standard component is the use of normalization layers (NL). In particular, in the image-classification do- main, batch normalization (BN; Ioffe, 2017) is crucial to the successful training of very deep convolutional networks. However, as we discuss in Section 3, standard assumptions of the meta-learning scenario violate the assumptions of BN and vice-versa, complicating the deployment of BN in meta-learning. Many papers proposing novel meta-learning approaches employ different forms of BN for the proposed models, and some forms make implicit assumptions that, while improving benchmark performance, may result in potentially undesirable behaviours. Moreover, as we demon- strate in Section 5, performance of the trained models can vary significantly based on the form of BN employed, con- founding comparisons across methods. Further, naive adop- tion of BN for meta-learning does not reflect the statistical structure of the data-distribution in this scenario. In contrast, we propose a novel variant of BN – TASKNORM – that explicitly accounts for the statistical structure of the data distribution. We demonstrate that by doing so, TASKNORM further accelerates training of models using meta-learning while achieving improved test-time performance. Our main contributions are as follows: • We identify and highlight several issues with BN schemes used in the recent meta-learning literature. • We propose TASKNORM, a novel variant of BN which is tailored for the meta-learning setting. • In experiments with fourteen datasets, we demonstrate that TASKNORM consistently outperforms competing meth- ods, while making less restrictive assumptions than its strongest competitor. 2. Background and Related Work In this section we lay the necessary groundwork for our investigation of batch normalization in the meta-learning scenario. Our focus in this work is on image classification. We denote images x ∈ R C×W×H where W is the image width, H the image height, C the number of image channels. Each image is associated with a label y ∈{1,...,M } arXiv:2003.03284v2 [stat.ML] 28 Jun 2020
Transcript
TASKNORM: Rethinking Batch Normalization for Meta-Learning
John Bronskill * 1 Jonathan Gordon * 1 James Requeima 1 2 Sebastian Nowozin 3 Richard E. Turner 1 3
AbstractModern meta-learning approaches for image clas-sification rely on increasingly deep networksto achieve state-of-the-art performance, makingbatch normalization an essential component ofmeta-learning pipelines. However, the hierarchi-cal nature of the meta-learning setting presentsseveral challenges that can render conventionalbatch normalization ineffective, giving rise to theneed to rethink normalization in this setting. Weevaluate a range of approaches to batch normal-ization for meta-learning scenarios, and developa novel approach that we call TASKNORM. Ex-periments on fourteen datasets demonstrate thatthe choice of batch normalization has a dramaticeffect on both classification accuracy and train-ing time for both gradient based- and gradient-free meta-learning approaches. Importantly, TAS-KNORM is found to consistently improve perfor-mance. Finally, we provide a set of best practicesfor normalization that will allow fair comparisonof meta-learning algorithms.
1. IntroductionMeta-learning, or learning to learn (Thrun & Pratt, 2012;Schmidhuber, 1987), is an appealing approach for design-ing learning systems. It enables practitioners to constructmodels and training procedures that explicitly target de-sirable charateristics such as sample-efficiency and out-of-distribution generalization. Meta-learning systems havebeen demonstrated to excel at complex learning tasks suchas few-shot learning (Snell et al., 2017; Finn et al., 2017)and continual learning (Nagabandi et al., 2019; Requeimaet al., 2019a; Jerfel et al., 2019).
Recent approaches to meta-learning rely on increasinglydeep neural network based architectures to achieve state-of-
*Equal contribution 1University of Cambridge 2Invenia Labs3Microsoft Research. Correspondence to: John Bronskill<[email protected]>.
Proceedings of the 37 th International Conference on MachineLearning, Vienna, Austria, PMLR 119, 2020. Copyright 2020 bythe author(s).
the-art performance in a range of benchmark tasks (Finnet al., 2017; Mishra et al., 2018; Triantafillou et al., 2020;Requeima et al., 2019a). When constructing very deep net-works, a standard component is the use of normalizationlayers (NL). In particular, in the image-classification do-main, batch normalization (BN; Ioffe, 2017) is crucial to thesuccessful training of very deep convolutional networks.
However, as we discuss in Section 3, standard assumptionsof the meta-learning scenario violate the assumptions ofBN and vice-versa, complicating the deployment of BN inmeta-learning. Many papers proposing novel meta-learningapproaches employ different forms of BN for the proposedmodels, and some forms make implicit assumptions that,while improving benchmark performance, may result inpotentially undesirable behaviours. Moreover, as we demon-strate in Section 5, performance of the trained models canvary significantly based on the form of BN employed, con-founding comparisons across methods. Further, naive adop-tion of BN for meta-learning does not reflect the statisticalstructure of the data-distribution in this scenario. In contrast,we propose a novel variant of BN – TASKNORM – thatexplicitly accounts for the statistical structure of the datadistribution. We demonstrate that by doing so, TASKNORMfurther accelerates training of models using meta-learningwhile achieving improved test-time performance. Our maincontributions are as follows:
• We identify and highlight several issues with BN schemesused in the recent meta-learning literature.
• We propose TASKNORM, a novel variant of BN which istailored for the meta-learning setting.
• In experiments with fourteen datasets, we demonstratethat TASKNORM consistently outperforms competing meth-ods, while making less restrictive assumptions than itsstrongest competitor.
2. Background and Related WorkIn this section we lay the necessary groundwork for ourinvestigation of batch normalization in the meta-learningscenario. Our focus in this work is on image classification.We denote images x ∈ RC×W×H where W is the imagewidth,H the image height, C the number of image channels.Each image is associated with a label y ∈ {1, . . . ,M}
arX
iv:2
003.
0328
4v2
[st
at.M
L]
28
Jun
2020
TASKNORM: Rethinking Batch Normalization for Meta-Learning
xτ∗myτ∗mψτyτnxτn
θ
Context Dτ Target T τ
m = 1, ...,Mτn = 1, ..., Nτ
τ=1,...
Figure 1. Directed graphical model for multi-task meta-learning.
where M is the number of image classes. Finally, a datasetis denoted D = {(xn, yn)}Nn=1.
2.1. Meta-Learning
We consider the meta-learning classification scenario.Rather than a single, large dataset D, we assume access to adataset D = {τt}Kt=1 comprising a large number of trainingtasks τt, drawn i.i.d. from a distribution p(τ). The data fora task τ consists of a context set Dτ = {(xτn, yτn)}
Nτn=1 with
Nτ elements with the inputs xτn and labels yτn observed, anda target set T τ = {(xτ∗m , yτ∗m )}Mτ
m=1 with Mτ elements forwhich we wish to make predictions. Here the inputs xτ∗ areobserved and the labels yτ∗ are only observed during meta-training (i.e., training of the meta-learning algorithm). Theexamples from a single task are assumed i.i.d., but examplesacross tasks are not. Note that the target set examples aredrawn from the same set of labels as the examples in thecontext set.
At meta-test time, the meta-learner is required to makepredictions for target set inputs of unseen tasks. Often, theassumption is that test tasks will include classes that havenot been seen during meta-training, and Dτ will containonly a few observations. The goal of the meta-learner is toprocess Dτ , and produce a model that can make predictionsfor any test inputs xτ∗ ∈ T τ∗ associated with the task.
Meta-Learning as Hierarchical Probabilistic ModellingA general and useful view of meta-learning is through theperspective of hierarchical probabilistic modelling (Heskes,2000; Bakker & Heskes, 2003; Grant et al., 2018; Gordonet al., 2019). A standard graphical representation of thismodelling approach is presented in Figure 1. Global param-eters θ encode information shared across all tasks, whilelocal parameters ψτ encode information specific to task τ .This model introduces a hierarchy of latent parameters, cor-responding to the hierarchical nature of the data distribution.
A general approach to meta-learning is to design inferenceprocedures for the task-specific parameters ψτ = fφ(D
τ )conditioned on the context set (Grant et al., 2018; Gor-don et al., 2019), where f is parameterized by additionalparameters φ. Thus, a meta-learning algorithm definesa predictive distribution parameterized by θ and φ as
p (yτ∗m |xτ∗m , fφ (Dτ ) ,θ) . This perspective relates to the in-ner and outer loops of meta-learning algorithms (Grantet al., 2018; Rajeswaran et al., 2019): the inner loop uses fφto provide local updates to ψ, while the outer loop providespredictions for target points. Below, we use this view tosummarize a range of meta-learning approaches.
Episodic Training The majority of modern meta-learningmethods employ episodic training (Vinyals et al., 2016).During meta-training, a task τ is drawn from p(τ) and ran-domly split into a context set Dτ and target set T τ . Themeta-learning algorithm’s inner-loop is then applied to thecontext set to produce ψτ . With θ and ψτ , the algorithmcan produce predictions for the target set inputs xτ∗m .
Given a differentiable loss function, and assuming that fφ isalso differentiable, the meta-learning algorithm can then betrained with stochastic gradient descent algorithms. Usinglog-likelihood as an example loss function, we may expressa meta-learning objective for θ and φ as
L(θ,φ) = Ep(τ)
[Mτ∑m=1
log p (yτ∗m |xτ∗m , fφ (Dτ ) ,θ)
]. (1)
Common Meta-Learning Algorithms There has beenan explosion of meta-learning algorithms proposed in re-cent years. For an in-depth review see Hospedales et al.(2020). Here, we briefly introduce several methods, focus-ing on those that are relevant to our experiments. Arguablythe most widely used is the gradient-based approach, thecanonical example for modern systems being MAML (Finnet al., 2017). MAML sets θ to be the initialization of theneural network parameters. The local parameters ψτ arethe network parameters after applying one or more gradientupdates based on Dτ . Thus, f in the case of MAML isa gradient-based procedure, which may or may not haveadditional parameters (e.g., learning rate).
Another widely used class of meta-learners are amortized-inference based approaches e.g, VERSA (Gordon et al.,2019) and CNAPS (Requeima et al., 2019a). In these meth-ods, θ parameterizes a shared feature extractor, and ψ aset of parameters used to adapt the network to the localtask, which include a linear classifier and possibly addi-tional parameters of the network. For these models, f isimplemented via hyper-networks (Ha et al., 2016) with pa-rameters φ. An important special case of this approachis Prototypical Networks (ProtoNets) (Snell et al., 2017),which replace ψ with nearest neighborhood classificationin the embedding space of a learned feature extractor gθ.
2.2. Normalization Layers in Deep Learning
Normalization layers (NL) for deep neural networks wereintroduced by Ioffe & Szegedy (2015) to accelerate the train-
TASKNORM: Rethinking Batch Normalization for Meta-Learning
ing of neural networks by allowing the use of higher learningrates and decreasing the sensitivity to network initialization.Since their introduction, they have proven to be crucialcomponents in the successful training of ever-deeper neuralarchitectures. Our focus is the few-shot image classificationsetting, and as such we concentrate on NLs for 2D convolu-tional networks. The input to a NL is A = (a1, . . . ,aB), abatch of B image-shaped activations or pre-activations, towhich the NL is applied as
a′n = γ
(an − µ√σ2 + ε
)+ β, (2)
where µ and σ are the normalization moments, γ and β arelearned parameters, ε is a small scalar to prevent division by0, and operations between vectors are element-wise. NLsdiffer primarily by how the normalization moments arecomputed. The first such layer – batch normalization (BN)– was introduced by Ioffe & Szegedy (2015). A BN layerdistinguishes between training and test modes. At trainingtime, BN computes the moments as
µBNc =1
BHW
B∑b=1
W∑w=1
H∑h=1
abwhc, (3)
σ2BNc =
1
BHW
B∑b=1
W∑w=1
H∑h=1
(abwhc − µBNc)2. (4)
Here, µBN ,σ2BN ,γ,β ∈ RC . As µBN and σ2
BN dependon the batch of observations, BN can be susceptible tofailures if the batches at test time differ significantly fromtraining batches, e.g., for streaming predictions. To coun-teract this, at training time, a running mean and variance,µr,σr ∈ RC , are also computed for each BN layer over alltraining tasks and stored. At test time, test activations a arenormalized using Equation (2) with the statistics µr and σrin place of the batch statistics. Importantly, BN relies on theimplicit assumption that D comprises i.i.d. samples fromsome underlying distribution.
More recently, additional NLs have been introduced. Manyof these methods differ from standard BN in that they nor-malize each instance independently of the remaining in-stances in the batch, making them more resilient to batchdistribution shifts at test time. These include instance nor-malization (Ulyanov et al., 2016), layer normalization (Baet al., 2016), and group normalization (Wu & He, 2018).These are discussed further in Section 3.3.
2.3. Desiderata for Meta-Learning NormalizationLayers
As modern approaches to meta-learning systems routinelyemploy deep networks, NLs become essential for efficienttraining and optimal classification performance. For BN inthe standard supervised settings, i.i.d. assumptions about the
data distribution imply that estimating moments from thetraining set will provide appropriate normalization statisticsfor test data. However, this does not hold in the meta-learning scenario, for which data points are only assumedto be i.i.d. within a specific task. Therefore, the choice ofwhat moments to use when applying a NL to the contextand target set data points, during both meta-training andmeta-test time, is key.
As a result, recent meta-learning approaches employ severalnormalization procedures that differ according to these de-sign choices. A range of choices are summarized in Figure 2.As we discuss in Section 3 and demonstrate with experi-mental results, some of these have implicit, undesirableassumptions which have significant impact on both predic-tive performance and training efficiency. We argue thatan appropriate NL for the meta-learning scenario requiresconsideration of the data-generating assumptions associatedwith the setting. In particular, we propose the followingdesiderata for a NL when used for meta-learning:
1. Improves speed and stability of training without harmingtest performance (test set accuracy or log-likelihood);
2. Works well across a range of context set sizes;
3. Is non-transductive, thus supporting inference at meta-test time in a variety of circumstances.
A non-transductive meta-learning system makes predictionsfor a single test set label conditioned only on a single in-put and the context set, while a transductive meta-learningsystem conditions on additional samples from the test set:
p(yτ∗i |xτ∗i , Dτ )
non-transductive
; p(yτ∗i |xτ∗i=1:m, Dτ )
transductive
. (5)
We argue that there are two key issues with transductivemeta-learners. The first is that transductive learning is sen-sitive to the distribution over the target set used duringmeta-training, and as such is less generally applicable thannon-transductive learning. For example, transductive learn-ers may fail to make good predictions if target sets containa different class balance than what was observed duringmeta-training, or if they are required to make predictionsfor one example at a time. Transductive learners can alsoviolate privacy constraints. In Table 1 and Appendix D, weprovide empirical demonstrations of these failure cases.
The second issue is that transductive learners have moreinformation available than non-transductive learners at pre-diction time, which may lead to unfair comparisons. It isworth noting that some meta-learning algorithms are specif-ically designed to leverage transductive inference (e.g., Renet al., 2018; Liu et al., 2019), though we do not discussthem in this work. In Section 5 we demonstrate that thereare significant performance differences for a model whentrained transductively versus non-transductively.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
𝑨 𝑨∗
𝑁 𝑀𝐵𝑁
𝑨′
𝑁𝑀𝐵𝑁
𝑨′∗
𝑨
𝑨′
𝑨∗
𝑨′∗
𝝁𝑟 , 𝝈𝑟𝑁𝑁
Conventional Batch
Normalization
(CBN)
𝑨
𝑨′
𝑨∗
𝑨′∗
𝑁 𝑀𝐵𝑁 𝑁
MetaBN
𝑨 𝑨∗
𝑁 𝑀𝐵𝑁
𝑨′
𝑁𝑀𝐵𝑁
𝑨′∗
Transductive Batch
Normalization (TBN)
𝑨 𝑨∗
𝑁 𝑀𝐿𝑁
𝑨′
𝑁𝑀𝐿𝑁
𝑨′∗
Layer
Normalization
(LN)
𝑨 𝑨∗
𝑁 𝑀𝐼𝑁
𝑨′
𝑁𝑀𝐼𝑁
𝑨′∗
Instance
Normalization
(IN)TaskNorm-I, RN𝑨
𝑁
𝑀𝐼𝑁
𝑨′
𝑀𝐵𝑁
1 − 𝛼
𝑨∗
𝑁
𝑀𝐼𝑁
𝑨∗′
𝛼1 − 𝛼
𝐵𝐶
Batch Normalization 𝐵𝑁
𝐵𝐶
Layer Normalization 𝐿𝑁𝐵𝐶
Instance Normalization 𝐼𝑁
TaskNorm-L𝑨
𝑁
𝑀𝐿𝑁
𝑨′
𝑀𝐵𝑁
1 − 𝛼
𝑨∗
𝑁
𝑀𝐿𝑁
𝑨∗′
𝛼1 − 𝛼
𝑁 Normalize with moments 𝑀 𝑀𝑋𝑋 Compute 𝑋𝑋 moments 𝐴, 𝐴′ : In, Out Context Activations 𝐴∗, 𝐴′∗: In, Out Target Activations
Meta-train
Meta-test
Figure 2. A range of options for batch normalization for meta-learning. The cubes on the left depict the dimensions over which differentmoments are calculated for normalization of 2D convolutional layers. The computational diagrams on the right show how context andtarget activations are processed for various normalization methods. For all methods except conventional BN (CBN), the processing isidentical at meta-train and meta-test time. Cube diagrams are derived from Wu & He (2018).
3. Normalization Layers for Meta-learningIn this section, we discuss several normalization schemesthat can and have been applied in the recent meta-learningliterature, highlighting the modelling assumptions and ef-fects of different design choices. Throughout, we assumethat the meta-learning algorithm is constructed such that thecontext-set inputs are passed through every neural-networkmodule that the target set inputs are passed through at predic-tion time. This implies that moments are readily availablefrom both the context and target set observations for anynormalization layer, and is the case for many widely-usedmeta-learning models (e.g., Finn et al., 2017; Snell et al.,2017; Gordon et al., 2019).
To illustrate our arguments, we provide experiments withMAML running simple, but widely used few-shot learningtasks from the Omniglot (Lake et al., 2011) and miniIma-genet (Ravi & Larochelle, 2017) datasets. The results ofthese experiments are provided in Table 1, and full experi-mental details in Appendix B.
3.1. Conventional Usage of Batch Normalization (CBN)
We refer to conventional batch normalization (CBN) as thatdefined by Ioffe & Szegedy (2015) and as outlined in Sec-tion 2.2. In the context of meta-learning, this involves nor-malizing tasks with computed moments at meta-train time,and using the accumulated running moments to normalizethe tasks at meta-test time (see CBN in Figure 2).
We highlight two important issues with the use of CBN formeta-learning. The first is that, from the graphical modelperspective, this is equivalent to lumping µ and σ with theglobal parameters θ, i.e., they are learned from the meta-training set and shared across all tasks at meta-test time.
We might expect CBN to perform poorly in meta-learningapplications since the running moments are global acrossall tasks while the task data is only i.i.d. locally withina task, i.e., CBN does not satisfy desiderata 1. This iscorroborated by our results (Table 1), where we demonstratethat using CBN with MAML results in very poor predictiveperformance - no better than chance. The second issue isthat, as demonstrated by Wu & He (2018), using small batchsizes leads to inaccurate moments, resulting in significantincreases in model error. Importantly, the small batch settingmay occur often in meta-learning, for example in the 1-shotscenario. Thus, CBN does not satisfy desiderata 2.
Despite these issues, CBN is sometimes used, e.g., by Snellet al. (2017), though testing was performed only on Om-niglot and miniImagenet where the distribution of tasks ishomogeneous (Triantafillou et al., 2020). In Section 5, weshow that Batch renormalization (BRN; Ioffe, 2017) can ex-hibit poor predictive performance in meta-learning scenarios(see Appendix A.1 for further details).
3.2. Transductive Batch Normalization (TBN)
Another approach is to do away with the running momentsused for normalization at meta-test time, and replace thesewith context / target set statistics. Here, context / target setstatistics are used for normalization, both at meta-train andmeta-test time. This is the approach taken by the authorsof MAML (Finn et al., 2017),1 and, as demonstrated in ourexperiments, seems to be crucial to achieve the reportedperformance. From the graphical model perspective, thisimplies associating the normalization statistics with neitherθ nor ψ, but rather with a special set of parameters that is
1See for example (Finn, 2017) for a reference implementation.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
local for each set (i.e., normalization statistics for T τ are in-dependent of Dτ ). We refer to this approach as transductivebatch normalization (TBN; see Figure 2).
Unsurprisingly, Nichol et al. (2018) found that using TBNprovides a significant performance boost in all cases theytested, which is corroborated by our results in Table 1. Inother words, TBN achieves desiderata 2, and, as we demon-strate in Section 5, desiderata 1 as well. However, it istransductive. Due to the ubiquity of MAML, many compete-tive meta-learning methods (e.g. Gordon et al., 2019) haveadopted TBN. However, in the case of TBN, transductivityis rarely stated as an explicit assumption, and may oftenconfound the comparison among methods (Nichol et al.,2018). Importantly, we argue that to ensure comparisons inexperimental papers are rigorous, meta-learning methodsthat are transductive must be labeled as such.
3.3. Instance-Based Normalization Schemes
An additional class of non-transductive NLs are instance-based NLs. Here, both at meta-train and meta-test time,moments are computed separately for each instance, anddo not depend on other observations. From a modellingperspective, this corresponds to treating µ and σ as localat the observation level. As instance-based NLs do notdepend on the context set size, they perform equally wellacross context-set sizes (desiderata 2). However, as wedemonstrate in Section 5, the improvements in predictiveperformance are modest compared to more suitable NLsand they are worse than CBN in terms of training efficiency(thus not meeting desiderata 1). Below, we discuss twoexamples, with a third discussed in Appendix A.2.
Layer Normalization (LN; Ba et al., 2016) LN (see Fig-ure 2) has been shown to improve performance comparedto CBN in recurrent neural networks, but does not offerthe same gains for convolutional neural networks (Ba et al.,2016). The LN moments are computed as:
µLNb =1
HWC
W∑w=1
H∑h=1
C∑c=1
abwhc, (6)
σ2LNb
=1
HWC
W∑w=1
H∑h=1
C∑c=1
(abwhc − µLNb)2 (7)
where µLN ,σ2LN ∈ RB . While non-transductive, Table 1
demonstrates that LN falls far short of TBN in terms ofaccuracy. Further, in Section 5 we demonstrate that LNlacks in training efficiency when compared to other NLs.
Instance Normalization (IN; Ulyanov et al., 2016) IN(see Figure 2) has been used in a wide variety of image
generation applications. The IN moments are computed as:
µINbc =1
HW
W∑w=1
H∑h=1
abwhc, (8)
σ2INbc
=1
HW
W∑w=1
H∑h=1
(abwhc − µINbc)2 (9)
where µIN ,σ2IN ∈ RB×C . Table 1 demonstrates that IN
has superior predictive performance to that of LN, but fallsconsiderably short of TBN. In Section 5 we show that INlacks in training efficiency when compared to other NLs.
4. Task NormalizationIn the previous section, we demonstrated that it is not im-mediately obvious how NLs should be designed for meta-learning applications. We now develop TASKNORM, thefirst NL that is specifically tailored towards this scenario.TASKNORM is motivated by the view of meta-learning as hi-erarchical probabilistic modelling, discussed in Section 2.1.Given this hierarchical view of the model parameters, thequestion that arises is, how should we treat the normalizationstatistics µ and σ? Figure 1 implies that the data associatedwith a task τ are i.i.d. only when conditioning on both θand ψτ . Thus, the normalization statistics µ and σ shouldbe local at the task level, i.e., absorbed into ψτ . Further, theview that ψτ should be inferred conditioned on Dτ impliesthat the normalization statistics for the target set should becomputed directly from the context set. Finally, our desirefor a non-transductive scheme implies that any contributionfrom points in the target should not affect the normalizationfor other points in the target set, i.e., when computing µand σ for a particular observation xτ∗ ∈ T τ , the NL shouldonly have access to Dτ and xτ∗.
4.1. Meta-Batch Normalization (METABN)
This perspective leads to our definition of METABN, whichis a simple adaptation of CBN for the meta-learning setting.In METABN, the context set alone is used to compute thenormalization statistics for both the context and target sets,both at meta-train and meta-test time (see Figure 2). Toour knowledge, METABN has not been described in anypublication, but concurrent to this work, it is used in theimplementation of Meta-Dataset (Triantafillou et al., 2019).
METABN meets almost all of our desiderata, it (i) is non–transductive since the normalization of a test input does notdepend on other test inputs in the target set, and (ii) as wedemonstrate in Section 5, it improves training speed whilemaintaining accuracy levels of meta-learning models. How-ever, as we demonstrate in Section 5, METABN performsless well for small context sets. This is because momentestimates will have high-variance when there is little data,
TASKNORM: Rethinking Batch Normalization for Meta-Learning
and is similar to the difficulty of using BN with small-batchtraining (Wu & He, 2018). To address this issue, we intro-duce the following extension to METABN, which yields ourproposed normalization scheme – TASKNORM.
4.2. TASKNORM
The key intuition behind TASKNORM is to normalize a taskwith the context set moments in combination with a setof non-transductive, secondary moments computed fromthe input being normalized. A blending factor α betweenthe two sets of moments is learned during meta-training.The motivation for TASKNORM is as follows: when thecontext setDτ is small (e.g. 1-shot or few-shot learning) thecontext set alone will lead to noisy and inaccurate estimatesof the “true” task statistics. In such cases, a secondary set ofmoments may improve the estimate of the moments, leadingto better training efficiency and predictive performance inthe low data regime. Further, this provides informationregarding xτ∗ at prediction time while maintaining non-transductivity. The pooled moments for TASKNORM arecomputed as:
µTN =αµBN + (1− α)µ+, (10)
σ2TN =α
(σ2BN + (µBN − µTN )2
)+ (1− α)
(σ2+ + (µ+ − µTN )2
), (11)
where µTN ,σTN ∈ RB×C , µ+, σ2+ are additional mo-
ments from a non-transductive NL such as LN or IN com-puted using activations from the example being normalized(see Figure 2), and µBN and σBN are computed from Dτ .Equation (11) is the standard pooled variance when com-bining the variance of two Gaussian estimators.
Importantly, we parameterize α = SIGMOID(SCALE|Dτ |+OFFSET), where the SIGMOID function ensures that 0 ≤α ≤ 1, and the scalars SCALE and OFFSET are learnedduring meta-training. This enables us to learn how mucheach set should contribute to the estimate of task statisticsas a function of the context-set size |Dτ |. Figure 3a depictsthe value of α as a function of context set size |Dτ | for arepresentative set of trained TASKNORM layers. In general,when the context size is suitably large (Nτ > 25), α is closeto unity, i.e., normalization is carried out entirely with thecontext set in those layers. When the context size is smaller,there is a mix of the two sets of moments.
Allowing each TASKNORM layer to separately adapt to thesize of the context set (as opposed to learning a fixed αper layer) is crucial in the meta-learning setting, where weexpect the size of Dτ to vary, and are often particularlyinterested in the “few-shot” regime. Figure 3b plots the lineSCALE|Dτ |+ OFFSET for same set of NLs as Figure 3a.The algorithm has learned that the SCALE parameter is non-zero and the OFFSET is almost zero in all cases indicating
100 101 102
Context Set Size0.600.650.700.750.800.850.900.951.00
Alph
a
Layer 1Layer 2Layer 3Layer 4
(a)
0 100 200 300 400 500Context Set Size
010203040506070
SCAL
E*(C
onte
xt S
ize)+
OFFS
ET
Layer 1Layer 2Layer 3Layer 4
(b)
Figure 3. Plots of: (a) α versus context set size, and (b) α versusSCALE|Dτ |+OFFSET for the first NL in each of the four layersin the feature extractor for the TASKNORM-I model.
the importance of having α being a function of context size.In Appendix E, we provide an ablation study demonstratingthe importance of our proposed parameterization of α. If thecontext size is fixed, we do not use the full parameterization,but learn a single value for alpha directly. The computationalcost of TASKNORM is marginally greater than CBN’s. As aresult, per-iteration time increases only slightly. However,as we show in Section 5, TASKNORM converges faster thanCBN.
In related work, Nam & Kim (2018) define Batch-InstanceNormalization (BIN) that combines the results of CBN andIN with a learned blending factor in order to attenuate unnec-essary styles from images. However, BIN blends the outputof the individual CBN and IN normalization operations asopposed to blending the moments. Finally, we note thatReptile (Nichol et al., 2018) uses a non-transductive form oftask normalization that involves normalizing examples fromthe target set one example at a time with the moments of thecontext set augmented with the single example. We refer tothis approach as reptile normalization or RN. It is easy toshow that RN is a special case of TASKNORM augmentedwith IN when α = |Dτ |/(1 + |Dτ |). In Section 5, we showthat reptile normalization falls short of TASKNORM, sup-porting the intuition that learning the value of α is preferableto fixing a value.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
5. ExperimentsIn this section, we evaluate TASKNORM along with a rangeof competitive normalization approaches.2 The goal ofthe experiments is to evaluate the following hypotheses:(i) Meta-learning algorithms are sensitive to the choice ofNL; (ii) TBN will, in general, outperform non-transductiveNLs; and (iii) NLs that consider the meta-learning dataassumptions (TASKNORM, METABN, RN) will outperformones that do not (CBN, BRN, IN, LN, etc.).
5.1. Small Scale Few-Shot Classification Experiments
We evaluate TASKNORM and a set of NLs using the firstorder MAML and ProtoNets algorithms on the Omniglotand miniImageNet datasets under various way (the numberof classes used in each task) and shot (the number of contextset examples used per class) configurations. This setting issmaller scale, and considers only fixed-sized context andtarget sets. Configuration and training details can be foundin Appendix B.
Accuracy Table 1 and Table C.1 show accuracy resultsfor various normalization methods on the Omniglot andminiImageNet datasets using the first order MAML and theProtoNets algorithms, respectively. We compute the averagerank in an identical manner to Triantafillou et al. (2020).
For MAML, TBN is clearly the best method in terms ofclassification accuracy. The best non-transductive approachis TASKNORM that uses IN augmentation (TASKNORM-I). The two methods using instance-based normalization(LN, IN) do significantly less well than methods designedwith meta-learning desiderata in mind (i.e. TASKNORM,MetaBN, and RN). The methods using running averagesat meta-test time (CBN, BRN) fare the worst. Figure 4acompares the performance of MAML on unseen tasks fromminiImageNet when trained with TBN, IN, METABN, andTASKNORM, as a function of the number of shots per classin Dτ , and demonstrates that these trends are consistentacross the low-shot range.
Note that when meta-testing occurs one example at a time(e.g. in the streaming data scenario) or one class at a time(unbalanced class distribution scenario), accuracy for TBNdrops dramatically compared to the case where all the exam-ples are tested at once. This is an important drawback of thetransductive approach. All of the other NLs in the table arenon-transductive and do not suffer a decrease in accuracywhen tested an example at a time or a class at a time.
Compared to MAML, the ProtoNets algorithm is much lesssensitive to the NL used. Table C.1 indicates that withthe exception of IN, all of the normalization methods yield
2Source code is available at https://github.com/cambridge-mlg/cnaps
good performance. We suspect that this is due to the factthat in ProtoNets employs a parameter-less nearest neigh-bor classifier and no gradient steps are taken at meta-testtime, reducing the importance of normalization. The topperformer is LN which narrowly edges out TaskNorm-L andCBN. Interestingly, TBN is not on top and TASKNORM-Ilags as IN is the least effective method.
Training Speed Figure 4b plots validation accuracy ver-sus training iteration for the first order MAML algorithmtraining on Omniglot 5-way-5-shot. TBN is the most ef-ficient in terms of training convergence. The best non-transductive method is again TASKNORM-I, which is onlymarginally worse than TBN and just slightly better thanTASKNORM-L. Importantly, TASKNORM-I is superior to ei-ther of MetaBN and IN alone in terms of training efficiency.Figure C.1a depicts the training curves for the ProtoNetsalgorithm. With the exception of IN which converges to alower validation accuracy, all NLs converge at the the samespeed.
For the MAML algorithm, the experimental results supportour hypotheses. Performance varies significantly acrossNLs. TBN outperformed all methods in terms of classifica-tion accuracy and training efficiency, and TASKNORM is thebest non-transductive approach. Finally, The meta-learningspecific methods outperformed the more general ones. Thepicture for ProtoNets is rather different. There is little vari-ability across NLs, TBN lagged the most consistent methodLN in terms of accuracy, and the NLs that considered meta-learning needs were not necessarily superior to those thatdid not.
5.2. Large Scale Few-Shot Classification Experiments
Next, we evaluate NLs on a demanding few-shot classi-fication challenge called Meta-Dataset, composed of thir-teen (eight train, five test) image classification datasets (Tri-antafillou et al., 2020). Experiments are carried out withCNAPS, which achieves state-of-the-art performance onMeta-Dataset (Requeima et al., 2019a) and ProtoNets. Thechallenge constructs few-shot learning tasks by drawingfrom the following distribution. First, one of the datasetsis sampled uniformly; second, the “way” and “shot” aresampled randomly according to a fixed procedure; third,the classes and context / target instances are sampled. As aresult, the context size Dτ will vary in the range between5 and 500 for each task. In the meta-test phase, the iden-tity of the original dataset is not revealed and tasks mustbe treated independently (i.e. no information can be trans-ferred between them). The meta-training set comprises adisjoint and dissimilar set of classes from those used formeta-test. Details provided in Appendix B and Triantafillouet al. (2020).
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Table 1. Accuracy results for different few-shot settings on Omniglot and miniImageNet using the MAML algorithm. All figures arepercentages and the ± sign indicates the 95% confidence interval. Bold indicates the highest scores. The numbers after the configurationname indicate the way and shots, respectively. The vertical lines enclose the transductive results. The TBN, examples, and class columnsindicate accuracy when tested with all target examples at once, one example at a time, and one class at a time, respectively. All other NLsare non-transductive and yield the same result when tested by example or class.
Configuration TBN example class CBN BRN LN IN RN MetaBN TaskNorm-L TaskNorm-I
Figure 4. (a) Accuracy vs shot for MAML on 5-way miniImagenet classification. (b) Plot of validation accuracy versus training iterationusing MAML for Omniglot 5-way, 5-shot corresponding to the results in Table 1. (c) Training Loss versus iteration corresponding to theresults using the CNAPS algorithm in Table 2. Note that TBN, CBN, and RN all share the same meta-training step.
Accuracy The classification accuracy results for CNAPSand ProtoNets on Meta-Dataset are shown in Table 2 andTable 3, respectively. In the case of ProtoNets, all the theNLs specifically designed for meta-learning scenarios out-perform TBN in terms of classification accuracy based ontheir average rank over all the datasets. For CNAPS, bothRN and TASKNORM-I meet or exceed the rank of TBN.This may be as |Dτ | (i) is quite large in Meta-Dataset, and(ii) may be imbalanced w.r.t. classes, making predictionharder with transductive NLs. TASKNORM-I comes outas the clear winner ranking first in 11 and 10 of the 13datasets using CNAPS and ProtoNets, respectively. Thissupports the hypothesis that augmenting the BN momentswith a second, instance based set of moments and learn-ing the blending factor α as a function of context set sizeis superior to fixing α to a constant value (as is the casewith RN). With both algorithms, the instance based NLs fallshort of the meta-learning specific ones. However, in thecase of CNAPS, they outperform the running average basedmethods (CBN, BRN), which perform poorly. In the caseof ProtoNets, BRN outperforms the instance based meth-ods, and IN fairs the worst of all. In general, ProtoNets isless sensitive to the NL used when compared to CNAPS.
The BASELINE column in Table 2 is taken from Requeimaet al. (2019a), where the method reported state-of-the-artresults on Meta-Dataset. The BASELINE algorithm uses therunning moments learned during pre-training of its featureextractor for normalization. Using meta-learning specificNLs (in particular TASKNORM) achieves significantly im-proved accuracy compared to BASELINE.
As an ablation, we have also added an additional variantof TASKNORM that blends the batch moments from thecontext set with the running moments accumulated dur-ing meta-training that we call TASKNORM-r. TASKNORM-r makes use of the global running moments to augmentthe local context statistic and it did not perform as well asthe TASKNORM variants that employed local moments (i.e.TASKNORM-I and TASKNORM-L).
Training Speed Figure 4c plots training loss versus train-ing iteration for the models in Table 2 that use the CNAPSalgorithm. The fastest training convergence is achieved byTASKNORM-I. The instance based methods (IN, LN) arethe slowest to converge. Note that TASKNORM convergeswithin 60k iterations while BASELINE takes 110k iterations
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Table 2. Few-shot classification results on META-DATASET using the CNAPS (top) and ProtoNets (bottom) algorithms. Meta-trainingperformed on datasets above the dashed line. Datasets below the dashed line are entirely held out. All figures are percentages and the ±sign indicates the 95% confidence interval over tasks. Bold indicates the highest scores. Vertical lines in the TBN column indicate thatthis method is transductive. Numbers in the BASELINE column are from (Requeima et al., 2019a).
Table 3. Few-shot classification results on META-DATASET using the Prototypical Networks algorithm. Datasets below the dashed lineare entirely held out. Meta-training performed on datasets above the dashed line. All figures are percentages and the ± sign indicatesthe 95% confidence interval over tasks. Bold indicates the highest scores. Vertical lines in the TBN column indicate that this method istransductive.
and IN takes 200k. Figure C.1b shows the training curvesfor the ProtoNets algorithm. The convergence speed trendsare very similar to CNAPS, with TASKNORM-I the fastest.
Our results demonstrate that TASKNORM is the best ap-proach for normalizing tasks on the large scale Meta-Datasetbenchmark in terms of classification accuracy and trainingefficiency. Here, we see high sensitivity of performanceacross NLs. Interestingly, in this setting TASKNORM-I out-performed TBN in classification accuracy, as did both RNand METABN. This refutes the hypothesis that TBN willalways outperform other methods due to its transductiveproperty, and implies that designing NL methods specifi-cally for meta-learning has significant value. In general, themeta-learning specific methods outperformed more generalNLs, supporting our third hypothesis. We suspect the reasonthat TASKNORM outperforms other methods is due to its
ability to adaptively leverage information from both Dτ andxτ∗ when computing moments, based on the size of Dτ .
6. ConclusionsWe have identified and specified several issues and chal-lenges with NLs for the meta-learning setting. We haveintroduced a novel variant of batch normalization – thatwe call TASKNORM – which is geared towards the meta-learning setting. Our experiments demonstrate that TAS-KNORM achieves performance gains in terms of both classi-fication accuracy and training speed, sometimes exceedingtransductive batch normalization. We recommend that fu-ture work in the few-shot / meta-learning community adoptTASKNORM, and if not, declare the form of normalizationused and implications thereof, especially where transductivemethods are applied.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
AcknowledgmentsThe authors would like to thank Elre Oldewage, Will Teb-butt, and the reviewers for their insightful comments andfeedback. Richard E. Turner is supported by Google, Ama-zon, ARM, Improbable and EPSRC grants EP/M0269571and EP/L000776/1.
ReferencesAbadi, M., Agarwal, A., Barham, P., Brevdo, E., Chen, Z.,
Citro, C., Corrado, G. S., Davis, A., Dean, J., Devin, M.,Ghemawat, S., Goodfellow, I., Harp, A., Irving, G., Isard,M., Jia, Y., Jozefowicz, R., Kaiser, L., Kudlur, M., Lev-enberg, J., Mané, D., Monga, R., Moore, S., Murray, D.,Olah, C., Schuster, M., Shlens, J., Steiner, B., Sutskever,I., Talwar, K., Tucker, P., Vanhoucke, V., Vasudevan,V., Viégas, F., Vinyals, O., Warden, P., Wattenberg, M.,Wicke, M., Yu, Y., and Zheng, X. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015.URL https://www.tensorflow.org/. Softwareavailable from tensorflow.org.
Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization.arXiv preprint arXiv:1607.06450, 2016.
Bakker, B. and Heskes, T. Task clustering and gating forBayesian multitask learning. Journal of Machine Learn-ing Research, 4:83–99, May 2003.
Chen, Y. A re-implementation of "prototypical networksfor few-shot learning". https://github.com/cyvius96/prototypical-network-pytorch,2018.
Finn, C., Abbeel, P., and Levine, S. Model-agnosticmeta-learning for fast adaptation of deep networks.In Precup, D. and Teh, Y. W. (eds.), Proceed-ings of the 34th International Conference on Ma-chine Learning, volume 70 of Proceedings of Ma-chine Learning Research, pp. 1126–1135, InternationalConvention Centre, Sydney, Australia, 06–11 Aug2017. PMLR. URL http://proceedings.mlr.press/v70/finn17a.html.
Finn, C. B. Code for "Model-agnostic meta-learning forfast adaptation of deep networks". https://github.com/cbfinn/maml, 2017.
Gordon, J., Bronskill, J., Bauer, M., Nowozin, S., andTurner, R. Meta-learning probabilistic inference forprediction. In International Conference on LearningRepresentations, 2019. URL https://openreview.net/forum?id=HkxStoC5F7.
Grant, E., Finn, C., Levine, S., Darrell, T., and Griffiths, T.Recasting gradient-based meta-learning as hierarchical
bayes. In International Conference on Learning Represen-tations, 2018. URL https://openreview.net/forum?id=BJ_UL-k0b.
Ha, D., Dai, A., and Le, Q. V. Hypernetworks. In In-ternational Conference on Learning Representations,2016. URL https://openreview.net/forum?id=rkpACe1lx.
Heskes, T. Empirical bayes for learning to learn. In Pro-ceedings of the Seventeenth International Conference onMachine Learning, ICML ’00, pp. 367–374, San Fran-cisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc.ISBN 1-55860-707-2. URL http://dl.acm.org/citation.cfm?id=645529.658133.
Hospedales, T., Antoniou, A., Micaelli, P., and Storkey,A. Meta-learning in neural networks: A survey. arXivpreprint arXiv:2004.05439, 2020.
Ioffe, S. Batch renormalization: Towards reducing mini-batch dependence in batch-normalized models. In Guyon,I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R.,Vishwanathan, S., and Garnett, R. (eds.), Advances inNeural Information Processing Systems 30, pp. 1945–1953. Curran Associates, Inc., 2017.
Ioffe, S. and Szegedy, C. Batch normalization: Accel-erating deep network training by reducing internal co-variate shift. In Bach, F. and Blei, D. (eds.), Proceed-ings of the 32nd International Conference on MachineLearning, volume 37 of Proceedings of Machine Learn-ing Research, pp. 448–456, Lille, France, 07–09 Jul2015. PMLR. URL http://proceedings.mlr.press/v37/ioffe15.html.
Jerfel, G., Grant, E., Griffiths, T., and Heller, K. A. Rec-onciling meta-learning and continual learning with on-line mixtures of tasks. In Wallach, H., Larochelle, H.,Beygelzimer, A., dÁlché-Buc, F., Fox, E., and Garnett, R.(eds.), Advances in Neural Information Processing Sys-tems 32, pp. 9119–9130. Curran Associates, Inc., 2019.
Krizhevsky, A. and Hinton, G. Learning multiple layersof features from tiny images. Technical report, Citeseer,2009.
Lake, B., Salakhutdinov, R., Gross, J., and Tenenbaum, J.One shot learning of simple visual concepts. In Proceed-ings of the annual meeting of the cognitive science society,volume 33, 2011.
LeCun, Y., Cortes, C., and Burges, C. MNIST hand-written digit database. AT&T Labs [Online]. Available:http://yann. lecun. com/exdb/mnist, 2:18, 2010.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Liu, Y., Lee, J., Park, M., Kim, S., Yang, E., Hwang,S., and Yang, Y. Learning to propagate labels: trans-ductive propagation network for few-shot learning. InInternational Conference on Learning Representations,2019. URL https://openreview.net/forum?id=SyVuRiC5K7.
Luo, C., Zhan, J., Xue, X., Wang, L., Ren, R., and Yang,Q. Cosine normalization: Using cosine similarity in-stead of dot product in neural networks. In InternationalConference on Artificial Neural Networks, pp. 382–391.Springer, 2018.
Mishra, N., Rohaninejad, M., Chen, X., and Abbeel,P. A simple neural attentive meta-learner. In In-ternational Conference on Learning Representations,2018. URL https://openreview.net/forum?id=B1DmUzWAW.
Nagabandi, A., Finn, C., and Levine, S. Deep online learn-ing via meta-learning: Continual adaptation for model-based RL. In International Conference on LearningRepresentations, 2019. URL https://openreview.net/forum?id=HyxAfnA5tm.
Nam, H. and Kim, H.-E. Batch-instance normalization foradaptively style-invariant neural networks. In Advances inNeural Information Processing Systems, pp. 2558–2567,2018.
Nichol, A., Achiam, J., and Schulman, J. Onfirst-order meta-learning algorithms. arXiv preprintarXiv:1803.02999, 2018.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J.,Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga,L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Rai-son, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang,L., Bai, J., and Chintala, S. Pytorch: An imperativestyle, high-performance deep learning library. In Wal-lach, H., Larochelle, H., Beygelzimer, A., d’ Alché-Buc,F., Fox, E., and Garnett, R. (eds.), Advances in Neural In-formation Processing Systems 32, pp. 8024–8035. CurranAssociates, Inc., 2019.
Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S.Meta-learning with implicit gradients. In Wallach, H.,Larochelle, H., Beygelzimer, A., d’ Alché-Buc, F., Fox,E., and Garnett, R. (eds.), Advances in Neural Informa-tion Processing Systems 32, pp. 113–124. Curran Asso-ciates, Inc., 2019.
Ravi, S. and Larochelle, H. Optimization as a model for few-shot learning. In Proceedings of the International Confer-ence on Learning Representations, 2017. URL https://openreview.net/pdf?id=rJY0-Kcll.
Ren, M., Ravi, S., Triantafillou, E., Snell, J., Swersky, K.,Tenenbaum, J. B., Larochelle, H., and Zemel, R. S. Meta-learning for semi-supervised few-shot classification. InInternational Conference on Learning Representations,2018. URL https://openreview.net/forum?id=HJcSzz-CZ.
Requeima, J., Gordon, J., Bronskill, J., Nowozin, S., andTurner, R. E. Fast and flexible multi-task classification us-ing conditional neural adaptive processes. In Wallach, H.,Larochelle, H., Beygelzimer, A., dÁlché-Buc, F., Fox, E.,and Garnett, R. (eds.), Advances in Neural InformationProcessing Systems 32, pp. 7957–7968. Curran Asso-ciates, Inc., 2019a.
Requeima, J., Gordon, J., Bronskill, J., Nowozin, S., andTurner, R. E. Code for "Fast and flexible multi-task clas-sification using conditional neural adaptive processes".https://github.com/cambridge-mlg/cnaps, 2019b.
Salimans, T. and Kingma, D. P. Weight normalization: Asimple reparameterization to accelerate training of deepneural networks. In Advances in neural information pro-cessing systems, pp. 901–909, 2016.
Schmidhuber, J. Evolutionary principles in self-referentiallearning, or on learning how to learn: the meta-meta-... hook. PhD thesis, Technische Universität München,1987.
Singh, S. and Krishnan, S. Filter response normalizationlayer: Eliminating batch dependence in the training ofdeep neural networks. arXiv preprint arXiv:1911.09737,2019.
Snell, J. Code for the nips 2017 paper "prototypicalnetworks for few-shot learning". https://github.com/jakesnell/prototypical-networks,2017.
Snell, J., Swersky, K., and Zemel, R. Prototypical networksfor few-shot learning. In Guyon, I., Luxburg, U. V., Ben-gio, S., Wallach, H., Fergus, R., Vishwanathan, S., andGarnett, R. (eds.), Advances in Neural Information Pro-cessing Systems 30, pp. 4077–4087. Curran Associates,Inc., 2017.
Thrun, S. and Pratt, L. Learning to learn. Springer Science& Business Media, 2012.
Triantafillou, E., Zhu, T., Dumoulin, V., Lamblin, P.,Xu, K., Goroshin, R., Gelada, C., Swersky, K., Man-zagol, P.-A., and Larochelle, H. Code for "Meta-dataset: A dataset of datasets for learning to learnfrom few examples". https://github.com/google-research/meta-dataset, 2019.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Triantafillou, E., Zhu, T., Dumoulin, V., Lamblin, P., Evci,U., Xu, K., Goroshin, R., Gelada, C., Swersky, K., Man-zagol, P.-A., and Larochelle, H. Meta-dataset: A datasetof datasets for learning to learn from few examples. InInternational Conference on Learning Representations,2020. URL https://openreview.net/forum?id=rkgAGAVKPr.
Ulyanov, D., Vedaldi, A., and Lempitsky, V. Instance nor-malization: The missing ingredient for fast stylization.arXiv preprint arXiv:1607.08022, 2016.
Vinyals, O., Blundell, C., Lillicrap, T., kavukcuoglu, k., andWierstra, D. Matching networks for one shot learning. InLee, D. D., Sugiyama, M., Luxburg, U. V., Guyon, I., andGarnett, R. (eds.), Advances in Neural Information Pro-cessing Systems 29, pp. 3630–3638. Curran Associates,Inc., 2016.
Wu, Y. and He, K. Group normalization. In Proceedings ofthe European Conference on Computer Vision (ECCV),pp. 3–19, September 2018.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
A. Additional Normalization LayersHere we discuss various additional NLs that are relevant tometa-learning.
A.1. Batch Renormalization (BRN)
Batch renormalization (BRN; Ioffe, 2017) is intended tomitigate the issue of non-identically distributed and/or smallbatches while retaining the training efficiency and stabilityof CBN. In BRN, the CBN algorithm is augmented with anaffine transform with batch-derived parameters which cor-rect for the batch statistics being different from the overallpopulation. The normalized activations of a BRN layer arecomputed as follows:
a′n = γ
(r
(an − µBNσBN + ε
)+ d
)+ β,
where
r =stop_grad
(clip[1/rmax,rmax]
(σBNσr
)),
d =stop_grad
(clip[−dmax,dmax]
(µBN − µr
σr
)).
Here stop_grad(·) denotes a gradient blocking operation,and clip[a,b] denotes an operation returning a value inthe range [a, b]. Like CBN, BRN is not well suited to themeta-learning scenario as it does not map directly to thehierarchical form of meta-learning models. In Section 5, weshow that using BRN can improve predictive performancecompared to CBN, but still performs significantly worsethan competitive approaches. Table 1 shows that batchrenormalization performs poorly when using MAML.
A.2. Group Normalization (GN)
A key insight of Wu & He (2018) is that CBN performancesuffers with small batch sizes. The goal of Group Normal-ization (GN; Wu & He, 2018) is thus to address the problemof normalization of small batch sizes, which, among othermatters, is crucial for training large models in a data-parallelfashion. This is achieved by dividing the image channelsinto a number of groups G and subsequently computing themoments for each group. GN is equivalent to LN when thereis only a single group (G = 1) and equivalent to IN whenthe number of groups is equal to the number of channels inthe layer (G = C).
A.3. Other NLs
There exist additional NLs including Weight Normaliza-tion (Salimans & Kingma, 2016), Cosine Normalization
(Luo et al., 2018), Filter Response Normalization (Singh &Krishnan, 2019), among many others.
Weight normalization reparameterizes weight vectors in aneural network to improve the conditioning for optimization.Weight normalization is non-transductive, but we don’t con-sider this approach further in this work as we focus on NLsthat modify activations as opposed to weights.
Filter Response Normalization (FRN) is another non-transductive NL that performs well for all batch sizes. How-ever we did not include it in our evaluation as FRN alsoencompasses the activation function as an essential part ofnormalization making it difficult to be a drop in replacementfor CBN in pre-trained networks as is the case for some ofour experiments.
Cosine normalization replaces the dot-product calculationin neural networks with cosine similarity for improved per-formance. We did not consider this method further in ourwork as it is not a simple drop-in replacement for CBN inpre-existing networks such as the ResNet-18 we use in ourexperiments.
B. Experimental DetailsIn this section, we provide the experimental details re-quired to reproduce our experiments. The experiments usingMAML (Finn et al., 2017) were implemented in TensorFlow(Abadi et al., 2015), the Prototypical Networks experimentswere implemented in Pytorch (Paszke et al., 2019), and theexperiments using CNAPS (Requeima et al., 2019a) wereimplemented using a combination of TensorFlow (Abadiet al., 2015) and Pytorch. All experiments were executed onNVIDIA Tesla P100-16GB GPUs.
B.1. MAML Experiments
We evaluate MAML using a range of normalization layerson:
1. Omniglot (Lake et al., 2011): a few-shot learningdataset consisting of 1623 handwritten characters (eachwith 20 instances) derived from 50 alphabets.
2. miniImageNet (Vinyals et al., 2016): a dataset of60,000 color images that is sub-divided into 100classes, each with 600 instances.
For all the MAML experiments, we used the codebase pro-vided by the MAML authors (Finn, 2017) with only smallmodifications to enable additional normalization techniques.Note that we used the first-order approximation versionof MAML for all experiments. MAML was invoked withthe command lines as specified in the main.py file inthe MAML codebase. No hyper-parameter tuning was per-formed and we took the results from a single run. All models
TASKNORM: Rethinking Batch Normalization for Meta-Learning
were trained for 60,000 iterations and then tested. No earlystopping was used. We did not select the model based on val-idation accuracy or other criteria. The MAML code employsten gradient steps at test time and computes classificationaccuracy after each step. We report the maximum accuracyacross those ten steps. To generate the plot in Figure 4a, weuse the same command line as Omniglot-5-1, but vary theupdate batch size from one to ten.
B.2. CNAPS Experiments
We evaluate CNAPS using a range of normalization lay-ers on a demanding few-shot classification challenge calledMeta-Dataset (Triantafillou et al., 2020). Meta-Dataset iscomposed of ten (eight train, two test) image classificationdatasets. We augment Meta-Dataset with three additionalheld-out datasets: MNIST (LeCun et al., 2010), CIFAR10(Krizhevsky & Hinton, 2009), and CIFAR100 (Krizhevsky& Hinton, 2009). The challenge constructs few-shot learn-ing tasks by drawing from the following distribution. First,one of the datasets is sampled uniformly; second, the “way”and “shot” are sampled randomly according to a fixed pro-cedure; third, the classes and context / target instances aresampled. Where a hierarchical structure exists in the data(ILSVRC or OMNIGLOT), task-sampling respects the hier-archy. In the meta-test phase, the identity of the originaldataset is not revealed and the tasks must be treated inde-pendently (i.e. no information can be transferred betweenthem). Notably, the meta-training set comprises a disjointand dissimilar set of classes from those used for meta-test.Full details are available in Triantafillou et al. (2020).
For all the CNAPS experiments, we use the code providedby the the CNAPS authors (Requeima et al., 2019b) withonly small modifications to enable additional normalizationtechniques. We follow an identical dataset configuration andtraining process as prescribed in Requeima et al. (2019b). Togenerate results in Table 2, we used the following CNAPSoptions: FiLM feature adaptation, a learning rate of 0.001,and TBN, CBN, BRN, and RN used 70,000 training itera-tions, IN used 200,000 iterations, LN used 110,000 itera-tions, and TASKNORM used 60,000 iterations. The CNAPScode generates two models: fully trained and best validation.We report the better of the two. We performed no hyper-parameter tuning and report the test results from the firstrun. Note that CBN, TBN, and RN share the same trainedmodel. They differ only in how meta-testing is done.
B.3. Prototypical Networks Experiments
We evaluate the Prototypical Networks (Snell et al., 2017)algorithm with a range of NLs using the same Omniglot,miniImageNet, and Meta-Dataset benchmarks.
For Omniglot, we used the codebase created by the Pro-totypical Networks authors (Snell, 2017). For miniIma-
geNet, we used the a different codebase ((Chen, 2018)) asthe first codebase did not support miniImageNet. Only smallmodifications were made to the two codebases to enableadditional NLs. For Omniglot and miniImageNet, we sethyper-parameters as prescribed in (Snell et al., 2017). Earlystopping was employed and the model that produced thebest validation was used for testing.
For Meta-Dataset, we use the code provided by the theCNAPS authors (Requeima et al., 2019b) with only smallmodifications to enable additional normalization techniquesand a new classifier adaptation layer to generate the linearclassifier weights per equation (8) in (Snell et al., 2017).We follow an identical dataset configuration and trainingprocess as prescribed in Requeima et al. (2019b). To gen-erate results in Table 3, we used the following CNAPSoptions: no feature adaptation, a learning rate of 0.001,60,000 training iterations for all NLs, and the pretrainedfeature extractor weights were not frozen and allowed toupdate during meta-training.
C. Additional Classification ResultsTable C.1 shows the classification accuracy results for theProtoNets algorithm on the Omniglot and miniImageNetdatasets. Figure C.1a and Figure C.1b show the trainingcurves for the ProtoNets algorithm on Omniglot and Meta-Dataset, respectively.
D. Additional Transduction TestsA non-transductive meta-learning system makes predictionsfor a single test set label conditioned only on a single inputand the context set. A transductive meta-learning systemalso conditions on additional samples from the test set.
Table D.2 demonstrates failure modes for transductive learn-ing. In addition to reporting the classification accuracy re-sults when the target set is evaluated all at once (first columnof results for each NL), we report the classification accuracywhen meta-testing is performed one target-set example ata time (second column of results for each NL), and onetarget-set class at a time (third column of results for eachNL). Table D.2 demonstrates that classification accuracydrops dramatically for TBN when testing is performed oneexample or one class at a time.
Importantly, in the case of TASKNORM-I (or any non-transductive NL – i.e. all of NLs evaluated in this work apartfrom TBN), the evaluation results are identical whether theyare meta-tested on the entire target set at once, one exampleat a time, or one class at a time. This shows that transductivelearning is sensitive to the distribution over the target setused during meta-training, demonstrating that transductivelearning is less generally applicable than non-transductive
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Table C.1. Accuracy results for different few-shot settings on Omniglot and miniImageNet using the Prototypical Networks algorithm. Allfigures are percentages and the ± sign indicates the 95% confidence interval. Bold indicates the highest scores. The numbers after theconfiguration name indicate the way and shots, respectively. The vertical lines in the TBN column indicate that this method is transductive.
Configuration TBN CBN BRN LN IN RN MetaBN TaskNorm-L TaskNorm-I
Figure C.1. (a) Plot of validation accuracy versus training iteration using ProtoNets for Omniglot 20-way, 1-shot corresponding to theresults in Table C.1. (b) Training Loss versus iteration corresponding to the results using the ProtoNets algorithm on META-DATASET inTable 3. Note that TBN, CBN, and RN all share the same meta-training step.
learning. In particular, transductive learners may fail tomake good predictions if target sets contains a differentclass balance than what was observed during meta-training,or if they are required to make predictions for one exampleat a time (e.g. in streaming applications).
E. Ablation Study: Choosing the bestparameterization for α
There are a number of possibilities for the parameterizationof the TASKNORM blending parameter α. We consider fourdifferent configurations for each NL:
1. α is learned separately for each channel (i.e. channelspecific) as an independent parameter.
2. α is learned shared across all channels as an indepen-dent parameter.
3. α is learned separately for each channel (i.e. channel
specific) as a function of context set size (i.e. α =SIGMOID(SCALE|Dτ |+ OFFSET)).
4. α is learned shared across all channels as a functionof context set size (i.e. α = SIGMOID(SCALE|Dτ | +OFFSET)).
Accuracy Table E.3 and Table E.4 show classificationaccuracy for the various parameterizations for MAML andthe CNAPS algorithms, respectively using the TASKNORM-I NL.
When using the MAML algorithm, there are only two op-tions to evaluate as the context size is fixed for each con-figuration of dataset, shot, and way and thus we need onlyevaluate the independent options (1 and 2 above). Table E.3indicates that the classification accuracy for the channel spe-cific and shared parameterizations are nearly identical, butthe shared parameterization is better in the Omniglot-5-1benchmark and hence has the best ranking overall.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Table D.2. Few-shot classification results for TBN and TASKNORM-I on META-DATASET using the CNAPS algorithm. For each NL, thefirst column of results "All" reports accuracy when meta-testing is performed on the entire target set at once. The second column of results"Example" reports accuracy when meta-testing is performed one example at a time. The third column of results "Class" reports accuracywhen meta-testing is performed one class at a time. All figures are percentages and the ± sign indicates the 95% confidence interval overtasks. Meta-training is performed on datasets above the dashed line, while datasets below the dashed line are entirely held out.
TBN TASKNORM-IDataset All Example Class All Example Class
Table E.3. Few-shot classification results for two α parameteriza-tions on Omniglot and miniImageNet using the MAML algorithm.All figures are percentages and the ± sign indicates the 95% confi-dence interval over tasks. Bold indicates the highest scores.
When using the CNAPS algorithm on the Meta-Datasetbenchmark, the best parameterization option in terms ofclassification accuracy is α shared across channels as a func-tion of context size. One justification for having α be afunction of context size can be seen in Figure 3b. Herewe plot the line SCALE|Dτ |+ OFFSET on a linear scalefor a representative set of NLs in the ResNet-18 used inthe CNAPS algorithm. The algorithm has learned that theSCALE parameter is non-zero and the OFFSET is almostzero in all cases. If a constant α would lead to better accu-racy, we would see the opposite (i.e the SCALE parameterwould be at or near zero and the OFFSET parameter beingsome non-zero value). From Table E.4 we can also see thataccuracy is better when the parameterization is a shared αopposed to having a channel-specific α.
Training Speed Figure E.2a and Figure E.2b show thelearning curves for the various parameterization optionsusing the MAML and the CNAPS algorithms, respectivelywith a TASKNORM-I NL.
For the MAML algorithm the training efficiency of theshared and channel specific parameterizations are almostidentical. For the CNAPS algorithm, Figure E.2b indicatesthe training efficiency of the independent parameterizationis considerably worse than the functional one. The twofunctional representations for the CNAPs algorithm havealmost identical training curves. Based on Figure E.2aand Figure E.2b, we conclude that the training speed ofthe functional parameterization is superior to that of theindependent parameterization and that there is little or nodifference in the training speeds between the functional,shared parameterization and the functional, channel specificparameterization.
In summary, the best parameterization for α when it islearned shared across channels as a function of context setsize (option 4, above). We use this parameterization in all ofthe CNAPS experiments in the main paper. For the MAMLexperiments, the functional parameterization is meaninglessgiven that all the test configurations have a fixed contextsize. In that case, we used the independent, shared acrosschannels parameterization for α for the experiments in themain paper.
TASKNORM: Rethinking Batch Normalization for Meta-Learning
Table E.4. Few-shot classification results for various α parameterizations on META-DATASET using the CNAPS algorithm. All figures arepercentages and the ± sign indicates the 95% confidence interval over tasks. Bold indicates the highest scores. Meta-training performedon datasets above the dashed line, while datasets below the dashed line are entirely held out.
Independent FunctionalDataset Channel Specific Shared Channel Specific Shared
Figure E.2. (a) Plots of validation accuracy versus training iteration corresponding to the parameterization experiments using the MAMLalgorithm in Table E.3. (b) Plot of training loss versus iteration corresponding to the parameterization experiments using the CNAPSalgorithm in Table E.4.
