Discrete InfoMax Codes for Meta-Learning

1 INTRODUCTION . Generalizing to unseen data is a problem of vital importance in machine learning . Many deep metalearning methods optimize for generalization by directly minimizing the loss of held-out validation data . Recent works have used this framework to achieve impressive feats such as learning to classify using one labeled image per class ( Snell et al. , 2017 ; Finn et al. , 2017 ) , learning unsupervised update rules that generalize to different domains ( Metz et al. , 2018 ) , and accelerating training procedures that are millions of steps long ( Flennerhag et al. , 2018 ) . However , the meta-learning setup introduces a new overfitting problem : the model may overfit to the distribution of tasks seen during training . In other words , the meta-learning framework decreases task-specific overfitting at the cost of introducing task-wise overfitting . The primary aim of this work is to elucidate this tradeoff between task-specific and task-wise overfitting in meta-learning . Specifically , we use tools from information theory to bound the overall generalization gap of a meta-learner . Analogously to how generalization bounds for standard learning algorithms reveal the role of dataset size in generalizing to new datapoints , our bound reveals the roles of both dataset size ( m ) and number of datasets ( n ) in generalizing to new tasks . The specific form of our generalization gap suggests that meta-learning models can be implicitly regularized by constructing them to express each datapoint with a small number of bits . To this end , we propose Discrete InfoMax COdes ( DIMCO ) , a deep neural network model that outputs a discrete representation of each datapoint . Note that a continuous representation x̃ ∈ Rn ( e.g . Snell et al . ( 2017 ) ) uses 32n bits per datapoint . This is wildly inefficient in terms of bit-efficiency : our experiments in Section 5 show that DIMCO ’ s discrete representation requires roughly 10× less bits per datapoint to achieve similar performance compared to continuous methods . DIMCO generalizes well to novel datasets because its learning objective encourages the model to compactly use all of its degrees of freedom , thus enabling it to effectively compare datapoints while only requiring a small number of bits per datapoint . Our specific contributions are : 1 . Derive a generalization bound for meta-learning that makes the tradeoff between taskspecific and task-wise overfitting concrete . 2 . Propose DIMCO , a neural network model that is designed to have a low value of a specific term in our generalization bound . 3 . Empirically demonstrate that DIMCO generalizes better than previous meta-learning methods when trained with small datasets , and that it is more memory- and time-efficient compared to previous image retrieval methods . This paper is organized as follows . We detail our problem setup and derive a generalization bound for meta-learning in Section 2 . Taking motivation from our bound , we propose our meta-learning model ( DIMCO ) in Section 3 . We put our analysis in the context of previous work in Section 4 . Notably , we suggest that certain previous meta-learning methods may have benefitted from implicit regularization . We present experiments in Section 5 and conclude the paper with a discussion about limitations and future directions of our approach in Section 6 . 2 A GENERALIZATION BOUND FOR META-LEARNING . Throughout this section , we denote data , model outputs , and labels as x , x̃ , and y , respectively . We use capital symbols X , X̃ , Y to denote the random variables corresponding to x , x̃ , y . The ( discrete ) Mutual Information between two random variables X1 , X2 is defined as I ( X1 ; X2 ) = H ( X1 ) −H ( X1|X2 ) = H ( X2 ) −H ( X2|X1 ) ( 1 ) I ( X1 ; X2 ) is a symmetric quantity which measures the amount of information shared between X1 and X2 . It has its lowest value 0 when X1 and X2 are independent and increases with the correlation between X1 and X2 . We refer the reader to ( Cover & Thomas , 2012 ) for further exposition . 2.1 PROBLEM SETUP . We begin by describing our meta-learning problem setup . Define a task T to be a distribution over Z = X × Y . Let tasks T 1 , . . . , Tn be sampled i.i.d . from a distribution of tasks τ . Associated with each task T , we define a dataset DT = z1T , . . . , z m T = ( x 1 T , y 1 T ) , . . . , ( x m T , y m T ) which is a set of m i.i.d . samples from the data distribution ( zjT ∼ T ) . We consider models with parameters θ that map datapoints X to representations X̃ ( X , θ ) . Our objective is the expecteation of the negative mutual information between representations an labels across all tasks : L ( τ , θ ) = −ET∼τ [ I ( X̃ ( XT , θ ) ; YT ) ] . ( 2 ) This objective is closely related to both previous loss functions and evaluation metrics in setups that involve testing on unseen classes ( e.g . few-shot classification , image retrieval ) . We show that previous loss functions can be seen as approximations to ( 2 ) in Appendix A , and that mutual information is strongly correlated with metrics such as few-shot accuracy and Recall @ 1 in Section 5.1 . An important difference between ( 2 ) and previous objectives for few-shot classification is that we do not split each task into support/query ( also called train/test ) sets . This property bypasses the pesky issue of using batch normalization ( BN ) during meta-training . As Nichol & Schulman ( 2018 ) points out , BN can leak information from the support set to the query set . This issue is not a problem in our setup since we do not assume a seperate `` query set '' . Additionally , not using support/query splits enables meta-learning with tasks consisting of one image per class , as we demonstrate in Section 5.4 . Note that the standard meta-learning setup can not learn from such tasks since it requires at least two images per class ( one for support , one for query ) to compute the loss function . Overall , our model demonstrates that the commonly used construction of a held-out test set within each task is not strictly necessary for meta-learning . 2.2 GENERALIZATION BOUND . We bound the difference between expected loss and empirical loss : Theorem 1 . Let τ , n , m , X , X̃ , Y , θ , L be defined as above . Let dΘ be the VC dimension of the encoder X̃ ( · ) . Let Î ( X̃ ( XT , θ ) ; YT ) be the empirical estimate of the mutual information using finite dataset DT , and define empirical loss as L̂ ( T 1 : n , θ ) = − 1 n n∑ i=1 Î ( X̃ ( XT i , θ ) ; YT i ) . ( 3 ) The following inequality holds with high probability : L ( τ , θ ) − L̂ ( T 1 : n , θ ) ≤ O ( √ dΘ n log n dΘ ) +O ( |X̃| log ( m ) √ m ) +O ( |X̃||Y | m ) ( 4 ) Proof . We use standard Chernoff bounds along with a finite sample bound for mutual information from Shamir et al . ( 2010 ) ; see Appendix B . The generalization gap has three terms , two of which decrease as m increases , and the other decreases as n increases . Typically for few-shot learning , n is very large while m is small : the typical miniImagenet 5-way 1-shot setup has n > 1010 and m = 5 . We therefore claim that the latter two terms are the main difficulties for generalizing to new tasks . We see from Theorem 1 that these terms can be reduced by using small |X̃| . Therefore , in the context of meta-learning , using short representations ( i.e . small |X̃| ) can compensate for having a small train set ( i.e . small m ) . Meta-learning is typically formulated as performing two levels of learning : task-general learning and task-specific learning . Theorem 1 implies that in the tasks considered in recent literature , task-general learning should have almost no generalization gap , making task-specific learning the main source of ( meta- ) overfitting . This can be problematic since in many works , the task-specific learning algorithm is usually just a byproduct of whatever clever meta-learning loss was proposed . Our theorem suggests that we should pay more attention to directly regularizing this task-specific learner , and our DIMCO model , described in the next section , can be seen as a minimal working example in this direction . 3 DISCRETE INFOMAX CODES ( DIMCO ) . We now present our model , Discrete InfoMax COdes ( DIMCO ) . Motivated by Section 2 , DIMCO produces a short discrete code X̃ and is trained by maximizing mutual information I ( X̃ ; Y ) . Figure 1 graphically shows the overall structure of DIMCO . 3.1 FACTORIZED DISCRETE CODES . We propose a factorized discrete representation scheme which enables us to represent discrete distributions with exponentially fewer parameters compared to listing the probability of each event . We represent each event as the product of d independent events , each of which consists of p different possibilities . We thus have pd events in total , but only require pd parameters to represent the probability of each event . Binary codes can be viewed as a special case of this scheme where p = 2 . This factorization trick allows us to consider representations of size |X̃| = 64256 ( Section 5 ) . This representation has the advantage of requiring only d log2 p bits per datapoint , whereas a Ddimensional continuous vector embedding requires 32d bits ( assuming 32-bit floats ) . 3.2 MODEL . Recall that we represent a given image using d independent discrete distributions , each of which has p possibilities . First , a ( convolutional ) neural network enc ( · ) takes image X as input and outputs a vector of length dp , which we reshape into a matrix of size d× p : enc ( X ) = l11 l12 l13 . . . l1p l21 l22 l23 . . . l2p ... ... ... . . . ... ld1 ld2 ld3 . . . ldp ( 5 ) Each row of this matrix represents the logits of a discrete distribution . We apply the softmax function to each row to get probabilities . softmax ( li1 , . . . , lip ) = pi1 , . . . , pip ( 6 ) The ith codeword is sampled according to the categorical distribution following these probabilites : x̃i ∼ Cat ( pi1 , . . . , pip ) . ( 7 ) We simply concatenate each x̃i to obtain the representation : x̃ = ( x̃1 , . . . , x̃d ) . 3.3 TRAINING . Recall that x̃ is a discrete random variable and X̃ is its distribution . Instead of sampling x̃ ∼ X̃ , we directly use X̃ to compute the objective : I ( X̃ ; Y ) = H ( X̃ ) −H ( X̃|Y ) . ( 8 ) The first term , H ( X̃ ) , can be calculated by taking the average of all probabilities and computing the entropy : H ( X̃ ) = d∑ i=1 H ( X̃i ) = d∑ i=1 H ( Cat ( ∑m j=1 p j i1 m , ∑m j=1 p j i2 m , . . . , ∑m j=1 p j ip m ) ) ( 9 ) The second term is H ( X̃|Y ) = ∑ck=1 p ( Y = k ) H ( X̃|Y = k ) where c is the number of classes . The marginal probability of Y ( p ( Y = k ) ) is the frequency of k in { y1 , . . . , ym } . H ( X̃|Y = k ) can be obtained by computing ( 9 ) using only xj for which yj = k. Though we have motivated the use of I ( X̃ ; Y ) as a loss function throughout this paper , we provide yet another perspective using the decomposition in ( 8 ) . Minimizing H ( X̃|Y ) encourages discriminatory behavior . This term encourages the average embedding of each class to be as concentrated as possible . Maximizing H ( X̃ ) incentivizes the model to overall use all possible values of X̃ . We emphasize that such closed-form computation of I ( X̃ ; Y ) is only possible because we are using discrete codes .

Summary: This paper proposes a method to learn classifier outputs for meta learning in the form of factorized discrete codes by maximizing the mutual information between the model's outputs and the ground truth labels. The authors further present an information theoretic generalization bound for meta learning in terms of the number of tasks, number of training samples per task and the expressively of the model. They further show empirically that their approach does not need a separate query set during meta-training and can generalize better than a few of the other metric-learning based meta-learning approaches specifically at lower shot values. They show that their method requires less memory than the N-pair meta-learning method.