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: In this paper, the authors claimed that uniformity in embedding space if the key for good generalization, and then propose an adversarial training based method to improve the uniformity of feature space. The claim is from previous work, thus the key contribution is the way to impose such regularization. The method itself makes sense to me.