Self-Supervised Learning Disentangled Group Representation as Feature

1 Introduction Deep learning is all about learning feature representations [ 5 ] . Compared to the conventional end-to-end supervised learning , Self-Supervised Learning ( SSL ) first learns a generic feature representation ( e.g. , a network backbone ) by training with unsupervised pretext tasks such as the prevailing contrastive objective [ 36 , 16 ] , and then the above stage-1 feature is expected to serve various stage-2 applications with proper fine-tuning . SSL for visual representation is so fascinating that it is the first time that we can obtain “ good ” visual features for free , just like the trending pre-training in NLP community [ 26 , 8 ] . However , most SSL works only care how much stage-2 performance an SSL feature can improve , but overlook what feature SSL is learning , why it can be learned , what can not be learned , what the gap between SSL and Supervised Learning ( SL ) is , and when SSL can surpass SL ? The crux of answering those questions is to formally understand what a feature representation is and what a good one is . We postulate the classic world model of visual generation and feature representation [ 1 , 69 ] as in Figure 1 . Let U be a set of ( unseen ) semantics , e.g. , attributes such as “ digit ” and “ color ” . There 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Entangle Semantics 1344D 1212D SimCLR SimCLR + PIRM Aug. Related L1 Norm = 0.06 L1 Norm = 0.06 ( a ) ( c ) Aug . Unrelated 20 40 60 80 SimCLR IP-IRM SimCLR IP-IRM Ac cu ra cy ( b ) IP-IRMSimCLR Layer 29 IP-IRMSimCLR Layer 18 IP-IRMSimCLR Layer 29 IP-IRMSimCLR Layer 18 IP-IRMSimCLR Layer 29 IP-IRMSimCLR Layer 18 0 1 2 3 Figure 2 : ( a ) The heat map visualizes feature dimensions related to augmentations ( aug. related ) and unrelated to augmentations ( aug. unrelated ) , whose respective classification accuracy is shown in the bar chart below . Dashed bar denotes the accuracy using full feature dimensions . Experiment was performed on STL10 [ 22 ] with representation learnt with SimCLR [ 16 ] and our IP-IRM . ( b ) Visualization of CNN activations [ 77 ] of 4 filters on layer 29 and 18 of VGG [ 75 ] trained on ImageNet100 [ 81 ] . The filters were chosen by first clustering the aug. unrelated filters with k-means ( k = 4 ) and then selecting the filters corresponding to the cluster centers . is a set of independent and causal mechanisms [ 66 ] ϕ : U → I , generating images from semantics , e.g. , writing a digit “ 0 ” when thinking of “ 0 ” [ 74 ] . A visual representation is the inference process φ : I → X that maps image pixels to vector space features , e.g. , a neural network . We define semantic representation as the functional composition f : U → I → X . In this paper , we are only interested in the parameterization of the inference process for feature extraction , but not the generation process , i.e. , we assume ∀I ∈ I , ∃u ∈ U , such that I = ϕ ( u ) is fixed as the observation of each image sample . Therefore , we consider semantic and visual representations the same as feature representation , or simply representation , and we slightly abuse φ ( I ) : = f ( ϕ−1 ( I ) ) , i.e. , φ and f share the same trainable parameters . We call the vector x = φ ( I ) as feature , where x ∈ X . We propose to use Higgins ’ definition of disentangled representation [ 40 ] to define what is “ good ” . Definition 1 . ( Disentangled Representation ) Let G be the group acting on U , i.e. , g · u ∈ U × U transforms u ∈ U , e.g. , a “ turn green ” group element changing the semantic from “ red ” to “ green ” . Suppose there is a direct product decomposition1 G = g1× . . .× gm and U = U1× . . .×Um , where gi acts on Ui respectively . A feature representation is disentangled if there exists a group G acting on X such that : 1 . Equivariant : ∀g ∈ G , ∀u ∈ U , f ( g ·u ) = g ·f ( u ) , e.g. , the feature of the changed semantic : “ red ” to “ green ” in U , is equivalent to directly change the color vector in X from “ red ” to “ green ” . 2 . Decomposable : there is a decomposition X = X1 × . . .×Xm , such that each Xi is fixed by the action of all gj , j 6= i and affected only by gi , e.g. , changing the “ color ” semantic in U does not affect the “ digit ” vector in X . Compared to the previous definition of feature representation which is a static mapping , the disentangled representation in Definition 1 is dynamic as it explicitly incorporate group representation [ 35 ] , which is a homomorphism from group to group actions on a space , e.g. , G → X × X , and it is common to use the feature space X as a shorthand—this is where our title stands . Definition 1 defines “ good ” features in the common views : 1 ) Robustness : a good feature should be invariant to the change of environmental semantics , such as external interventions [ 45 , 87 ] or domain shifts [ 32 ] . By the above definition , a change is always retained in a subspace Xi , while others are not affected . Hence , the subsequent classifier will focus on the invariant features and ignore the ever-changing Xi . 2 ) Zero-shot Generalization : even if a new combination of semantics is unseen in training , each semantic has been learned as features . So , the metrics of each Xi trained by seen samples remain valid for unseen samples [ 95 ] . Are the existing SSL methods learning disentangled representations ? No . We show in Section 4 that they can only disentangle representations according to the hand-crafted augmentations , e.g. , color jitter and rotation . For example , in Figure 2 ( a ) , even if we only use the augmentation-related feature , 1Note that gi can also denote a cyclic subgroup Gi such as rotation [ 0◦ : 1◦ : 360◦ ] , or a countable one but treated as cyclic such as translation [ ( 0 , 0 ) : ( 1 , 1 ) : ( width , height ) ] and color [ 0 : 1 : 255 ] . the classification accuracy of a standard SSL ( SimCLR [ 16 ] ) does not lose much as compared to the full feature use . Figure 2 ( b ) visualizes that the CNN features in each layer are indeed entangled ( e.g. , tyre , motor , and background in the motorcycle image ) . In contrast , our approach IP-IRM , to be introduced below , disentangles more useful features beyond augmentations . In this paper , we propose Iterative Partition-based Invariant Risk Minimization ( IP-IRM [ ai '' p @ : m ] ) that guarantees to learn disentangled representations in an SSL fashion . We present the algorithm in Section 3 , followed by the theoretical justifications in Section 4 . In a nutshell , at each iteration , IP-IRM first partitions the training data into two disjoint subsets , each of which is an orbit of the already disentangled group , and the cross-orbit group corresponds to an entangled group element gi . Then , we adopt the Invariant Risk Minimization ( IRM ) [ 2 ] to implement a partition-based SSL , which disentangles the representation Xi w.r.t . gi . Iterating the above two steps eventually converges to a fully disentangled representation w.r.t . ∏m i=1 gi . In Section 5 , we show promising experimental results on various feature disentanglement and SSL benchmarks . 2 Related Work . Self-Supervised Learning . SSL aims to learn representations from unlabeled data with hand-crafted pretext tasks [ 28 , 63 , 33 ] . Recently , Contrastive learning [ 65 , 61 , 38 , 80 , 16 ] prevails in most state-ofthe-art methods . The key is to map positive samples closer , while pushing apart negative ones in the feature space . Specifically , the positive samples are from the augmented views [ 82 , 3 , 94 , 42 ] of each instance and the negative ones are other instances . Along this direction , follow-up methods are mainly four-fold : 1 ) Memory-bank [ 90 , 61 , 36 , 18 ] : storing the prototypes of all the instances computed previously into a memory bank to benefit from a large number of negative samples . 2 ) Using siamese network [ 7 ] to avoid representation collapse [ 34 , 19 , 83 ] . 3 ) Assigning clusters to samples to integrate inter-instance similarity into contrastive learning [ 11 , 12 , 13 , 88 , 56 ] . 4 ) Seeking hard negative samples with adversarial training or better sampling strategies [ 73 , 20 , 44 , 48 ] . In contrast , our proposed IP-IRM jumps out of the above frame and introduces the disentangled representation into SSL with group theory to show the limitations of existing SSL and how to break through them . Disentangled Representation . This notion dates back to [ 4 ] , and henceforward becomes a highlevel goal of separating the factors of variations in the data [ 84 , 79 , 86 , 58 ] . Several works aim to provide a more precise description [ 27 , 29 , 72 ] by adopting an information-theoretic view [ 17 , 27 ] and measuring the properties of a disentangled representation explicitly [ 29 , 72 ] . We adopt the recent group-theoretic definition from Higgins et al . [ 40 ] , which not only unifies the existing , but also resolves the previous controversial points [ 78 , 59 ] . Although supervised learning of disentangled representation is a well-studied field [ 100 , 43 , 10 , 70 , 49 ] , unsupervised disentanglement based on GAN [ 17 , 64 , 57 , 71 ] or VAE [ 39 , 15 , 99 , 50 ] is still believed to be theoretically challenging [ 59 ] . Thanks to the Higgins ’ definition , we prove that the proposed IP-IRM converges with full-semantic disentanglement using group representation theory . Notably , IP-IRM learns a disentangled representation with an inference process , without using generative models as in all the existing unsupervised methods , making IP-IRM applicable even on large-scale datasets . Group Representation Learning . A group representation has two elements [ 47 , 35 ] : 1 ) a homomorphism ( e.g. , a mapping function ) from the group to its group action acting on a vector space , and 2 ) the vector space . Usually , when there is no ambiguity , we can use either element as the definition . Most existing works focus on learning the first element . They first define the group of interest , such as spherical rotations [ 24 ] or image scaling [ 89 , 76 ] , and then learn the parameters of the group actions [ 23 , 46 , 68 ] . In contrast , we focus on the second element ; more specifically , we are interested in learning a map between two vector spaces : image pixel space and feature vector space . Our representation learning is flexible because it delays the group action learning to downstream tasks on demand . For example , in a classification task , a classifier can be seen as a group action that is invariant to class-agnostic groups but equivariant to class-specific groups ( see Section 4 ) .

Summary: This paper proposes a self-supervised learning model which is a variant of the Invariant Risk Minimization with an iterative partition mechanism. The method can encourage the self-supervised learning of a disentangled representation. Empirical results validate the effectiveness of the proposed model on self-supervised learning and disentangled representation learning.