Masked Prediction: A Parameter Identiﬁability View
Anonymous Author(s)
AfﬁliationAddress email
Abstract
The vast majority of work in self-supervised learning have focused on assessingrecovered features by a chosen set of downstream tasks.
While there are severalcommonly used benchmark datasets, this lens of feature learning requires assumptions on the downstream tasks which are not inherent to the data distribution itself.
In this paper, we present an alternative lens, one of parameter identiﬁability: assuming data comes from a parametric probabilistic model, we train a self-supervisedlearning predictor with a suitable parametric form, and ask whether the parametersof the optimal predictor can be used to extract the parameters of the ground truthgenerative model.
Speciﬁcally, we focus on latent-variable models capturing sequential structures,namely Hidden Markov Models with both discrete and conditionally Gaussianobservations.
We focus on masked prediction as the self-supervised learning taskand study the optimal masked predictor.
We show that parameter identiﬁability isgoverned by the task difﬁculty, which is determined by the choice of data modeland the amount of tokens to predict.
Technique-wise, we uncover close connectionswith the uniqueness of tensor rank decompositions , a widely used tool in studyingidentiﬁability through the lens of the method of moments.
Introduction
Self-supervised learning (SSL) is a relatively new approach to unsupervised learning, where thelearning algorithm learns to predict auxiliary labels generated automatically from the data withouthuman annotators.
The hope is that with a properly designed prediction task, a successfully learnedpredictor would capture some knowledge about the underlying data.
While SSL has been enjoying arapid growth on the empirical front, theoretical understanding of why and when SSL works is stillnascent.
In no small part, this is because formalizing the desired guarantees seems challenging.
Forinstance, the focus of SSL has largely been on learning good features , which in practice has beenquantiﬁed by downstream performance on various benchmark datasets [Wang et al., 2018, 2019,
Deng et al., 2009, Zhai et al., 2019, Tamkin et al., 2021].
To provide theoretical underpinning to this,one needs to make extra assumptions on the relationship between the self-supervised prediction taskand the downstream tasks [Arora et al., 2019, Saunshi et al., 2020, HaoChen et al., 2021, Lee et al.
,2021a, Wang et al., 2021, Wei et al., 2021, Wen and Li, 2021].
While associating SSL with downstream supervised tasks is a useful perspective and has led to severalvery interesting theoretical results, we take a step back and revisit a more general goal of SSL, whichis to learn some informative functionals of the data distribution.
Naturally, the key question here iswhat functionals should be considered informative .
While downstream performance is a notable validchoice, in this work, we choose an alternative criterion that is meaningful even without referencingany downstream tasks.
The alternative lens we are interested in is whether the functionals of the data distribution extractedby the SSL predictors can be simply stitched together to obtain the data distribution itself, givenadditional side-information about the family from which the data distribution is drawn.
While thismight seem like a tall order, masked prediction based SSL algorithms (which is essentially whatpseudo-likelihood corresponds to) have classically been used for learning parametric graphical modelssuch as Ising models [Ravikumar et al., 2010, Bresler, 2015, Vuffray et al., 2016].
But can this bedone for broader classes of parametric models?
In this paper, we take a preliminary step towards this and ask the question of parameter identiﬁability :assuming the data comes from a ground truth parametric probabilistic model, can common selfsupervised tasks uniquely identify the parameters of the ground truth model?
More precisely, are theparameters of the model uniquely determined by the optimal predictor for the SSL task (Deﬁnition 1)?
An appeal of this identiﬁability perspective is that when a SSL task is sufﬁcient for parameteridentiﬁability, the model parameters can then be recovered straightforwardly from the parametersfrom the optimal SSL predictor.
Parameter identiﬁcation also has the desirable property of beingindependent of any downstream task.
A prior, it is unclear whether we can achieve such model parameter identiﬁability via self-supervisedtasks, since it requires recovering the full (parametric) generative model which is arguably moredifﬁcult than learning generic latent representations.
This work provides a positive answer for broadclasses of HMMs: we show that the commonly-used masked prediction task [Pathak et al., 2016,
Devlin et al., 2018, He et al., 2021, Lee et al., 2021a], wherein a model is trained to predict amasked-out part of a sample given the rest of the sample, can identify the parameters of a HMM.
As noted earlier, while such masked prediction for parameter learning has been applied in classicalsettings such as Ising models [Ravikumar et al., 2010, Bresler, 2015, Vuffray et al., 2016], the HMMsetup in this work is more challenging due to the presence of latent variables.
HMMs are also moresuitable for modeling practical sequential data, and have been commonly adopted in theoreticalanalyses as a clean proxy for languages [Wei et al., 2021, Xie et al., 2021].
Concretely, the two HMM models we consider in this work are 1) the classic HMM with discrete latentand discrete observables, and 2) a HMM variant with discrete latents and continuous observables thatare conditionally Gaussian given the latent, which we abbreviate as G-HMMs.
We show that:
• Parameter identiﬁability is governed by the difﬁculty of the masked prediction task.
The taskdifﬁculty is related to the amount of information provided by the combination of the model and theprediction task—where the difﬁculty can be increased by using a more complicated model, or bypredicting more tokens.
For instance, predicting the conditional mean of one token given anotherdoes not yield identiﬁability for a discrete HMM (Theorem 2), but does so when data comes froma G-HMM (Theorem 3).
Moreover, the identiﬁability in the latter case quite strongly leveragesstructural properties of the posterior of the latent variables (Section 3.1).• Tools for characterizing the uniqueness of tensor decompositions (e.g., Kruskal’s Theorem [Kruskal,1977, Allman et al., 2009]) can be leveraged to prove identiﬁability: For both HMM (Theorem 5)and G-HMM (Theorem 6), if we have predictors of the tensor product of tokens (e.g., E [ x 2 ⊗ x 3 | x 1 ] ),we can use the predictor output to construct a 3-tensor whose rank-1 components are uniquelydetermined and reveal the parameters of the model.
The rest of the paper is structured as follows.
Section 2 provides relevant deﬁnitions, preliminaries andassumptions.
Section 3 states the main results of this work.
Main proofs, including the identiﬁabilityproof via tensor decomposition, are provided in Section 4, with the rest deferred to the appendix.
We then discuss related works in Section 5.
Finally, we emphasize that this work is a ﬁrst-cut studyon this lens of parameter recovery for analyzing SSL tasks based on masked prediction, and ourencouraging results suggest interesting open directions in this thread of analyzing self-supervisedlearning via parameter recovery, which are discussed brieﬂy in the conclusion.
Setup
This work focuses on two classes of latent-variable sequence models.
The ﬁrst are fully discretehidden Markov models (HMMs), and the second are HMMs whose observables marginally followa mixtures of Gaussians with identity covariance.
We denote the observations and hidden statesrespectively by { x t } t ≥ 1 and { h t } t ≥ 1 for both classes.
The hidden states h 1 → h 2 → · · · form a Markov chain, and conditional on h t , the observable x t is independent of all other variables.
Throughout, we refer to { x t } t ≥ 1 as tokens, following the nomenclature from language models.
2.1 Models
Discrete Hidden Markov Model
We ﬁrst describe the parameterization of the standard HMMswith discrete latents and observations.
Let X := { 1 , . . .
, d } =
[ d ] denote the observation space, andlet H :=
[ k ] be the state space.
1 The parameters of interest are the transition matrix T ∈
R k × k andthe emission matrix O ∈ R d × k , deﬁned in the standard way as
[Equation]
Conditionally-Gaussian HMM (G-HMM)
We next describe the parameterization ofconditionally-Gaussian HMMs (G-HMMs) .
The state space H :=
[ k ] is the same as in the previous case, while the observation space is now continuous with X := R d .
The parameters of interestare T ∈ R k × k , the transition matrix, and { µ i } i ∈
[ k ] ⊂ R d , the means of the k identity-covariance
Gaussians.
Precisely,
[Equation]
We use M :=
[ µ 1 , . .
. , µ k ] ∈ R
d × k to denote the matrix whose columns are the Gaussian means.
2.2 Masked prediction tasks
We are interested in the (regression) task of predicting one or more “masked out” tokens as a functionof another observed token, with the goal of minimizing expected squared loss under a distributiongiven by an HMM or G-HMM (equation 1).
In the case of the discrete HMMs, we will speciﬁcallybe predicting the one-hot encoding vectors of the observations.
Thus, both for HMM and G-HMM,predicting a single token will correspond to predicting a vector.
For notational convenience, wewill simply associate the discrete states or observations via their one-hot vectors { e 1 , e 2 , . . . } in theappropriate space and interchangeably write h
=
i or h = e i , and similarly for x .
For the task ofpredicting the tensor product of (one-hot encoding vectors of) tokens ⊗ τ ∈T x τ from another tokenx t (where T is some index set and t / ∈ T ), the optimal predictor with respect to the squared losscalculates the conditional expectation:
[Equation]
where “vec" returns the vectorized form of a tensor.
We use the shorthand “ ⊗ τ ∈T x τ | x t ” to refer to this prediction task.
For instance, consider the case ofpredicting x 2 given x 1 under the HMM with parameters ( O, T ) .
The optimal predictor, denoted byf 2 | 1 , can be written in terms of ( O, T ) as
[Equation]
= (cid:88) i ∈
[ k ] (cid:88) j ∈
[Equation]
k ]
[Equation]
(cid:124)
[Equation]
(cid:125) :=[
[Equation]
= (cid:88) i ∈
[ k ] (cid:88)
j ∈ [ k ]
O i T ij
O
[Equation]
Here φ : R d → R k denotes the posterior distribution of a hidden state h t given the correspondingobservation x t , i.e., φ ( x t ) =
E [ h t | x t ] .
Our goal is to study the parameter identiﬁability from the prediction tasks, when the predictors havethe correct parametric form.
Formally, we deﬁne identiﬁability from a prediction task as follows:
[Equation]
[Equation] [Equation]
Deﬁnition 1 (Identiﬁability from a prediction task, HMM) .
A prediction task sufﬁces for identiﬁabilityif, for any two HMMs with parameters ( O, T ) and ( ˜ O, ˜ T ) , equality of their optimal predictors forthis task implies that there is a permutation matrix Π
such that O =
˜ O Π and T
= Π (cid:62) ˜ T Π .
In other words, the mapping from (the natural equivalence classes of) HMM distributions to optimalpredictors for a task is injective, up to a permutation of the hidden state labels.
By identiﬁabilityfrom a collection of prediction tasks, we refer to the injectiveness of the mapping from HMMdistributions to the collections of optimal predictors for the tasks.
Identiﬁability for G-HMMs isdeﬁned analogously with O, ˜ O changed to M, ˜ M .
2.3 Assumptions
We now state the assumptions used in our results.
The ﬁrst assumption is that the transition matricesof the HMMs are doubly stochastic.
Assumption 1 (Doubly stochastic transitions) .
The transition matrix T is doubly stochastic, and themarginal distribution of the initial hidden state h 1 is stationary with respect to T .
This assumption guarantees that the stationary distribution of the latent distribution is uniform for anyt , and the transition matrix for the reversed chain is simply T (cid:62) .
Moreover, this assumption reducesthe parameter space and hence will make the non-identiﬁability results stronger.
We require the following conditions on the parameters for the discrete HMM:
Assumption 2 (Non-redundancy, discrete HMM) .
Every row of O is non-zero.
Assumption 2 can be interpreted as requiring each token to have a non-zero probability of beingobserved, which is a mild assumption.
We also require the following non-degeneracy condition:
Assumption 3 (Non-degeneracy, discrete HMM) .
rank ( T ) = rank ( O ) =
k ≤ d .
Note that Assumption 3 only requires the parameters to be non-degenerate, rather than have singularvalues bounded away from 0.
The reason is that this work will focus on population level quantitiesand make no claims on ﬁnite sample behaviors or robustness.
For G-HMM, we similarly require the parameters to be non-degenerate:
Assumption 4 (Non-degeneracy, G-HMM) .
rank ( T ) = rank ( M ) =
k ≤ d .
Moreover, we assume that the norms of the means are known and equal:
Assumption 5 (Equal norms of the means) .
For each i ∈ [ k ] , µ i is a unit vector.
Assumptions 1-4 are fairly standard [see, e.g., Anandkumar et al., 2012]; in particular, Assumption 3,4 are required to enable efﬁcient learning, since learning degenerate HMMs can be computationallyhard [Mossel and Roch, 2005].
Assumption 5 may be an artifact of our proofs, and it would beinteresting to relax in future work.
Our notion of identiﬁability from a prediction task (or a collection of prediction tasks) will restrictattention to HMMs satisfying Assumptions 1, 2, 3 and G-HMMs satisfying Assumptions 1, 4, 5.
2.4 Uniqueness of tensor rank decompositions
Some of our identiﬁability results rely on the uniqueness of tensor rank-1 decompositions [Hitchcock,1927].
An order- t tensor (or t -tensor ) is an t -way multidimensional array; a matrix is a 2-tensor.
The tensor rank of a tensor W is the minimum number R such that W can be written as a sum of Rrank-1 tensors.
That is, if a t -tensor W has rank- R , it means that W = (cid:80) i ∈
[ r ] ⊗ j ∈ [ t ] U ( j ) i for somematrices U ( j ) ∈ R n j × r , where U ( j ) i denotes the i th column of matrix U ( j ) .
In this work, we only need to work with 3-tensors of the form W = (cid:80) i ∈
[ R ] A i ⊗ B i ⊗ C i for somematrices A ∈ R n 1 × R , B ∈ R n 2 × R , C ∈ R n 3 × R , as 3-tensors will sufﬁce for identiﬁability in all ofour settings of interest.
5 A classic work by Kruskal [1977] gives a sufﬁcient condition under which
A, B, C can be recovered up to column-wise permutation and scaling.
The condition is stated in terms of the Kruskal rank , which is the maximum number r such that every r columns of the matrixare linearly independent.
Let k A denote the Kruskal rank of matrix A , then:
Proposition 1 (Kruskal’s theorem, Kruskal [1977]) .
The components A, B, C of a 3-tensor W := 168 (cid:80) i ∈
[ R ] A i ⊗ B i ⊗ C
i are identiﬁable up to a shared column-wise permutation and column-wisescaling if k A + k B
+ k C ≥ 2 R + 2 .
We note that this work focuses on identiﬁability results rather than providing an algorithm or samplecomplexity bounds, though the proofs can be adapted into algorithms [see, e.g., Harshman, 1970]under slightly more restrictive conditions (which will be satisﬁed by all of our identiﬁability results).
Identiﬁability from masked prediction tasks
We now present the main (non-)identiﬁability results, and show that the combination of the datagenerative models and the prediction task directly impacts the sufﬁciency of identiﬁability.
3.1 Pairwise prediction
We begin with the simplest prediction task: namely predicting one token from another.
We refer tosuch tasks as pairwise prediction tasks .
For HMMs, this task fails to provide parameter identiﬁability:
Theorem 2 (Nonidentiﬁability of HMM from predicting x t | x 1 ) .
For any t ∈ Z , t ≥ 2 , there exists apair of HMM distributions with parameters ( O, T ) and ( ˜ O, ˜ T ) , each satisfying Assumptions 1, 2 and3, such that the optimal predictors for the task x t | x 1 are the same under each distribution, but thereis no permutation matrix Π ∈ R k × k such that ˜
O
= O Π and ˜ T
= Π (cid:62) T Π are both satisﬁed.
Theorem 2 follows from the fact that the optimal predictor has the form of a product of (stochastic)matrices, and generally, one cannot uniquely recover matrices from their product sans additionalconditions
[Donoho and Elad, 2003, Candes et al., 2006, Spielman et al., 2012, Arora et al., 2014,
Georgiev et al., 2005, Aharon et al., 2006, Cohen and Gillis, 2019].
Speciﬁcally, by equation 1, theoptimal predictor is f ( x 1 ) =
E [ x t | x 1 ] =
OT t − 1 φ ( x 1 ) (where φ ( x 1 ) :
=
E [ h 1 | x t ] is the posterior).
When t = 2 , we can ﬁnd a non-permutation matrix R such that ˜ O =
OR ,
˜ T
= R (cid:62) TR give the samepredictor as O, T .
For t > 2 , even if ˜ O
=
O , we show that the matrix power T t − 1 is not identiﬁable:
Claim 1 (Nonidentiﬁability of matrix powers) .
For any t ∈ Z , t ≥ 2 , there exist stochastic matrices
T, ˜ T satisfying Assumption 1, 3, such that T (cid:54) =
˜ T and T t =
˜ T t .
On the other hand, pairwise prediction actually does sufﬁce for identiﬁability for G-HMM:
Theorem 3 (Identiﬁability of G-HMM from predicting x 2 | x 1 ) .
Under Assumption 1, 4, and 5, if theoptimal predictors for the task x 2 | x 1 under the G-HMM distributions with parameters ( M, T ) and( ˜ M, ˜ T ) are the same, then ( M, T )
=
( ˜ M, ˜ T ) up to a permutation of the hidden state labels.
Comparing Theorem 2 and 3 shows that the speciﬁc parametric form of the generative model matters.
Note that HMM and G-HMM have a similar form when conditioning on the latent variable; that is,with t = 2 , the predictor conditioned on the hidden variable h 2 is P ( x 2 | h 2 = i ) =
OT i for HMM,and P ( x 2 | h 2 = i ) =
MT i for G-HMM.
The salient difference between these two setups lies in theposterior function: while the posterior function for HMM is linear in the observable, the posteriorfunction for G-HMM is more complicated and “reveals” more information about the parameter.
To formalize the above intuition, ﬁrst recall that the GHMM posterior has entries [ φ ( x t )]
i =
[Equation]
(cid:1) , ∀ i ∈ [ k ] .
We will show that for G-HMM, even matching the posteriorfunction nearly sufﬁces to identify M : if M, ˜ M parameterize two posterior functions φ, ˜ φ whereφ =
˜ φ , then up to a permutation, ˜ M must be equal to either M or a unique (and somewhat special)transformation of M .
The next step is to further exclude the (special) transformation, which isachieved using the constraint that T, ˜ T are stochastic matrices.
The ﬁrst step of the proof sketch iscaptured by following lemma:
Lemma 1.
For d ≥ k ≥ 2 , under Assumption 4, 5, φ = ˜ φ implies ˜ M = M or ˜ M =
HM , where His a Householder transformation of the form H :
=
I d − 2ˆ v ˆ v (cid:62) ∈ R
d × d , with ˆ v :=
( M † ) (cid:62) 1 √ 1 (cid:62) M † ( M † ) (cid:62) 1 .
To provide some geometric intuition about how H acts on M , note that ˆ v is a unit vector in thecolumn space of M and perpendicular to the afﬁne hull of A :=
{ µ i : i ∈ [ k ] } , which means ˆ v (cid:62) µ i isthe same for all i ∈ [ k ] .
As a result, ˜ M
= [˜ µ 1 , ..., ˜ µ k ] =
[ Hµ 1 , ..., Hµ k ] =
M − 2(ˆ v (cid:62) µ 1 )[ˆ v, ..., ˆ v ]is a translation of M along the direction of ˆ v , such that the translated points { ˜ µ i } i ∈ [ k ] lie on theopposite side of the origin.
It is non-trivial to argue that HM is the only solution (other than Mitself) that preserves φ , and we defer the proof to Appendix A.1.2.
It is, however, easy to see that
HM indeed results in a matching posterior, whose sufﬁcient conditions are 1)
˜ M is a translation of
M , and 2) (cid:107) ˜ µ i (cid:107) 2 − (cid:107) µ i (cid:107) 2 is the same for all i ∈ [ k ] .
˜ M :
=
HM indeed satisﬁes both conditions.
Proof sketch for Theorem 3 : We ﬁrst show that if M, T and ˜ M, ˜ T produce the same predictor, thentheir posterior function must be equal up to a permutation (Lemma 2).
We can then apply Lemma1 to recover M up to a permutation and a Householder transformation H .
Then, we show that if˜ M = HM , then the corresponding ˜ T must have negative entries and thus would not be a validstochastic matrix.
Hence it must be that ˜ M, M are equal up to permutation.
Finally, by way of remarks, another way to think of the difference between the two setups is thatfor HMM, P ( x 2 | x 1 ) is a mixture of categorical distributions, which itself is also a categoricaldistribution.
This also implies that the nonidentiﬁability from pairwise prediction in the HMM casecannot be resolved by changing the squared loss to another proper loss function.
On the other hand,for G-HMM, the conditional distribution P ( x 2 | x 1 ) is a mixture of Gaussians, which is well knownto be identiﬁable.
In fact, if we were given access to the entire conditional distribution P ( x 2 | x 1 )(instead of just the conditional mean), it is even easier to prove identiﬁability for G-HMM.
Thoughthis is already implied from identiﬁability from the conditional means, we provided a (much simpler)proof in Appendix A.3 assuming access to the full conditional distribution.
3.2 Beyond pairwise prediction
The conclusion from Theorem 2 is that a single pairwise prediction task does not sufﬁce for identiﬁability on HMMs.
The next question is then: can we modify the task to obtain identiﬁability?
A naturalidea is to force the model to “predict more”, and one straightforward way to do so is to combinemultiple pairwise prediction tasks.
It turns out that this does not resolve the nonidentiﬁability issue,as we can show that the parameters are not identiﬁable even when considering all possible pairwisetasks involved 3 (adjacent) tokens:
Theorem 4 (Nonidentiﬁability of HMM from all pairwise predictions on 3 tokens) .
There exists apair of HMM distributions with parameters ( O, T ) and ( ˜ O, ˜ T ) , each satisfying Assumptions 1, 2and 3, and also
˜ O
(cid:54)
=
O , such that, for each of the tasks x 2 | x 1 , x 1 | x 2 , x 3 | x 1 , and x 1 | x 3 , the optimalpredictors are the same under each distribution.
We brieﬂy remark that the reason for only considering adjacent time steps is that when the tokensare at least two time steps apart, matching predictors only matches powers of the transition matrices,which in general does not ensure the transition matrices themselves are matched as shown in Claim 1.
For the intuition of the nonidentiﬁability result in Theorem 4, recall that the limitation of pairwisepredictions on HMMs comes from non-uniqueness of matrix factorization.
While adding additionalpairwise prediction tasks introduces more equations on the product of matrices, these equations arehighly dependent, and the proof works by providing counterexamples that can simultaneously satisfyall these equations.
The above intuition leads to another way of forcing the model to “predict more”, that is, to increasethe number of predicted toens.
The hope is that doing so results in equations on tensors—as opposedto matrices— for which there is a lot of classical machinery delineating tensors for which the rank-1decomposition is unique, as discussed in Section 2.4.
This intuition proves to be true and we showthat increase the number from 1 to 2 already sufﬁces for identiﬁability:
Theorem 5 (Identiﬁability from masked prediction on three tokens, HMM) .
Let ( t 1 , t 2 , t 3 ) be anypermutation of (1 , 2 , 3) , and consider the prediction task x
t 2 ⊗ x t 3 | x t 1 .
Under Assumption 1, 2, 3,if the optimal predictors under the HMM distributions with parameters ( O, T ) and ( ˜ O, ˜ T ) are thesame, then ( O, T ) =
( ˜ O, ˜ T ) up to a permutation of the hidden state labels.
Compared to prior results on identiﬁability from third order moments [Allman et al., 2009, Anandkumar et al., 2012, 2014], the difﬁculty in our setup is that we only have access to the conditional2-tensors (i.e. matrices) given by the predictors.
The proof idea is to construct a third-order tensor bylinearly combining the conditional 2-tensors for each possible value of the token being conditionedon, such that Kruskal’s theorem applies and gives identiﬁability.
Note, importantly, that the weightsfor the linear combination cannot depend on the marginal probabilities of the token being conditionedon, since we do not have access to these marginals, and it is unclear whether we could extract uniquemarginals given the conditional probabilities we are predicting.
Thus, the above theorem cannot besimply derived from results showing parameter identiﬁability from the 3rd order moments.
It can be show that this tensor decomposition argument can also be applied to G-HMM, with the helpof Lemma 1.
We leave the details to Theorem 6 in Appendix A.
Proofs
We now discuss proofs for some of the main results.
Section 4.1 proves the identiﬁability of HMMparameters from the task of predicting two tokens (Theorem 5) using ideas from tensor decomposition,and Section 4.2 shows the identiﬁability proof of pairwise prediction on G-HMM.
The rest of theproofs are deferred to the appendix.
4.1 Proof of Theorem 5: identiﬁability of predicting two tokens for HMM
There are three cases for the two-token prediction task, i.e. 1) x 2 ⊗ x 3 | x 1 , 2) x 1 ⊗ x 3 | x 2 , and 3)x 1 ⊗ x 2 | x 3 .
We will prove for the ﬁrst two cases, as the third case is proved the same way as theﬁrst case by symmetry.
In all cases, the idea is to use the predictor to construct a 3-tensor whosecomponents are each of rank- k , so that applying Kruskal’s theorem gives identiﬁability.
Case 1, x 2 ⊗ x 3 | x 1 : O, T and ˜ O, ˜ T producing the same predictor means f 2 ⊗ 3 | 1 ( x 1 ) :
= E [ x 2 ⊗x 3 | x 1 ] = ˜
E [ x 2 ⊗ x 3 | x 1 ] :=
˜ f 2 ⊗ 3 | 1 ( x 1 ) , where E , ˜ E are parameterized by the corresponding parameters.
Let X := { e i : i ∈
[ d ] } , and consider the following 3-tensor:
[Equation]
where O i denotes the i th column of O , and similarly for ( OT ) i .
Note that W can also be written as
[Equation]
We want to apply Kruskal’s theorem for identiﬁability.
In particular, we will show that each componentin equation 2 forms a matrix of Kruskal rank k the second and third components clearly satisfy thiscondition by Assumption 3.
For the ﬁrst component, recall that φ ( x ) =
O (cid:62) x (cid:107) O (cid:62) x [Equation] 1 and write a i as
[Equation]
Putting a i into a matrix form, we get A :=
[ a 1 , ..., a k ] = diag (cid:0) [1 / (cid:107) ( e ( d ) j ) (cid:62) O (cid:107) 1 ] j ∈
d ] (cid:1) OT (cid:62) , 7 whichis of rank k by Assumption 3.
Hence components W are all of Kruskal rank k , and columns of OT, Oare identiﬁed up to column-wise permutation and scaling by Kruskal’s theorem.
The indeterminacyin scaling is further removed noting that columns of O, T need to sum up to 1.
Lastly, T is recoveredas T
= O † OT .
Case 2, x 1 ⊗ x 3 | x 2 : The optimal predictor for the task of predicting x 1 , x 3 given x 2 takes the form
[Equation]
Similarly as the previous case, we would like to construct a 3-tensor whose components can beuniquely determined by Kruskal’s theorem.
Let X be the same as before, and consider the 3-tensor
[Equation]
where the ﬁrst component can be simpliﬁed to
[Equation]
The matrix A :=
[ a 1 , ..., a k ]
=
D − 1 O is of rank k , hence we can identify (up to permutation) columnsof each component of W by Kruskal’s theorem.
This means if O, T and ˜ O, ˜ T produce the samepredictor, then we have OT =
˜
O ˜ T , OT (cid:62) =
˜ O ˜ T (cid:62) , and that O, ˜ O are matched up to a scaling of rows(i.e.
D − 1 ).
Next, to determine D , note that T, ˜ T are doubly stochastic by Assumption 1, which meansthe all-one vector 1 ∈
R k satisﬁes T 1 = ˜ T 1 = 1 .
Hence ˜ O ˜ T 1 = OT 1 =
O 1 =
[ (cid:107) O (cid:62) j (cid:107) 1 ] j ∈
[ d ] .
Wecan then compute D as D = diag
( OT 1 ) , and recover O as O = DA .
Finally, T is also recoveredsince ˜ O ˜
T
= O ˜ T =
OT ⇒ ˜
TT − 1 =
I k ⇒ ˜
T = T .4.
Proof of Theorem 3: identiﬁability of predicting x 2 given x 1 for G-HMM
For G-HMM, the predictor for x 2 given x 1 is parameterized as f 2 | 1 ( x 1 ) =
E [ x 2 | x 1 ] = MTφ ( x 1 ) .
If M, T and ˜ M, ˜ T produce the same predictor, then
[Equation]
Let R :=
( ˜ M ˜ T ) † (
MT ) ∈
R k × k , then ˜ φ ( x ) =
Rφ ( x ) .
The following lemma (proof deferred to
Appendix A.1) says that φ, ˜ φ must then be equal up to a permutation of coordinates:
Lemma 2.
If there exists a non-singular matrix R ∈ R k × k such that φ ( x ) =
R ˜ φ ( x ) , ∀ x ∈ R d , then
R must be a permutation matrix.
Combined with Lemma 1, we have ˜ M is equal to (up to a permutation) either M or HM , where His the Householder reﬂection given in Lemma 1.
The remaining step is to show that HM can be ruled out by requiring ˜ T to be a stochastic matrix.
Note that matching both the predictor and the posterior function means we also have ˜ M ˜ T = MT ,or ˜ T
= ( ˜ M † M ) T .
Recall that H :
= I d − 2ˆ v ˆ v (cid:62) for ˆ v = ( M † ) (cid:62) 1 √ 1 (cid:62) M † ( M † ) (cid:62) 1 .
When ˜ M = H ˜ M , thecolumn sum of ˜ M † M is
[Equation]
[Equation]
This means the column sum of ˜ T is 1 (cid:62) ˜ T = 1 (cid:62) ( ˜ M † M ) T = − 1 (cid:62) T = − 1 (cid:62) , which violates theconstraint that ˜ T should be a stochastic matrix with positive entries and column sum 1.
Hence it mustbe that M =
˜
M and hence also T = ˜ T (up to permutation), proving the theorem statement.
Related works
Self-supervised learning On the empirical side, self-supervised methods have gained a greatamount of popularity across many domains, including language understanding [Mikolov et al., 2013,
Vaswani et al., 2017, Devlin et al., 2018], visual understanding [Doersch et al., 2015, Pathak et al.
,2016], and distribution learning [Gutmann and Hyvärinen, 2010, Gao et al., 2020].
Classic ideas suchas contrastive learning [Hadsell et al., 2006, Gutmann and Hyvärinen, 2010, Dosovitskiy et al., 2014]and masked prediction [Mikolov et al., 2013] remain powerful in their modern realizations [Hénaffet al., 2019, Chen et al., 2020b, Devlin et al., 2018, Radford et al., 2019, Chen et al., 2020a, He et al.
,2021], pushing the state of the art performance and even surpassing supervised pretraining in variousaspects [Lee et al., 2021b, Liu et al., 2021].
On the theoretical front, there have been analyses on both masked predictions [Lee et al., 2021a]and contrastive methods [Arora et al., 2019, Tosh et al., 2020a,b, Wang and Isola, 2020, HaoChenet al., 2021, Wen and Li, 2021], with a focus on characterizing the quality of the learned featuresfor downstream tasks [Saunshi et al., 2020, Wei et al., 2021].
These approaches usually rely onquite strong assumptions to tie the self-supervised learning objective to the downstream tasks ofinterest.
In contrast, our work takes the view of parameter identiﬁability, for which there is no needfor downstream assumptions but instead the speciﬁc parametric form is key.
Note also that while theparameter recovery lens is a new contribution of our work, Wen and Li [2021] argue (as a side-productof their analysis) that some aspects of a generative model are recovered in their setup.
Their datamodel, however, is substantially different from ours and has very different identiﬁability properties(owing to its basis in sparse coding).
Latent variable models and tensor methods Latent variable models have been widely studied inthe literature.
One important area of research is independent component analysis (ICA), where thedata are assumed to be given as a transformation (mixing) of unknown independent sources which
ICA methods aim to identify.
In nonlinear ICA data models, both the sources and the mixing functionare generally not identiﬁable.
Hyvarinen and Morioka [2016, 2017] however, under some additionalassumptions (on the dependency structure of the different time steps, among other things), providesome identiﬁability results on the sources.
Unlike our setup, the mixing function in these models isdeterministic and also not the object of recovery.
More related to this work is the line of work on learning latent variable models with tensor methods.
Speciﬁc to learning HMMs, Mossel and Roch [2005] and Anandkumar et al.
[2012, 2014] providealgorithms based on third-order moments.
A major difference between these prior works on tensormethods and ours is that previous results operate on joint moments, while the results in this work arebased on conditional moments given by the optimal predictors for the masked tokens.
Conclusion
In this work, we take a model parameter identiﬁability view of self-supervised learning, which offersa complementary perspective to the current focus of feature learning for downstream performance.
By analyzing the masked prediction task in the setup of HMMs and its conditionally-Gaussian variant
G-HMM, we showed that parameter recovery is determined by the task difﬁculty, which can be tunedby both changing the parametric form of the data generative model, and by changing the maskedprediction task.
We emphasize that this is a ﬁrst-cut effort in the research program of analyzing SSL through the lensof model identiﬁability; we aim to build on this foundation to extend our analyses from HMMs tomore complicated latent sequence and latent variable models.
We also note that we have focusedhere on population analyses, and model identiﬁability.
It would be of interest to build off this todevelop and analyze the corresponding ﬁnite-sample learning algorithms for parametric generativemodels given SSL tasks, with sample complexity results, both in the realizable case, as well as inthe agnostic case where we have model mis-speciﬁcation.
Given the use of conditional MLEs andregressions in SSL, and the natural robustness of these to model mis-speciﬁcations, we conjecturethat these approaches will be much more robust when compared to say spectral methods.
Overall, we hope this work on an alternative lens to analyze SSL inspires further research.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] The limitations and future
directions are discussed in the conclusion section.
(c) Did you discuss ctions any potential negative societal impacts of your work?
[N/A]
This work is purely theoretical.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] Assumptions are listed in Section 2. together with a discussion on their implications
(b) Did you include complete proofs of all theoretical results?
[Yes] Proofs for Theorem 3 and 5 are presented in Section 4.
Missing proofs are provided in Appendix A forGHMMs, and Appendix B for HMMs.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[N/A]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[N/A]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[N/A]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[N/A]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N [Table] A [Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]