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 introduces a setup with the goal to adapt from a coarse pretrained model to unseen fine-grained labels. This problem is formulated as a superclass-subclass latent model and learned with maximum likelihood via expectation-maximization. The proposed approach, super-class conditional Gaussian mixture (SCGM) model, defines a hierarchical Gaussian distribution on the class hierarchy that models both the superclasses and subclasses. SCGM is evaluated on 6 image-net related datasets and a a real Dialysis-Event dataset collected by hospitals, and demonstrates competitive results under two evaluation setups: generalizing to seen and unseen superclass.