Superclass-Conditional Gaussian Mixture Model For Learning Fine-Grained Embeddings

1 INTRODUCTION . Training deep models with sufficient generalizability is of fundamental importance , which demands immense training data with fine-grained annotations ( Krizhevsky et al. , 2012 ; Brown et al. , 2020 ) . In many fields , however , data labeling requires domain-specific knowledge , such as medicine ( Sohoni et al. , 2020 ) , thus is prohibitive , and infeasible to be exhaustive . In this case , data for model training may only be “ coarsely ” labeled , while later the model is tested on a finer-grained classification task ( Bukchin et al. , 2021 ) . For example , consider an event prediction task for dialysis patients ( Inaguma et al. , 2019 ) . Hemodialysis is a major renal replacement therapy for patients with end-stage renal failure . These patients have to take hemodialysis thrice a week , each lasts for 4-5 hours . During the treatment , unexpected events , such as muscle cramp , perspiration , and dizziness , may happen as a result of lowering blood pressure , which need intensive medical care , thus should always be avoided . It is therefore an important medical issue to predict such events before an initiation of hemodialysis . In this task , binary labels , which mark the incidence of an event , can be collected . In contrast , finergrained labels that annotate different subtypes of events are seldom recorded . Since distinguishing different subtypes facilitates precise diagnoses , and helps physicians assess the risk for deciding whether to perform a hemodialysis ( with certain precautions ) , it is desirable that a model trained with coarse ( binary ) labels can perform well on a finer-grained multi-class ( subtypes ) task . To fill the gap of granularity between the training and testing scenarios , a practical way is to collect a few new records for a patient , with their fine-grained annotations . These data constitute a support set for fine-tuning a pre-trained model to the specific data distribution induced by the annotations of the target patient , for whom the adapted model is used for future predictions . Although massive finegrained annotation is impractical , annotating a few-shot set is feasible . In this work , we are interested in such a Cross-Granularity Few-Shot ( CGFS ) learning problem , where a model pre-trained on a set of coarse classes ( denoted as superclasses ) , needs to adapt to an unseen set of fine-grained target classes ( denoted as subclasses ) . The target subclasses could be descendants of the superclasses ( as in the aforementioned example ) , or could descend from other superclasses that are unobserved during pre-training . To be practical , the adaptation should only use a few samples from the subclasses . The CGFS problem is not limited to the above application . It occurs in a model ’ s lifespan when an application requires separating some subclasses from the superclasses yet when the training dataset was created these subclasses were unannotated . For example , it could occur in detecting rare pathology or variants using medical images that were coarsely described ( Oakden-Rayner et al. , 2020 ) , or personalizing a generic model trained on all historical users to a specific customer ( Luo et al. , 2020 ) . Despite its significance , CGFS can not be trivially solved by regularly training models with coarse labels , because typical losses for supervised learning aim to maximize inter-class boundaries but neglect intra-class variation . Thus , subclasses may arbitrarily and unevenly spread in every superclass . Recently , Bukchin et al . ( 2021 ) proposed to integrate coarse supervision and contrastive learning ( within superclasses ) for solving CGFS . However , their approach can not be readily used for medical records which typically contain static profiles and time series ( Che et al. , 2016 ) . This is due to the absence of a standard augmentation method for generating constrastive pairs on some data other than images . Also , since contrastive learning does not model subclass explicitly , their solution could be suboptimal ( as evaluated in Sec . 4 ) . Moreover , as their model was built upon MoCo ( He et al. , 2020 ) , and maintains many dictionaries for data sampling , its computational costs are high ( Sec . 4 ) . In this work , we propose a novel Superclass-Conditional Gaussian Mixture model ( SCGM ) to learn fine-grained embeddings for the CGFS problem . Our contributions are summarized as follows . • SCGM is agnostic to the encoder , thus is flexible to different applications . It models the generation of samples from hierarchical classes , and explicitly represents the unobserved subclasses by latent variables , without assuming their identities . It dynamically computes a Gaussian mixture for every sample conditioned on its superclass , and the model forms a hierarchy of Gaussian mixtures . • SCGM only adds a small overhead to an encoder , for parameterizing its distributions , thus is efficient . The model parameters are learned end-to-end by maximum likelihood estimation via a principled Expectation-Maximization ( EM ) algorithm . We also theoretically linked our loss function to InfoNCE ( Oord et al. , 2018 ) , explaining its effectiveness from a contrastive perspective . • In the experiments , we evaluated SCGM on both benchmark image datasets and a real-life medical dataset . Since SCGM is compatible with contrastive learning , we also tested it with a momentum encoder ( He et al. , 2020 ) . The results demonstrate SCGM on generic encoders has already outperformed the state-of-the-art ( SOTA ) baselines , with less computational costs , and it achieves boosted performance in some cases with momentum contrast integrated . 2 RELATED WORK . To our best knowledge , this is the first work to develop an SCGM model for underlying a framework that enables tackling CGFS across domains . The most relevant work ( Bukchin et al. , 2021 ) combines superclass-wise contrastive learning with coarse classification for preserving intra-class variation . Another work ( Yang et al. , 2021 ) used a three-step approach that pseudo-labels the embeddings ( pretrained by coarse classification and batch-wise contrastive learning ) by clustering every superclass . The pseudo-fine-labels were used to re-train the encoder . A similar three-step method ( Sohoni et al. , 2020 ) used a different loss for maximizing the worst-case expected accuracy over the pseudo-labeled subclasses . As discussed before , the former two methods require non-trivial searches for a suitable data augmentation method , which may be unavailable for some non-image data . By splitting training steps , the latter two methods could lead to suboptimal pseudo-labeling , and misleading labels could confuse the downstream steps . In contrast , our model is end-to-end . It explicitly infers the posterior of the subclass during coarse training . Moreover , it can achieve better performance without using as many computational resources as the model of ( Bukchin et al. , 2021 ) ( as evaluated in Sec . 4 ) . Learning with coarse supervision . Several works deal with coarse/fine labels from the perspective of weakly supervised learning ( Zhou , 2018 ) than to tackle the CGFS problem , including training methods that take advantage of a mixture of either balanced ( Ristin et al. , 2015 ; Guo et al. , 2018 ; Taherkhani et al. , 2019 ) or unbalanced ( Hsieh et al. , 2019 ; Liu et al. , 2019 ; Robinson et al. , 2020 ) coarse and fine labels . Among them , Liu et al . ( 2019 ) addressed a few-shot learning problem , but the model training assumes access to some fine labels and a graph of class hierarchy , which are unavailable in our problem . Thus it can not be adapted to solve our problem . Few-shot learning . Meta-learning has become a popular idea to handle the few-shot learning problem , which derives metric-based ( Vinyals et al. , 2016 ; Snell et al. , 2017 ) and optimization-based methods ( Finn et al. , 2017 ; Nichol et al. , 2018 ) . The idea has been extended to semi-supervised ( Ren et al. , 2018 ) , unsupervised ( Hsu et al. , 2018 ) and semantics-augmented ( Xing et al. , 2019 ) scenarios , when labels are scarce . However , none of them explores coarse labels for cross-granularity learning . Recently , many works observed that learning embeddings on all classes ( without episodic training ) , followed by simple fine-tuning , is superior to SOTA meta-learning methods ( Wang et al. , 2019 ; Dhillon et al. , 2020 ; Tian et al. , 2020 ) . In this work , similar to ( Bukchin et al. , 2021 ) , we focus on this paradigm to learn useful embeddings , and do not use meta-learning for pre-training . Embedding methods . The emerging self-supervised methods , such as contrastive methods ( Oord et al. , 2018 ; Chen et al. , 2020a ; He et al. , 2020 ) , are appealing in their ability to attain comparable embeddings to the supervised counterparts , and even surpass them when transferring to other tasks ( He et al. , 2020 ; Tian et al. , 2020 ) . Beyond instance-wise contrast , recent methods have explored instance-cluster contrast to boost performance ( Caron et al. , 2020 ; Li et al. , 2020 ) . In unsupervised methods , joint embedding and clustering has been found beneficial ( Caron et al. , 2018 ; Asano et al. , 2020 ) . Whereas , these methods never exploited coarse supervision . Thus , their embeddings/clusters do not necessarily reflect intra-class variation , which is important to the CGFS task . Vanilla Gaussian mixture ( GM ) model is unfit for the CGFS scenario . A recent work ( Manduchi et al. , 2021 ) extended GM to constrained clustering by conditioning every sample with a prior clustering preference . It is neither supervised by coarse classes nor hierarchically structured . Conventional hierarchical GM is used for hierarchical clustering ( Goldberger & Roweis , 2005 ; Olech & Paradowski , 2016 ; Athey et al. , 2019 ) by applying GM agglomeratively or divisively . These unsupervised methods only infer clusters , but do not pre-train embedding models for task adaptation . 3 SUPERCLASS-CONDITIONAL GAUSSIAN MIXTURE MODEL . Firstly , a word about some notations . Let Dtrain = { ( xi , yi ) } ni=1 be n sample-label training pairs , where yi ∈ Ysuper = { 1 , ... , c } is a superclass label . Each xi is associated with a latent ( unobserved ) subclass label ŷi ∈ Ysub = { 1 , ... , s } . Ysub relates to Ysuper by a hierarchical structure , i.e. , Ysub can be partitioned into c disjoint sets Ysub-1 , ... , Ysub-c , such that if ŷi ∈ Ysub-j , then yi = j ( 1 ≤ j ≤ c ) . Let fθ be an encoder ( i.e. , backbone network ) that is trained on Dtrain ( without knowing Ysub ) . It maps xi to a d-dimensional feature fθ ( xi ) . At test time , given a k-shot support set for a subset Ymsub ⊆ Ysub of m subclasses , i.e. , Dsupport = { ( xi , ŷi ) |ŷi ∈ Ymsub } mki=1 , the task is to train a classifier C : Rd → Ymsub with optimal accuracy on a test set of Ymsub subclasses . In our experiments , we also explored the case when the subclasses belong to superclasses that are not used for pre-training . As discussed in Sec . 2 , our focus is to train fθ for good embeddings without modifying fθ during the adaptation–a paradigm with SOTA few-shot performance ( Dhillon et al. , 2020 ; Tian et al. , 2020 ) . Formally , let fθ ( xi ) = vi , our goal is to find the model parameter θ that maximizes the likelihood of the posterior distribution pθ ( yi|vi ) on the observed data in Dtrain for classification tasks ( Grathwohl et al. , 2019 ) . To model the unobserved subclasses , we associate every vi with a latent variable zi to indicate to which subclass vi belongs . Suppose there are r possible subclasses , the log-likelihood to maximize can be rewritten by marginalizing out the latent variables . ℓ ( Dtrain ; θ ) = 1 n n∑ i=1 log [ pθ ( yi|vi ) ] = 1 n n∑ i=1 log [ r∑ zi=1 p ( yi|zi ) pθ ( zi|vi ) ] ( 1 ) where the distribution pθ ( zi|vi ) specifies the subclass membership of vi , and p ( yi|zi ) associates zi with a subclass partition Ysub-yi . Unlike some previous works ( Sohoni et al. , 2020 ) , which searched the number of subclasses for every superclass using a quality metric , we only assume a total number of r subclasses , and seek to infer their relationship with superclasses , i.e. , p ( yi|zi ) , without any prior on subclass partitions , so that the model is more generalizable .

Summary: The paper presents SCGM, a new technique for solving the Cross-Granularity Few-Shot learning (CGFS) problem.  CGFS is defined as the problem of adapting a classification model trained on coarse (“superclass”) labels to perform well on fine-grained labels, which consist of multiple “subclass” labels within each superclass.  The paper presents a new generative modeling approach that enables end-to-end classifier training in a manner that (a) incorporates information about superclass-subclass hierarchies (b) does not rely on explicit subclass enumeration (c) provides improved empirical performance on CGFS and (d) provides systems benefits with respect to existing baselines.  The authors provide detailed description of their approach and characterize its relationship with existing methods, and present empirical results that support their methodological choices.