On Linear Identifiability of Learned Representations

1 INTRODUCTION . An increasingly common methodology in machine learning is to improve performance on a primary down-stream task by first learning a high-dimensional representation of the data on a related , proxy task . In this paradigm , training a model reduces to fine-tuning the learned representations for optimal performance on a particular sub-task ( Erhan et al. , 2010 ) . Deep neural networks ( DNNs ) , as flexible function approximators , have been surprisingly successful in discovering effective high-dimensional representations for use in downstream tasks such as image classification ( Sharif Razavian et al. , 2014 ) , text generation ( Radford et al. , 2018 ; Devlin et al. , 2018 ) , and sequential decision making ( Oord et al. , 2018 ) . When learning representations for downstream tasks , it would be useful if the representations were reproducible , in the sense that every time a network relearns the representation function on the same data distribution , they were approximately the same , regardless of small deviations in the initialization of the parameters or the optimization procedure . In some applications , such as learning real-world causal relationships from data , such reproducible learned representations are crucial for accurate and robust inference ( Johansson et al. , 2016 ; Louizos et al. , 2017 ) . A rigorous way to achieve reproducibility is to choose a model whose representation function is identifiable in function space . Informally speaking , identifiability in function space is achieved when , in the limit of infinite data , there exists a single , global optimum in function space . Interestingly , Figure 1 exhibits learned representation functions that appear to be the same up to a linear transformation , even on finite data and optimized without convergence guarantees ( see Appendix A.1 for training details ) . In this paper , we account for Figure 1 by making precise the relationship it exemplifies . We prove that a large class of discriminative and autoregressive models are identifiable in function space , up to a linear transformation . Our results extend recent advances in the theory of nonlinear Independent Components Analysis ( ICA ) , which have recently provided strong identifiability results for generative models of data ( Hyvärinen et al. , 2018 ; Khemakhem et al. , 2019 ; 2020 ; Sorrenson et al. , 2020 ) . Our key contribution is to bridge the gap between these results and discriminative models , commonly used for representation learning ( e.g. , ( Hénaff et al. , 2019 ; Brown et al. , 2020 ) ) . The rest of the paper is organized as follows . In Section 2 , we describe a general discriminative model family , defined by its canonical mathematical form , which generalizes many supervised , self- supervised , and contrastive learning frameworks . In Section 3 , we prove that learned representations in this family have an asymptotic property desirable for representation learning : equality up to a linear transformation . In Section 4 , we show that this family includes a number of highly performant models , state-of-the-art at publication for their problem domains , including CPC ( Oord et al. , 2018 ) , BERT ( Devlin et al. , 2018 ) , and GPT-2 and GPT-3 ( Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ) . Section 5 investigates the actually realizable regime of finite data and partial optimization , showing that representations learned by members of the identifiable model family approach equality up to a linear transformation as a function of dataset size , neural network capacity , and optimization progress . 2 MODEL FAMILY AND DATA DISTRIBUTION . The learned embeddings of a DNN are a function not only of the parameters , but also the network architecture and size of dataset ( viewed as a sample from the underlying data distribution ) . This renders any analysis in full generality challenging . To make such an analysis tractable , in this section , we begin by specifying a set of assumptions about the underlying data distribution and model family that must hold for the learned representations to be similar up to a linear transformation . These assumptions are , in fact , satisfied by a number of already published , highly performant models . We establish definitions in this section , and discuss these existing approaches in depth in Section 4 . Data Distribution We assume the existence of a generalized dataset in the form of an empirical distribution pD ( x , y , S ) over random variables x , y and S with the following properties : • The random variable x is an input variable , typically high-dimensional , such as text or an image . • The random variable y is a target variable whose value the model predicts . In case of object classification , this would be some semantically meaningful class label . However , in our model family , y may also be a high-dimensional context variable , such a text , image , or sentence fragment . • S is a set containing the possible values of y given x , so pD ( y|x , S ) > 0 ⇐⇒ y ∈ S. Note that the set of labels S is not fixed , but a random variable . This allows supervised , contrastive , and self-supervised learning frameworks to be analyzed together : the meaning of S encodes the task . For supervised classification , S is deterministic and contains class labels . For self-supervised pretraining , S contains randomly-sampled high-dimensional variables such as image embeddings . For deep metric learning ( Hoffer and Ailon , 2015 ; Sohn , 2016 ) , the set S contains one positive and k negative samples of the class to which x belongs . Canonical Discriminative Form Given a data distribution as above , a generalized discriminative model family may be defined by its parameterization of the probability of a target variable y conditioned on an observed variable x and a set S that contains not only the true target label y , but also a collection of distractors y′ : pθ ( y|x , S ) = exp ( fθ ( x ) > gθ ( y ) ) ∑ y′∈S exp ( fθ ( x ) > gθ ( y′ ) ) , ( 1 ) The codomain of the functions fθ ( x ) and gθ ( y ) is RM , and the domains vary according to modelling task . For notational convenience both are parameterized by θ ∈ Θ , but f and g may use disjoint parts of θ , meaning that they do not necessarily share parameters . With F and G we denote the function spaces of fθ and gθ respectively . Our primary domain of interest is when fθ and gθ are highly flexible function approximators , such as DNNs . This brings certain analytical challenges . In neural networks , different choices of parameters θ can result in the same functions fθ and gθ , hence the map Θ → F × G is many-to-one . In the context of representation learning , the function fθ is typically viewed as a nonlinear feature extractor , e.g. , the learned representation of the input data . While other choices meet the membership conditions for the family defined by the canonical form of Equation ( 1 ) , in the remainder , we will focus on DNNs in the remainder . We next present a definition of identifiability suitable for DNNs , and prove that members of the above family satisfy it under additional assumptions . 3 MODEL IDENTIFIABILITY . In this section , we derive identifiability conditions for models in the family defined in Section 2 . 3.1 IDENTIFIABILITY IN PARAMETER SPACE . Identifiability analysis answers the question of whether it is theoretically possible to learn the parameters of a statistical model exactly . Specifically , given some estimator θ′ for model parameters θ∗ , identifiability is the property that , for any { θ′ , θ∗ } ⊂ Θ , pθ′ = pθ∗ =⇒ θ′ = θ∗ . ( 2 ) Models that do not have this property are said to be non-identifiable . This happens when different values { θ′ , θ∗ } ⊂ Θ can give rise to the same model distribution pθ′ ( y|x , S ) = pθ∗ ( y|x , S ) . In such a case , observing an empirical distribution pθ∗ ( y|x , S ) , and fitting a model pθ′ ( y|x , S ) to it perfectly does not guarantee that θ′ = θ∗ . Neural networks exhibit various symmetries in parameter space such that there is almost always a many-to-one correspondence between a choice of θ and resulting probability function pθ . A simple example in neural networks is that one can swap the ( incoming and outgoing ) connections of two neurons in a hidden layer . This changes the value of the parameters , but does not change the network ’ s function . Thus , when representation functions fθ or gθ are parameterized as DNNs , equation 2 is not satisfiable . 3.2 IDENTIFIABILITY IN FUNCTION SPACE . For reliable and efficient representation learning , we want learned representations fθ from two identifiable models to be sufficiently similar for interchangeable use in downstream tasks . The most general property we wish to preserve among learned representations is their ability to discriminate among statistical patterns corresponding to categorical groupings . In the model family defined in Section 2 , the data and context functions fθ and gθ parameterize pθ ( y|x , S ) , the probability of label assignment , through a normalized inner product . This induces a hyperplane boundary , for discrimination , in a joint space of learned representations for data x and context y . Therefore , in the following , we will derive identifiability conditions up to a linear transformation , using a notion of similarity in parameter space inspired by Hyvärinen et al . ( 2018 ) . Definition 1 . Let L∼ be a pairwise relation on Θ defined as : θ′ L∼ θ∗ ⇐⇒ fθ ′ ( x ) = Afθ∗ ( x ) gθ′ ( y ) = Bgθ∗ ( y ) ( 3 ) whereA andB are invertible M ×M matrices . See Appendix B for proof that L∼ is an equivalence relation . In the remainder , we refer to identifiability up to the equivalence relation L∼ as L∼-identifiable or linearly identifiable . 3.3 LINEAR IDENTIFIABILITY OF LEARNED REPRESENTATIONS . We next present a simple derivation of the L∼-identifiability of members of the generalized discriminative family defined in Section 2 . This result reveals sufficient conditions under which a discriminative probabilistic model pθ ( y|x , S ) has a useful property : the learned representations of the input x and target random variables y for any two pairs of parameters ( θ′ , θ∗ ) are related as θ′ L∼ θ∗ , that is , fθ′ ( x ) = Afθ∗ ( x ) and gθ′ ( y ) = Bgθ∗ ( y ) . We first review the notation for the proof , which is introduced in detail in Section 2 . We then highlight an important requirement on the diversity of the data distribution , which must be satisfied for the proof statement to hold . We prove the result immediately after . Notation . The target random variables y , associated with input random variables x , may be class labels ( as in supervised classification ) , or they could be stochastically generated from datapoints x as , e.g. , perturbed image patches ( as in self-supervised learning ) . We account for this additional stochasticity as a set-valued random variable S , containing all possible values of y conditioned on some x . For brevity , we will use shorthands that drop the parameters θ : p′ : = pθ′ , p∗ : = pθ∗ , f∗ : = fθ∗ , f ′ : = fθ′ , g ′ : = gθ′ . Diversity condition . We assume that for any ( θ′ , θ∗ ) for which it holds that p′ = p∗ , and for any given x , by repeated sampling S ∼ pD ( S|x ) and picking yA , yB ∈ S , we can construct a set of M distinct tuples { ( y ( i ) A , y ( i ) B ) } Mi=1 such that the matrices L′ and L∗ are invertible , where L′ consists of columns ( g′ ( y ( i ) A ) − g′ ( y ( i ) B ) ) , and L ∗ consists of columns g∗ ( y ( i ) A ) − g∗ ( y ( i ) B ) , i ∈ { 1 , . . . , M } . See Section 3.4 for detailed discussion . Theorem 1 . Under the diversity condition , models in the family defined by Equation ( 1 ) are linearly identifiable . That is , for any θ′ , θ∗ ∈ Θ , and f∗ , f ′ , g∗ , g′ , p∗ , p′ defined as in Section 2 , p′ = p∗ =⇒ θ′ L∼ θ∗ . ( 4 ) To establish the result , we proceed by directly constructing an invertible linear transformation that satisfies Definition 1 . Consider yA , yB ∈ S. The likelihood ratios for these points p′ ( yA|x , S ) p′ ( yB |x , S ) = p∗ ( yA|x , S ) p∗ ( yB |x , S ) ( 5 ) are equal . Substituting our model definition from equation ( 1 ) , we find : exp ( f ′ ( x ) > g′ ( yA ) ) exp ( f ′ ( x ) > g′ ( yB ) ) = exp ( f∗ ( x ) > g∗ ( yA ) ) exp ( f∗ ( x ) > g∗ ( yB ) ) , ( 6 ) where the normalizing constants cancelled out on the left- and right-hand sides . Taking the logarithm , this simplifies to : ( g′ ( yA ) − g′ ( yB ) ) > f ′ ( x ) = ( g∗ ( yA ) − g∗ ( yB ) ) > f∗ ( x ) . ( 7 ) Note that this equation is true for any triple ( x , yA , yB ) for which pD ( x , yB , yB ) > 0 . We next collect M distinct tuples ( y ( i ) A , y ( i ) B ) so that by repeating Equation ( 7 ) M times and by the diversity condition noted above , the resulting difference vectors are linearly independent . We collect these vectors together as the columns of ( M ×M ) -dimensional matrices L′ and L∗ , forming the following system of M linear equations : L′ > f ′ ( x ) = L∗ > f∗ ( x ) . Since L′ and L∗ are invertible , we rearrange : f ′ ( x ) = ( L∗L′−1 ) > f∗ ( x ) . ( 8 ) Hence , f ′ ( x ) = Af∗ ( x ) where A = ( L∗L′−1 ) . This completes the first half of the proof . See Appendix C for the second half of the proof , which is similar , and handles the function g .

Summary: In this paper, the authors address the model identifiability in a general setting that can be adapted to several recent deep learning models (DNN supervised learning, CPC, BERT and GPT). Since model parameters (NN weights) are not identifiable, the authors hypothesize that vector f and g can be identifiably up to a linear transformation. Although the purpose of the work is appealing, there are some issues related to the current structure of the paper, proposed theory and its relationship to the provided experiments. See my detailed comments and questions below: