Memory Efficient Meta-Learning with Large Images

1 Introduction . Meta-learning approaches to few-shot classification are very computationally efficient . Once metatrained , they can learn a new task at test time with as few as 1-5 optimization steps [ 1 , 2 ] or a single forward pass through the model [ 3–5 ] and with minimal or no hyper-parameter tuning . In contrast , transfer learning approaches based on fine-tuning typically rely on a large pre-trained feature extractor , and instead take 100s-1000s of optimization steps at test time in order to learn a task [ 6 ] , thus incurring a high computational cost for each new task encountered . This makes meta-learned solutions attractive in compute-constrained deployments , or scenarios where the model must learn multiple different tasks or update on-the-fly ( e.g . in continual and online learning settings [ 7–10 ] ) . However , a crucial barrier to progress is that meta-learning approaches are memory-intensive to train and thus can not easily leverage large images for a performance boost , as recent fine-tuning approaches have done . This limitation arises because a meta-learner must back-propagate through all ∗Authors contributed equally 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the examples in a task ’ s support set ( i.e . the task ’ s training set ) that contribute to a prediction on a query example . In some cases , this can be as many as 1000 images [ 11 ] . As a result , the amount of memory required for the computational graph grows linearly with the number of support images , and quadratically with their dimension . In contrast , transfer learning approaches can employ standard batch processing techniques to scale to larger images when under memory constraints — a feature which has contributed significantly to their recent success on few-shot benchmarks [ 11 ] . Current solutions for training meta-learners on large images include 1 ) parallelizing the model across multiple GPUs , which may not be available or convenient , 2 ) considering tasks with fewer support images , and 3 ) employing gradient/activation checkpointing methods [ 12 ] which incur longer training times and still fall short of the task sizes required in key benchmarks . Instead , most existing work [ 1 , 3 , 2 , 13 , 4 , 14 ] has opted for training on large tasks but small images which translates poorly into real-world applications and limits competitiveness on few-shot benchmarks . In this work , we improve on these alternatives by proposing LITE , a Large Image and Task Episodic training scheme for meta-learning models that enables training with large images and large tasks , on a single GPU . We achieve this through the simple observation that meta-learners typically aggregate a task ’ s support examples using a permutation invariant sum operation . This structure ensures invariance to the ordering of the support set . Consequently , the gradients for a task can be decomposed as a sum of gradient contributions from the task ’ s support examples . This enables us to perform a forward pass on a task ’ s entire support set , but realize significant savings in memory by back-propagating only a random subset of these images which we show is an unbiased approximation of the true gradient . Fig . 1 illustrates the key trade-offs , with LITE enabling meta-learners to benefit from large images for improved classification accuracy but still remain computationally efficient at test time . We use LITE to train key meta-learning methods and show that our best performing instantiation – Simple CNAPs [ 5 ] with LITE – achieves state-of-the-art results relative to all meta-learners on two challenging few-shot benchmarks : VTAB+MD [ 11 ] , an extensive suite of both meta-learning and transfer learning tasks , and ORBIT [ 14 ] , an object recognition benchmark of high-variation real-world videos . Our results showcase the unique advantage of meta-learning methods – that when properly trained they can be competitive with transfer learning approaches in terms of accuracy for a fraction of the computational cost at test time — and they serve as a counterpoint to the recent narrative that transfer learning is all you need for few-shot classification . Our contributions 1 . LITE , a general and memory-efficient episodic training scheme which enables meta-learning models to be trained on large images and large tasks on a single GPU . 2 . A mathematical justification for approximating the true gradient with a random subset of a task ’ s support examples which applies to common classes of meta-learning methods . 3 . Instantiations of LITE on key classes of meta-learners to demonstrate its versatility . 4 . State-of-the-art performance using Simple CNAPs with LITE compared to other leading meta-learners on two challenging few-shot benchmarks , VTAB+MD and ORBIT 2 2 Why Meta-Learning with Large Images and Tasks is Difficult . Meta-learning preliminaries In few-shot image classification , the goal is to recognize new classes when given only a few training ( or support ) images of each class . Meta-learners typically achieve this through episodic training [ 15 ] . Here , an episode or task τ contains a support setDτS = { ( xτn , yτn ) } Nτ n=1 and a query set DτQ = { ( xτ∗m , yτ∗m ) } Mτ m=1 , where ( x , y ) is an image-label pair , Nτ is the number of ( labeled ) support elements given to learn the new classes , and Mτ is the number of query elements requiring predictions . Note that in a given task , elements inDτQ are drawn from the same set of classes as the elements in DτS . For brevity we may use the shorthand DS = { x , y } and DQ = { x∗ , y∗ } . During meta-training , a meta-learner is exposed to a large number of training tasks { τ } . For each task τ , the meta-learner takes as input the support setDS and outputs the parameters of a classifier that has been adapted to the current task θτ = θφ ( DS ) . The classifier can now make task-specific probabilistic predictions f ( x∗ , θτ = θφ ( DS ) ) for any query input x∗ ( see Fig . 2 ) . A function L ( y∗ , f ( x∗ , θτ ) ) computes the loss between the adapted classifier ’ s predictions for the query input and the true label y∗ which is observed during meta-training . Assuming that L , f , and θφ ( DS ) are differentiable , the meta-learner can then be trained with stochastic gradient descent by back-propagating the loss and updating the parameters φ . At meta-test time , the trained meta-learner is given a set of unseen test tasks , which typically contain classes that have not been seen during meta-training . For each task , the meta-learner is given its support set DS , and is then evaluated on its predictions for all the query inputs x∗ ( Fig . 2 , left ) . Large memory requirements for meta-training The primary bottleneck to using large images ( i.e . ≥ 224× 224 pixels ) in meta-learning approaches is the large amount of ( GPU ) memory required to process a task ’ s support set DS during meta-training . Specifically , the meta-learner θφ ( DS ) must perform a forward pass with a task ’ s entire support set before it can back-propagate the loss for query elements ( x∗ , y ) ∈ DQ ( and release the computation graph ) , thus preventing the use of conventional batch processing . The amount of memory required scales linearly with the number of support images Nτ and quadratically with their dimensions . If Nτ is large ( e.g . the recent VTAB+MD benchmark [ 11 ] requires a task ’ s support set to be as large as 1000 images ) , memory on a single GPU is thus quickly exceeded for large images . Note , the number of query elements in the task Mτ is not a bottleneck when using large images as the loss decomposes over elements of the query set DQ and is therefore amenable to mini-batching . By contrast , as the classifier itself is a non-linear function of the support set , the loss does not decompose and so it is not obvious how to apply similar ideas to allow scaling of DS in a principled way . Current ad hoc solutions to this problem are : ( i ) parallelize the meta-learner across multiple GPUs which may not be convenient or available and can involve significant engineering effort ; ( ii ) train on tasks with smaller ( or sub-sampled ) support sets which may adversely affect performance on test tasks with more classes and/or large numbers of samples per class ; ( iii ) train on tasks with smaller images ( e.g . 84× 84 pixels in miniImageNet [ 16 ] ) which limits performance and translates poorly to many real-world applications ; or ( iv ) trade memory usage for additional computation [ 12 ] by 2Source code for ORBIT experiments is available at https : //github.com/microsoft/ORBIT-Dataset and for the VTAB+MD experiments at https : //github.com/cambridge-mlg/LITE . Algorithm 1 LITE for a meta-training task τ Require : DS : task support set ; DQ : task query set ; N : number of support examples in DS ; H : number of elements in DS to back-propagate ; M : number of query examples in DQ ; Mb : batch size for DQ ; backward ( ) ≡ function to back-propagate a loss ; step ( ) ≡ function to update parameters with a gradient step . 1 : B ← ceil ( M/Mb ) . number of query batches 2 : for all b ∈ 1 , . . . , B do 3 : DQb ← { x∗m , y∗m } Mb m=1 . get query batch from DQ 4 : H ← { ( xnh , ynh ) } Hh=1 where { nh } Hh=1 ∼ U ( 1 , N ) .H to back-propagate 5 : H ← DS ∩H .H to not back-propagate 6 : DSb ← H∪H 7 : θτ ← θ ( DSb ) 8 : Lb ← 1MbΣ Mb m=1L ( y∗m , f ( x∗m , θτ ) ) . get loss of query batch 9 : backward ( Lb ) . back-propagate loss on query batch 10 : end for 11 : φ← step ( φ , N/H ) . update φ using weighting factor N/H employing activation/gradient checkpointing ( i.e . during training store only a subset of intermediate activations in a network needed for backpropagation and recompute the rest with additional forward computations when needed ) which allows for training on larger tasks at the expense of training time , but still falls well short of the memory needed to accommodate the task sizes required for key benchmarks ( e.g . VTAB+MD ) . Although training meta-learners has large memory requirements , meta-testing is generally memory efficient , requiring only a small number of gradient operations , or none at all , compared to transfer learning approaches that would perform large numbers of gradient-based updates at test time . 3 Large Image and Task Episodic ( LITE ) training . In this section , we introduce our general and memory-efficient solution for training meta-learners episodically on tasks with large support sets and large images . We call our approach Large Image and Task Episodic training or LITE . In Section 3.1 , we describe how LITE can be applied to key classes of meta-learners . Approach The fundamental idea underlying LITE is to perform a forward pass using the entire support set DS , but to compute the gradient contribution on only a small random subset of the examples in the support set . By doing this , we realize large savings in memory that includes gradients , activations , and the computation graph for the elements of DS that are not back-propagated . This is an approximation of the true gradient that would result if back-propagation was performed on all of the examples in DS . In the following section , we show that this approximation is an unbiased estimate of the true gradient . The approach for a general meta-learner is detailed in Algorithm 1 and shown diagrammatically in Fig . 2 . Mathematical justification The parameters of the meta-learner φ are found by minimizing the expected loss over all tasks . argmin φ T∑ τ=1 Mτ∑ m=1 L ( yτ∗m , f ( xτ∗m , θφ ( DτS ) ) ) . ( 1 ) In most meta-learning approaches , the support set enters into the loss through a sum over the N individual contributions from each data point it contains . This structure enables the meta-learners to be invariant to the ordering of the support set and allows all members of the support set to contribute to the adapted parameters ( unlike alternative permutation invariant operators like max or min ) . Below we show in blue how this sum arises in popular brands of meta-learners . In amortization methods ( e.g . CNAPS [ 4 ] and VERSA [ 17 ] ) , the aggregation of support set points is built in directly via a deep set encoder eφ1 ( · ) . This encodes the support set into an embedding vector which is mapped to the classifier parameters by a hyper-network tφ0 ( · ) . θφ ( DS ) = tφ0 ( N∑ n=1 eφ1 ( xn , yn ) ) . ( 2 ) In gradient-based methods ( e.g . MAML [ 1 ] ) , the classifier parameters are adapted from an initial value φ0 using a sum of derivatives of an inner-loop loss computed for each data point in the support set . The derivatives play the role of the deep set encoder in amortization methods . θφ ( DS ) = φ0 + φ1 N∑ n=1 d dφ Linner ( yn , f ( xn , φ ) ) ∣∣∣ φ=φ0 ( 3 ) Metric-based methods ( e.g . ProtoNets [ 3 ] ) comprise a body formed of a feature extractor and a head formed from a distance-based classifier . The classifier ’ s body parameters are not adapted in a task specific way θ ( body ) φ , c ( DS ) = φ0 . The classifier ’ s head is adapted by averaging the activations for each class in the support set to form prototypes . Letting kc denote the number of support examples of class c , the adapted head parameters are given by θ ( head ) φ , c ( DS ) = 1 kc kc∑ i=1 f ( x ( c ) i , φ0 ) = 1 kc N∑ n=1 1 ( yn = c ) f ( xn , φ0 ) . ( 4 ) Query points can then be classified using their distance from these prototypes d ( f ( x∗n , φ0 ) , θ ( head ) φ , c ) . We have established that in many meta-learners , each support set affects the classifier parameters and therefore the loss through a sum of contributions from each of its elements . We now focus on the consequences of this structure on the gradients of the loss with respect to the meta-learner ’ s parameters . To reduce clutter , we consider the contribution from just a single query point from a single task and suppress the dependence of the loss on the classifier and the query data point , writing L ( y∗ , f ( x∗ , θφ ( DS ) ) ) = L ( eφ ( DS ) ) where eφ ( DS ) = N∑ n=1 eφ ( xn , yn ) = N∑ n=1 e ( n ) φ . ( 5 ) As a consequence of the summation , the derivative of the loss is given by d dφ L ( eφ ( DS ) ) = L′ ( eφ ( DS ) ) × ( N∑ n=1 de ( n ) φ dφ ) where L′ ( eφ ( DS ) ) = dL ( e ) ) de ∣∣∣ e=eφ ( DS ) ( 6 ) which is a product of the sensitivity of the loss to the encoding of the data points and the sensitivity of the contribution to the encoding from each data point w.r.t . the meta-learner ’ s parameters . This second term is the source of the memory overhead when training meta-learners , but importantly , it can be rewritten as an expectation w.r.t . a uniform distribution over the support set data-point indices , d dφ L ( eφ ( DS ) ) = N L′ ( eφ ( DS ) ) En∼U ( 1 , N ) [ de ( n ) φ dφ ] . ( 7 ) We can now define the LITE estimator of the loss-derivative by approximating the expectation by Monte Carlo sampling H times , d dφ L ( eφ ( DS ) ) ≈ N H L′ ( eφ ( DS ) ) H∑ h=1 de ( nh ) φ dφ = d dφ L̂ ( eφ ( DS ) ) where { nh } Hh=1 ∼ U ( 1 , N ) . ( 8 ) This estimator is unbiased , converging to the true gradient as H →∞ . The estimator does not simply involve subsampling of the support set – parts of it depend on all the support set data points DS – and this is essential for it to be unbiased . The expectation and variance of this estimator are E { nh } ∼U ( 1 , N ) [ dL̂ dφ ] = dL dφ and V { nh } ∼U ( 1 , N ) [ dL̂ dφ ] = N2 H ( L′ ) 2 V { nh } ∼U ( 1 , N ) [ de ( nh ) φ dφ ] . In Section 5.3 , we empirically show that the LITE gradient estimate is unbiased and that its standard deviation is smaller than that of the naive estimator formed by sub-sampling the full support set . LITE provides memory savings by subsampling H examples from the support set , with H < N , and back-propagating only them . Crucially , a forward pass is still performed with the complementary set of points , with cardinality N −H , but these are not back-propagated .

Summary: The authors propose a method that enables few-shot learners to be trained on 1 GPU without memory issues. This is very important since most few-shot learners need to process the whole train set of a task before evaluating the loss on the task. They show mathematically that their method gives an unbiased estimation of the true gradient (if all samples in training were back-propagated to update the model). They evaluate their method on CNAPS and ProtoNets and show that it can achieve better results since it can be trained with larger images.