Uniform Priors for Data-Efficient Transfer

1 INTRODUCTION . Deep Neural Networks have enabled great success in various machine learning domains such as computer vision ( Girshick , 2015 ; He et al. , 2016 ; Long et al. , 2015 ) , natural language processing ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Brown et al. , 2020 ) , decision making ( Schulman et al. , 2015 ; 2017 ; Fujimoto et al. , 2018 ) or in medical applications ( Ronneberger et al. , 2015 ; Hesamian et al. , 2019 ) . This can be largely attributed to the ability of networks to extract abstract features from data , which , given sufficient data , can effectively generalize to held-out test sets . However , the degree of generalization scales with the semantic difference between test and training tasks , caused e.g . by domain or distributional shifts between training and test data . Understanding how to achieve generalization under such shifts is an active area of research in fields like MetaLearning ( Snell et al. , 2017 ; Finn et al. , 2017 ; Chen et al. , 2020 ) , Deep Metric Learning ( DML ) ( Roth et al. , 2020b ; Hadsell et al. , 2006 ) , Zero-Shot Domain Adaptation ( ZSDA ) ( Tzeng et al. , 2017 ; Kodirov et al. , 2015 ) or low-level vision tasks ( Tang et al. , 2020 ) . In the few-shot Meta-Learning setting , a meta-learner is tasked to quickly adapt to novel test data given its training experience and a limited labeled data budget ; similarly fields like DML and ZSDA study generalization at the limit of such adaptation , where predictions on novel test data are made without any test-time finetuning . Yet , despite the motivational differences , each of these fields require representations to be learned from the training data that allow for better generalization and adaptation to novel tasks and data . Although there exists a large corpus of domain-specific training methods , in this paper we seek to investigate what fundamental properties learned features and feature spaces should have to facilitate such generalization . Fortunately , recent literature provides pointers towards one such property : the notion of “ feature uniformity ” for improved generalization . For Unsupervised Representation Learning , Wang & Isola ( 2020 ) highlight a link between the uniform distribution of hyperspherical feature representations and the transfer performance in downstream tasks , which has been implicitly adapted in the design of modern contrastive learning methods ( Bachman et al. , 2019 ; Tian et al. , 2020a ; b ) . Similarly , Roth et al . ( 2020b ) show that for Deep Metric Learning , uniformity in hyperspherical embedding space coverage as well as uniform singular value distribution embedding spaces are strongly connected to zero-shot generalization performance . Both Wang & Isola ( 2020 ) and Roth et al . ( 2020b ) link the uniformity in the feature representation space to the preservation of maximal information and reduced overfitting . This suggests that actively imposing a uniformity prior on learned feature representations should encourage better transfer properties by retaining more information and reducing bias towards training tasks , which in turn facilitate better adaptation to novel tasks . However , while both Wang & Isola ( 2020 ) and Roth et al . ( 2020b ) propose methods to incorporate this notion of uniformity , they are defined only for hyperspherical embedding spaces or contrastive learning approaches1 , thus severely limiting the applicability to other domains . To address these limitations and leverage the benefits of uniformity for any type of novel task and data adaptation for deep neural networks , we propose uniformity regularization , which places a uniform hypercube prior on the learned features space during training , without being limited to the contrastive training approaches or a hyperspherical representation space . Unlike e.g . a multivariate Gaussian , the uniform prior puts equal likelihood over the feature space , which then enables the network to make fewer assumptions about the data , limiting model overfitting to the training task . This incentivizes the model to learn more task-agnostic and reusable features , which in turn improve generalization ( Raghu et al. , 2019 ) . Our uniformity regularization follows an adversarial learning framework that allows us to apply our proposed uniformity prior , since a uniform distribution does not have a closed-form divergence minimization scheme . Using this setup , we experimentally demonstrate that uniformity regularization aids generalization in zero-shot setups such as Deep Metric Learning , Domain Adaptation , Out-of-Distribution Detection as well as few-shot Meta-Learning . Furthermore , for Deep Metric learning and few-shot Meta-Learning , we are even able to set a new state-of-the-art over benchmark datasets . Overall , our contributions can be summarized as : • We propose to perform uniformity regularization in the embedding spaces of a deep neural network , using a GAN-like alternating optimization scheme , to increase the transferability of learned features and the ability for better adaptation to novel tasks and data . • Using our proposed regularization , we achieve strong improvements over baseline methods in Deep Metric Learning , Zero-Shot Domain Adaptation , Out-of-Distribution Detection and Meta-Learning . Furthermore , uniformity regularization allows us to set a new state-ofthe-art in Meta-Learning on the Meta-Dataset ( Triantafillou et al. , 2019 ) as well as in Deep Metric Learning over two benchmark datasets ( Welinder et al. , 2010 ; Krause et al. , 2013 ) . 2 BACKGROUND . 2.1 GENERATIVE ADVERSARIAL NETWORKS ( GANS ) . Generative Adversarial Networks ( GANs , Goodfellow et al . ( 2014 ) ) were proposed as a generative model which utilizes an alternative optimization scheme that solves a minimax two-player game between a generator , G , and a discriminator , D. The generator G ( z ) is trained to map samples from a prior z ∼ p ( z ) to the target space , while the discriminator is trained to be an arbiter between the target data distribution p ( x ) and the generator distribution . The generator is trained to trick the discriminator into predicting that samples from G ( z ) actually stem from the target distribution . While many different GAN objectives have been proposed , the standard “ Non-Saturating Cost ” generator objective as well as the discriminator objective can be written as LD = max D Ez∼p ( z ) [ 1− logD ( G ( z ) ) ] + Ex∼p ( x ) [ logD ( x ) ] ( 1 ) LG = min G Ez∼p ( z ) [ 1− logD ( G ( z ) ) ] ( 2 ) with p ( z ) the generator prior and p ( x ) a defined target distribution ( e.g . natural images ) . 2.2 FAST ADAPTATION AND GENERALIZATION . Throughout this work , we use the notion of “ fast adaptation ” to novel tasks to measure the transferability of learned features , and as such the generalization and adaptation capacities of a model . Fast adaptation has recently been popularized by different meta-learning strategies ( Finn et al. , 2017 ; Snell et al. , 2017 ) . These methods assume distinct meta-training and meta-testing task distributions , 1By imposing a Gaussian potential over hyperspherical embedding distances or pairwise sample relations . where the goal of a meta-learner is to adapt fast to a novel task given limited samples for learning it . Specifically , a few-shot meta-learner is evaluated to perform n-way classification given k ‘ shots ’ , corresponding to k examples taken from n previously unseen classes . Generally , one distinguishes two types of meta-learners : ones requiringm training iterations for finetuning ( Finn et al. , 2017 ; Rajeswaran et al. , 2019 ) , and ones that do not ( Snell et al. , 2017 ; Lee et al. , 2019 ) . In the meta-learning phase , the meta-learner is trained to solve entire tasks as ( meta-training ) datapoints . Its generalization is measured by how well it can quickly adapt to novel test tasks . Many different strategies have been introduced to maximize the effectiveness of the meta-learning phase such as episodic training , where the model is trained by simulating ‘ test-like ’ conditions ( Vinyals et al. , 2016 ) , or finetuning , where the model performs up to m gradient steps on the new task ( Finn et al. , 2017 ) . While such meta-learning approaches assume the availability of a finetuning budget for adaptation at test time , zero-shot approaches introduce the limit scenario of fast adaptation , in which generalization has to be achieved without access to any examples . Such a setting can be found in metric learning ( Yang & Jin , 2006 ; Suárez et al. , 2018 ) , where a model is evaluated on the ability to perform zero-shot retrieval on novel data . Most commonly , metric models are trained on a training data distribution Dtrain and evaluated on a testing distribution Dtest which share no classes . However , the data generating function is assumed to be similar between Dtrain and Dtest , such as natural images of birds ( Welinder et al. , 2010 ) . While vanilla metric learning learns a parametrized metric on fixed feature extractors , Deep Metric Learning ( DML ) leverages deep neural networks to train the feature extractors concurrently ( Roth et al. , 2020b ) . Such deep abstractions further allow for simplified and computationally cheap metrics such as euclidean distances which make large-scale retrieval applications with fast similarity searches possible ( e.g . Johnson et al . ( 2017 ) ) . Similar to DML , Zero-Shot Domain Adaptation ( ZSDA ) introduces a learner that is also trained and evaluated on two distinctDtrain andDtest in a zero-shot setting . However , unlike DML , in ZSDA , the labels between the data distributions are shared . Instead , training and test distribution come from distinct data generative functions , such as natural images of digits ( Goodfellow et al. , 2013 ) and handwritten images of digits ( LeCun , 1998 ) . 3 EXTENDING NETWORK TRAINING WITH A UNIFORMITY PRIOR . In this section , we introduce the proposed uniformity regularization and detail the employed alternating GAN-like optimization scheme to perform it in a computationally tractable manner . Prior Matching . Given a neural network q ( y|x ) that is parameterized by θ we formally define the training objective as LT ( q ( y|x ) , y ) where LT is any task-specific loss such as a cross-entropy loss , ( x , y ) are samples from the training distribution Dtrain and q ( y|x ) the probability of predicting label y under q . This is a simplified formulation ; in practice , there are many different ways to train a neural network , such as ranking-based training with tuples ( Chopra et al. , 2005 ) . We define the embedding space z as the output of the final convolutional layer of a deep network . Accordingly , we ’ ll note q ( z|x ) as the conditional distribution for that embedding space which , due to the convnet being a deterministic mapping , is a dirac delta distribution at the value of the final convolutional layer . Section 4.1 further details how to apply uniformity regularization in practice . As we ultimately seek to impose a uniformity prior over the learned aggregate feature/embedding “ posterior ” q ( z ) = ∫ x q ( z|x ) p ( x ) dx , we begin by augmenting the generic task-objective to allow for the placement of a prior r ( z ) . For priors r ( z ) with closed-form KL-divergences D , one can define a prior-regularized task objective as L = min θ E ( x , y ) ∼Dtrain [ LT ( q ( y|x ) , y ) ] +Dx∼Dtrain ( q ( z|x ) ‖r ( z ) ) ( 3 ) similar to the Variational Autoencoder formulation in Kingma & Welling ( 2013 ) . However , to improve the generalization of a network by encouraging uniformity in the learned embeddings , we require regularization by matching the learned embedding space to a uniform distribution prior U ( −α , β ) , defined by the lower and upper bounds α and β , respectively . Unfortunately , such a regularization does not have a simple solution in practice , as a bounded uniform distribution has no closed-form KL divergence metric to minimize . Uniformity Regularization . To address the practical limitation of solving Eqn . 3 , we draw upon the GAN literature , in which alternate adversarial optimization has been successfully used to match a generated distribution to a defined target distribution using implicit divergence minimization . Latent variable models such as the Adversarial Autoencoder ( Makhzani et al. , 2015 ) have successfully used such a GAN-style adversarial loss , instead of a KL divergence , in the latent space of the autoencoder to learn a rich posterior . Such implicit divergence minimization allows us to match any well-defined distribution as a prior , but more specifically , ensures that we can successfully match learned embedding spaces to U ( −α , α ) , which we set to the unit hypercube U ( −1 , 1 ) by default . To this end , we adapt the GAN objective in Eqn . 1 and 2 for uniformity regularization optimization and train a discriminator , D , to be an arbiter between which samples are from the learned distribution q ( z|x ) and from the uniform prior r ( z ) . As such , the task model q ( parameterized by θ ) aims to fool the discriminator D into thinking that learned features , q ( z|x ) , come from the chosen uniform target distribution , r ( z ) , while the discriminator D learns to distinguish between learned features and samples taken from the prior , z̃ ∼ r ( z ) . Note that while the task-model defines a deterministic mapping for q ( z|x ) instead of a stochastic one , the aggregate feature “ posterior ” ∫ x q ( z|x ) p ( x ) dx , on which we apply our uniformity prior , is indeed a stochastic distribution ( Makhzani et al. , 2015 ) . Concretely for our uniformity regularization , we rewrite the discriminator objective from Eqn . 1 to account for the uniform prior matching , giving LD = max D Ex∼Dtrain [ log ( 1−D ( q ( z|x ) ) ) ] + Ez̃∼U ( −1,1 ) [ logD ( z̃ ) ] ( 4 ) Consequently , we reformulate the generator objective from Eqn . 2 to reflect the task-model q , Lmax = min θ Ex∼Dtrain [ log ( 1−D ( q ( z|x ) ) ) ] ( 5 ) where we used the notation Lmax to reflect that optimization maximizes the feature uniformity by learning to fool D. Our final min-max uniformity regularized objective for θ and the Discriminator is then given as L = min θ max D E ( x , y ) ∼Dtrain [ LT ( qθ ( y|x ) , y ) ] +γEx∼Dtrain [ log ( 1−D ( qθ ( z|x ) ) ) ] +Ez̃∼U ( −1,1 ) [ logD ( z̃ ) ] ( 6 ) with task-objective LT and training data distribution Dtrain . Using this objective , the learned feature space is implicitly encouraged to become more uniform . The amount of regularization is controlled by the hyperparameter γ , balancing generalization of the model to new tasks and performance on the training task at hand . Large γ values hinder effective feature learning from training data , while values of γ too small result in weak regularization , leading to a non-uniform learned feature distribution with reduced generalization capabilities .

Summary: The authors argue that uniform priors for the high-level latent representations improve transferability, which is beneficial in a number of tasks involving transference. The approach is evaluated on deep metric learning, zero-shot domain adaptation and few-shot meta-learning. The authors propose a uniformity regularization term on the latent representation, implemented as an adversarial discrepancy. The results show consistent improvement in the different tasks.