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: In this paper, the authors focus on the problem of learning unsupervised representations that are also disentangled. To this end, they build on top of the idea of invariant risk minimization and propose Iterative Partition-based Invariant Risk Minimization (IP-IRM) which iteratively assigns examples to partitions and constrains invariance in the loss among these partitions. They justify their approach mathematically according to a formal definition of disentanglement. The then perform multiple experiments qualitatively and quantitatively demonstrating the learned representations are disentangled. Furthermore, they show promising results on Out-of-distribution classification tasks.