A Universal Representation Transformer Layer for Few-Shot Image Classification

1 INTRODUCTION . Learning tasks from small data remains a challenge for machine learning systems , which show a noticeable gap compared to the ability of humans to understand new concepts from few examples . A promising direction to address this challenge is developing methods that are capable of performing transfer learning across the collective data of many tasks . Since machine learning systems generally improve with the availability of more data , a natural assumption is that few-shot learning systems should benefit from leveraging data across many different tasks and domains—even if each individual task has limited training data available . This research direction is well captured by the problem of multi-domain few-shot classification . In this setting , training and test data spans a number of different domains , each represented by a different source dataset . A successful approach in this multi-domain setting must not only address the regular challenge of few-shot classification—i.e. , the challenge of having only a handful of examples per class . It must also discover how to leverage ( or ignore ) what is learned from different domains , achieving generalization and avoiding cross-domain interference . Recently , Triantafillou et al . ( 2020 ) proposed a benchmark for multi-domain few-shot classification , Meta-Dataset , and highlighted some of the challenges that current methods face when training data is heterogeneous . Crucially , they found that methods which trained on all available domains would normally obtain improved performance on some domains at the expense of others . Following on their work , progress has been made , which includes the design of adapted hyper-parameter optimization strategies ( Saikia et al. , 2020 ) and more flexible meta-learning algorithms ( Requeima et al. , 2019 ) . Most notable is SUR ( Selecting Universal Representation ) ( Dvornik et al. , 2020 ) , a method that relies on a so-called universal representation , extracting from a collection of pre-trained and domain-specific neural network backbones . SUR prescribes a hand-crafted feature-selection procedure to infer how to weight each backbone for each task at hand , and produces an adapted representation for each task . This was shown to lead to some of the best performances on Meta-Dataset . ∗This work was done while Lu Liu was a research intern with Mila . †Canada CIFAR AI Chair In SUR , the classification procedure for each task is fixed and not learned . Thus , except for the underlying universal representation , there is no transfer learning performed with regards to how classification rules are inferred across tasks and domains . Yet , cross-domain generalization might be beneficial in that area as well , in particular when tasks have only few examples per class . Present work . To explore this question , we propose a Universal Representation Transformer ( URT ) layer , which can effectively learn to transform a universal representation into task-adapted representations . The URT layer is inspired from Transformer ( Vaswani et al. , 2017 ) and uses an attention mechanism to learn to retrieve or blend the appropriate backbones to use for each task . By training this layer across few-shot tasks from many domains , it can support transfer across these tasks . We show that our URT layer on top of a universal representation ’ s pre-trained backbones sets a new state-of-the-art performance on Meta-Dataset . It succeeds at outperforming SUR on 4 dataset sources without impairing accuracy on the others . This leads to top performance on 7 dataset sources when comparing to a set of competing methods . To interpret the strategy that URT learns to weigh the backbones from different domains , we visualize the attention scores for both seen and unseen domains and find that our model generates meaningful weights for the pre-trained domains . A comprehensive analysis on variants and ablations of the URT layer is provided to show the importance of various components of URT , notably the number of attention heads . 2 FEW-SHOT CLASSIFICATION . 2.1 PROBLEM SETTING . In this section , we will introduce the problem setting for few-shot classification and the formulation of meta-learning for few-shot classification . Few-shot classification aims to classify samples where only few examples are available for each class . We describe a few-shot learning classification task as the pair of examples , comprising of a support set S to define the classification task and the query set Q of samples to be classified . Meta-learning is a technique that aims to model the problem of few-shot classification as learning to learn from instances of few-shot classification tasks . The most popular way to train a meta-learning model is with episodic training . Here , tasks T = ( Q , S ) are sampled from a larger dataset by taking subsets of the dataset to build a support set S and a query set Q for the task . A common approach is to sample N -way-K-shot tasks , each time selecting a random subset of N classes from the original dataset and choosing only K examples for each class to add to the support set S. The meta-learning problem can then be formulated by the following optimization : min Θ E ( S , Q ) ∼p ( T ) [ L ( S , Q , Θ ) ] , L ( S , Q , Θ ) = 1 |Q| ∑ ( x , y ) ∼Q − log p ( y|x , S ; Θ ) + λΩ ( Θ ) , ( 1 ) where p ( T ) is the distribution of tasks , Θ are the parameters of the model and p ( y|x , S ; Θ ) is the probability assigned by the model to label y of query example x ( given the support set S ) , and Ω ( Θ ) is an optional regularization term on the model parameters with factor λ . Conventional few-shot classification targets the setting of N -way-K-shot , where the number of classes and examples are fixed in each episode . Popular benchmarks following this approach include Omniglot ( Lake et al. , 2015 ) or benchmarks made of subsets of ImageNet , such as miniImageNet ( Vinyals et al. , 2016 ) and tieredImageNet ( Ren et al. , 2018 ) . In such benchmarks , the tasks for training cover a set of classes that is disjoint from the classes in the test set of tasks . However , with the training and test sets tasks coming from a single dataset/domain , the distribution of tasks found in either sets is similar and lacks variability , which may be unrealistic in practice . It is in this context that Triantafillou et al . ( 2020 ) proposed Meta-Dataset , as a further step towards large-scale , multi-domain few shot classification . Meta-Dataset includes ten datasets ( domains ) , with eight of them available for training . Additionally , each task sampled in the benchmark varies in the number of classes N , with each class also varying in the number of shots K. As in all few-shot learning benchmarks , the classes used for training and testing do not overlap . 2.2 BACKGROUND AND RELATED WORK . Meta-Learning A promising approach for few-shot classification is to use meta-learning to more directly train a model to learn to perform few-shot classification , in an end-to-end way . The two most popular methods are Prototypical Networks ( Snell et al. , 2017 ) and Model Agnostic Meta-Learning ( MAML ) ( Finn et al. , 2017 ) . Triantafillou et al . ( 2020 ) showed that prototypical networks and MAML could be combined by leveraging prototypes for the initialization of the output weights value in the inner loop . Requeima et al . ( 2019 ) also proposed Conditional Neural Adaptive Processes ( CNAPs ) for few-shot classification , which can be seen as extending prototypical networks with a more sophisticated architecture that allows for improved task adaptation . This architecture was later improved further by Bateni et al . ( 2020 ) with Simple CNAPS , leading to one of the current best methods on Meta-Dataset . Another line of work which leverages the idea of “ transfer by fine-tuning ” can be found in Appendix B . Universal Representations In contrast , our work instead builds on that of Dvornik et al . ( 2020 ) and their method SUR ( Selecting from Universal Representations ) . Bilen & Vedaldi ( 2017 ) introduced the term universal representation to refer to a representation that supports good performance in multiple domains . One proposal towards such a representation is to train different neural networks backbones separately on the data of each available domain , then simply to concatenate the representation learned by each . Another is to introduce some parameter sharing between the backbones , by having a single network conditioned on the domain of the provenance of each batch of training data ( Rebuffi et al. , 2018 ) , e.g . using Feature-wise Linear Modulate ( FiLM ) ( Perez et al. , 2018 ) . SUR proposes to leverage a universal representation in few-shot learning tasks with a feature selection procedure that assigns different weights to each of the domain-specific subvectors of the universal representation . The objective is to assign high weights only to the domain-specific representations that are specifically useful for each few-shot task at hand . The weights are inferred by optimizing a loss on the support set that encourages high accuracy of a nearest-centroid classifier . As such , the method does not involve any meta-learning—a choice motivated by the concern that meta-learning may struggle in generalizing to domains that are dissimilar to the training domains . SUR achieved some of the best performances on Meta-Dataset . However , a contribution of our work is to provide evidence that meta-learning can actually be used to replace SUR ’ s hand-designed inference procedure and improve performance further . Task Adaptive Representations Another line of work tries to retrieve task adaptive representations for each task . Task specific representations can be conditioned on a representation of the current task ( Oreshkin et al. , 2018 ; Wang et al. , 2019 ) , projected to another space ( Yoon et al. , 2019 ) , or masked based on inter-class commonality and inter-class uniqueness ( Li et al. , 2019 ) . While the representation extracted from URT is also task adaptive , it is adaptive to a set of pretrained backbones and can be applied to more complicated multi-domain scenarios . Wang & Hebert ( 2016 ) proposed to improve a CNN by adding extra layers and train it using unsupervised data while our contribution mainly lies in composing representations instead of an improved CNN . Alet et al . ( 2018 ) introduced a modular meta-learning method , which learns a repertoire of modules that serves as nodes to construct a tree structure to solve a new robotic-related task . Comparatively , URT is a one-for-all layer which doesn ’ t need to construct different module structures for each task . Transformer Networks Our meta-learning approach to leverage universal representations is inspired directly from Transformer networks ( Vaswani et al. , 2017 ) . Our model structure is inspired by the structure of the dot-product self-attention in the Transformer , which we adapted here to multidomain few-shot learning by designing appropriate parametrizations for queries , keys and values . Self-attention was explored in the single-domain training regime by Ye et al . ( 2020 ) ; Liu et al . ( 2019b ; a ; 2020 ) , however for a different purpose , where each representation of individual examples in a task support set is influenced by all other examples . Rather than using self-attention between individual examples in the support set , our model uses self-attention to select between different domain-specific backbones .

Summary: The paper addresses the problem of multi-domain few-shot image classification (where unseen classes and examples come from diverse data sources), and proposes a Universal Representation Transformer (URT) layer, which learns to transform a universal representation into task-adapted representations. The method proposed builds on top of SUR [Dvornik et al 2020], where a universal representation is extracted from the outputs of a collection of pre-trained and domain-specific backbones and a selection procedure infers how to weight each backbone for a given task at hand. While SUR inferred those weights by optimising a loss on the support set (the few examples provided in a task), the authors in this paper introduce an attention-based layer (inspired by Vaswani et al Transformer) that learns to weight the appropriate backbones for each task. This layer has the main advantage that it can be learned across few-shot tasks from many domains so it can support transfer across these tasks.