Masked Prediction: A Parameter Identiﬁability View
Abstract
The vast majority of work in self-supervised learning have focused on assessing recovered features by a chosen set of downstream tasks.
While there are several commonly 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-supervised learning predictor with a suitable parametric form, and ask whether the parameters of the optimal predictor can be used to extract the parameters of the ground truth generative model.
Speciﬁcally, we focus on latent-variable models capturing sequential structures, namely Hidden Markov Models with both discrete and conditionally Gaussian observations.
We focus on masked prediction as the self-supervised learning task and study the optimal masked predictor.
We show that parameter identiﬁability is governed by the task difﬁculty, which is determined by the choice of data model and the amount of tokens to predict.
Technique-wise, we uncover close connections with the uniqueness of tensor rank decompositions , a widely used tool in studying identiﬁability through the lens of the method of moments.
Introduction
Self-supervised learning (SSL) is a relatively new approach to unsupervised learning, where the learning algorithm learns to predict auxiliary labels generated automatically from the data without human annotators.
The hope is that with a properly designed prediction task, a successfully learned predictor would capture some knowledge about the underlying data.
While SSL has been enjoying a rapid growth on the empirical front, theoretical understanding of why and when SSL works is still nascent.
In no small part, this is because formalizing the desired guarantees seems challenging.
For instance, the focus of SSL has largely been on learning good features , which in practice has been quantiﬁ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 task and 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 several very interesting theoretical results, we take a step back and revisit a more general goal of SSL, which is to learn some informative functionals of the data distribution.
Naturally, the key question here is what functionals should be considered informative .
While downstream performance is a notable valid choice, in this work, we choose an alternative criterion that is meaningful even without referencing any downstream tasks.
The alternative lens we are interested in is whether the functionals of the data distribution extracted by the SSL predictors can be simply stitched together to obtain the data distribution itself, given additional side-information about the family from which the data distribution is drawn.
While this might seem like a tall order, masked prediction based SSL algorithms (which is essentially what pseudo-likelihood corresponds to) have classically been used for learning parametric graphical models such as Ising models [Ravikumar et al., 2010, Bresler, 2015, Vuffray et al., 2016].
But can this be done 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 the parameters 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 parameter identiﬁability, the model parameters can then be recovered straightforwardly from the parameters from the optimal SSL predictor.
Parameter identiﬁcation also has the desirable property of being independent of any downstream task.
A priori, it is unclear whether we can achieve such model parameter identiﬁability via self-supervised tasks, since it requires recovering the full (parametric) generative model which is arguably more difﬁcult than learning generic latent representations.
This work provides a positive answer for broad classes 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 a masked-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 classical settings such as Ising models [Ravikumar et al., 2010, Bresler, 2015, Vuffray et al., 2016], the HMM setup in this work is more challenging due to the presence of latent variables.
HMMs are also more suitable for modeling practical sequential data, and have been commonly adopted in theoretical analyses 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 latent and discrete observables, and 2) a HMM variant with discrete latents and continuous observables that are 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 task difﬁculty is related to the amount of information provided by the combination of the model and the prediction task—where the difﬁculty can be increased by using a more complicated model, or by predicting more tokens.
For instance, predicting the conditional mean of one token given another does not yield identiﬁability for a discrete HMM (Theorem 2), but does so when data comes from a G-HMM (Theorem 3).
Moreover, the identiﬁability in the latter case quite strongly leverages structural 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 uniquely determined and reveal the parameters of the model.
The rest of the paper is structured as follows.
Section 2 provides relevant deﬁnitions, preliminaries and assumptions.
Section 3 states the main results of this work.
Main proofs, including the identiﬁability proof 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 study on this lens of parameter recovery for analyzing SSL tasks based on masked prediction, and our encouraging results suggest interesting open directions in this thread of analyzing self-supervised learning 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 discrete hidden Markov models (HMMs), and the second are HMMs whose observables marginally follow a mixtures of Gaussians with identity covariance.
We denote the observations and hidden states respectively 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 HMMs with discrete latents and observations.
Let X := { 1 , . . .
, d } =
[ d ] denote the observation space, and let H :
=
[ k ] be the state space.
1 The parameters of interest are the transition matrix T ∈ R k × k and the emission matrix O ∈ R d × k , deﬁned in the standard way as
[Equation]
Conditionally-Gaussian HMM (G-HMM)
We next describe the parameterization of conditionally-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 interest are 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 function of another observed token, with the goal of minimizing expected squared loss under a distribution given by an HMM or G-HMM (equation 1).
In the case of the discrete HMMs, we will speciﬁcally be 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, we will simply associate the discrete states or observations via their one-hot vectors { e 1 , e 2 , . . .
} in the appropriate space and interchangeably write h
=
i or h = e i , and similarly for x .
For the task of predicting the tensor product of (one-hot encoding vectors of) tokens ⊗ τ ∈T x τ from another token x t (where T is some index set and t / ∈ T ), the optimal predictor with respect to the squared loss calculates 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 of predicting x 2 given x 1 under the HMM with parameters ( O, T ) .
The optimal predictor, denoted by f 2 | 1 , can be written in terms of ( O, T ) as
[Equation]
1 = j )
[Equation]
x
[Equation] [Equation]
[Equation] [Equation]
] j .
Here φ : R d → R k denotes the posterior distribution of a hidden state h t given the corresponding observation 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 have the correct parametric form.
Formally, we deﬁne identiﬁability from a prediction task as follows:
Deﬁnition 1 (Identiﬁability from a prediction task, HMM) .
A prediction task sufﬁces for identiﬁability if, for any two HMMs with parameters ( O, T ) and ( ˜ O, ˜ T ) , equality of their optimal predictors for this 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 optimal predictors for a task is injective, up to a permutation of the hidden state labels.
By identiﬁability from a collection of prediction tasks, we refer to the injectiveness of the mapping from HMM distributions to the collections of optimal predictors for the tasks.
Identiﬁability for G-HMMs is deﬁ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 matrices of the HMMs are doubly stochastic.
Assumption 1 (Doubly stochastic transitions) .
The transition matrix T is doubly stochastic, and the marginal 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 any t , and the transition matrix for the reversed chain is simply T (cid:62) .
Moreover, this assumption reduces the 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 being observed, 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 singular values bounded away from 0.
The reason is that this work will focus on population level quantities and 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 computationally hard [Mossel and Roch, 2005].
Assumption 5 may be an artifact of our proofs, and it would be interesting to relax in future work.
Our notion of identiﬁability from a prediction task (or a collection of prediction tasks) will restrict attention 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 R 3 For discrete
[Equation]
(cid:1)
[Table]
does not need to be indexed by t due to the stationarity assumption in Section 2.3.
Assumption 5 can be changed to (cid:107) µ i (cid:107) 2 = c for all i ∈
[ k ] , for any other ﬁxed number c > 0 .
rank-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 some matrices 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 some matrices 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 of our 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 matrix are 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 :=
(cid:80)
i ∈
[ R ] A i ⊗ B i ⊗ C
i are identiﬁable up to a shared column-wise permutation and column-wise scaling 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 sample complexity 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 data generative 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 to such 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 a pair of HMM distributions with parameters ( O, T ) and ( ˜ O, ˜ T ) , each satisfying Assumptions 1, 2 and 3, such that the optimal predictors for the task x
t | x 1 are the same under each distribution, but there is 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 additional conditions [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, the optimal 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 same predictor 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 the optimal 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 the posterior function: while the posterior function for HMM is linear in the observable, the posterior function 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 posterior function 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 is achieved using the constraint that T, ˜ T are stochastic matrices.
The ﬁrst step of the proof sketch is captured by following lemma:
Lemma 1.
For d ≥ k ≥ 2 , under Assumption 4, 5, φ = ˜ φ implies ˜ M = M or ˜ M =
HM , where H is a Householder transformation of the form H :
=
I d − 2ˆ v ˆ v (cid:62) ∈ R
d × d , with ˆ v :=
[Equation]
To provide some geometric intuition about how H acts on M , note that ˆ v is a unit vector in the column space of M and perpendicular to the afﬁne hull of A :=
{ µ i : i ∈ [ k ] } , which means ˆ v (cid:62) µ i is the 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 the opposite side of the origin.
It is non-trivial to argue that HM is the only solution (other than M itself) 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, then their posterior function must be equal up to a permutation (Lemma 2).
We can then apply Lemma 1 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 valid stochastic 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 that for HMM, P ( x 2 | x 1 ) is a mixture of categorical distributions, which itself is also a categorical distribution.
This also implies that the nonidentiﬁability from pairwise prediction in the HMM case cannot 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 known to 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.
Though this 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 natural idea is to force the model to “predict more”, and one straightforward way to do so is to combine multiple 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 pairwise tasks involved 3 (adjacent) tokens:
Theorem 4 (Nonidentiﬁability of HMM from all pairwise predictions on 3 tokens) .
There exists a pair of HMM distributions with parameters ( O, T ) and ( ˜ O, ˜ T ) , each satisfying Assumptions 1,and 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 optimal predictors are the same under each distribution.
We brieﬂy remark that the reason for only considering adjacent time steps is that when the tokens are 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 pairwise predictions on HMMs comes from non-uniqueness of matrix factorization.
While adding additional pairwise prediction tasks introduces more equations on the product of matrices, these equations are highly dependent, and the proof works by providing counterexamples that can simultaneously satisfy all these equations.
The above intuition leads to another way of forcing the model to “predict more”, that is, to increase the number of predicted tokens.
The hope is that doing so results in equations on tensors—as opposed to matrices— for which there is a lot of classical machinery delineating tensors for which the rankdecomposition is unique, as discussed in Section 2.4.
This intuition proves to be true and we show that increasing 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 any permutation 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 the same, 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 conditional 2-tensors (i.e. matrices) given by the predictors.
The proof idea is to construct a third-order tensor by linearly combining the conditional 2-tensors for each possible value of the token being conditioned on, such that Kruskal’s theorem applies and gives identiﬁability.
Note, importantly, that the weights for the linear combination cannot depend on the marginal probabilities of the token being conditioned on, since we do not have access to these marginals, and it is unclear whether we could extract unique marginals given the conditional probabilities we are predicting.
Thus, the above theorem cannot be simply 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 help of 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 HMM parameters 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 the proofs 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 whose components 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 component in equation 2 forms a matrix of Kruskal rank k .
The second and third components clearly satisfy this condition by Assumption 3.
For the ﬁrst component, recall that φ ( x ) =
O (cid:62)
[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 /
[Equation]
( e
[Equation]
d ) j ) (cid:62)
O
[Equation]
1 ] j ∈ [ d ] (cid:1) OT (cid:62) , 7 which is of rank k by Assumption 3.
Hence components W are all of Kruskal rank k , and columns of OT, O are identiﬁed up to column-wise permutation and scaling by Kruskal’s theorem.
The indeterminacy in scaling is further removed noting that columns of O, T need to sum up to 1.
Lastly, T is recovered as 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
[Equation]
we would like to construct a 3-tensor whose components can be uniquely determined by Kruskal’s theorem.
Let X be the same as before
[Equation]
and consider the 3-tensor
W :=
[Equation]
where the ﬁrst component can be simpliﬁed to a
[Equation]
The matrix A :=
[ a 1 , ..., a k ]
=
D − 1 O is of rank k , hence we can identify (up to permutation) columns of each component of W by Kruskal’s theorem.
This means if O, T and ˜ O, ˜ T produce the same predictor, 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 means the 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 ] .
We can then compute D as D = diag
( OT 1 ) , and recover O as O = DA .
Finally, T is also recovered since ˜ O ˜ T
=
O ˜ T =
OT ⇒ ˜
TT − 1 =
I k ⇒ ˜
T = T .
4.2 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.
i .
Combined with Lemma 1, we have ˜ M is equal to (up to a permutation) either M or HM , where H is 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 , the column sum of ˜ M † M is
[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 the constraint that ˜ T should be a stochastic matrix with positive entries and column sum 1.
Hence it must be 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 great amount 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 such as 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énaff et 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 various aspects [Lee et al., 2021b, Liu et al., 2021].
On the theoretical front, there have been analyses on both masked predictions [Lee et al., 2021a,
Zhang and Hashimoto, 2021] and contrastive methods [Arora et al., 2019, Tosh et al., 2020a,b,
Wang and Isola, 2020, HaoChen et al., 2021, Wen and Li, 2021], with a focus on characterizing the quality of the learned features for downstream tasks [Saunshi et al., 2020, Wei et al., 2021].
These approaches usually rely on quite strong assumptions to tie the self-supervised learning objective to the downstream tasks of interest.
In contrast, our work takes the view of parameter identiﬁability, for which there is no need for downstream assumptions but instead the speciﬁc parametric form is key.
Note also that while the parameter recovery lens is a new contribution of our work, Wen and
Li [2021] argue (as a side-product of their analysis) that some aspects of a generative model are recovered in their setup.
Their data model, 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 in the literature.
One important area of research is independent component analysis (ICA), where the data is assumed to be given as a transformation (mixing) of unknown independent sources which
ICA aims to identify.
In nonlinear ICA data models, both the sources and the mixing function are generally not identiﬁable.
However, identiﬁability of the sources can be shown under some additional assumptions (e.g. on the dependency structure of different time steps)
[Hyvarinen and Morioka, 2016, 2017, Hälvä and Hyvarinen, 2020].
Similar ideas have also been applied in the self-supervised setting, where the latent variables can be identiﬁed under suitable assumptions on the conditional distribution of the latent [Zimmermann et al., 2021] or on data augmentations [Von Kügelgen et al., 2021].
Unlike our setup though, the mixing function in these models is deterministic and 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] provide algorithms based on third-order moments.
A major difference between these prior works on tensor methods and ours is that previous results operate on joint moments, while the results in this work are based 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 offers a 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 tuned by both changing the parametric form of the data generative model, and by changing the masked prediction task.
We emphasize that this is a ﬁrst-cut effort in the research program of analyzing SSL through the lens of model identiﬁability; we aim to build on this foundation to extend our analyses from HMMs to more complicated latent sequence and latent variable models.
We also note that we have focused here on population analyses, and model identiﬁability.
It would be of interest to build off this to develop and analyze the corresponding ﬁnite-sample learning algorithms for parametric generative models given SSL tasks, with sample complexity results, both in the realizable case, as well as in the agnostic case where we have model mis-speciﬁcation.
Given the use of conditional MLEs and regressions in SSL, and the natural robustness of these to model mis-speciﬁcations, we conjecture that 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.
Acknowledgement
The authors would like to thank Elan Rosenfeld for insightful discussions and comments.
We acknowledge the support of ONR via N000141812861, NSF via IIS-1909816, IIS-1955532, IIS2211907, and the CMU/PwC DT&I Center.
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] 2.
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.3, 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 for GHMMs, and Appendix B for HMMs.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental 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 they were 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., type of GPUs, internal cluster, or cloud provider)?
[N/A] 4.
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] 5.
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/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Missing proofs for G-HMM
This section provides missing proofs for results on G-HMM.
We will ﬁrst prove the two lemmas on properties of the posterior function (Lemma 1, 2), then show the proof for the three-token prediction task (Theorem 6) using a tensor decomposition idea similar to that of Theorem 5.
At the end, we show the identiﬁability from pairwise conditional distributions (as opposed to conditional expectation as in masked prediction tasks), which is proved by reducing parameter recovery to the identiﬁability of Gaussian mixtures (Theorem 7).
A.1 Proofs of helper lemmas A.1.1 Proof for Lemma
Given the form of the predictor, matching two predictors f, ˜ f means that the corresponding posteriors φ, ˜ φ are matched up to a linear transformation.
We will now prove the following lemma, which says that in this case, φ, ˜ φ can in fact only differ by a permutation of coordinates:
Lemma (Lemma 2 restated) .
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.
Proof.
We will prove the lemma by matching the Jacobian w.r.t.
x on both sides.
Let’s ﬁrst quickly recall the Jacobian of the posterior vector φ ( x ) ∈ R k , where [ φ ( x )]
i =
[Equation]
] exp( −
[Equation]
Denote o ( x ) := (cid:2) − (cid:107) x − µ 1 (cid:107) 2, ..., − (cid:107) x − µ k (cid:107) 2(cid:3) ∈
R k , then ∇ x φ ( x ) = ∇ o ( x ) softmax ( o ( x )) ·
∇ x o ( x ) , where ∇ o
[ softmax ( o )] i =
[ softmax ( o )]
i
[Equation]
Hence the Jacobian is
[Equation]
Denote ∆ :=
M − [ x, x, ..., x ] ∈ R d × k , and similarly ˜∆ = ˜
M − [ x, x, ..., x ] .
Matching ∇ x ˜ φ ( x ) = ∇ x Rφ ( x ) gives
[Equation]
Let’s take x = x ( i ) c :=
cµ i for c > 1 .
We claim that this x ( i ) c satisﬁes lim c →∞ φ ( x ( i ) c ) → e i .
This is because ∀ j (cid:54) =
i , lim
[Equation] [Equation]
where the last equality is because µ (cid:62) i ( µ i − µ j ) > 0 for any µ i , µ j lying on the same hypersphere.
With such choices of x , the two sides of equation 12 are now:
[Equation]
Since x := cµ i for c → ∞ lies outside the afﬁne hull of { ˜ µ i } i ∈ [ k ] , ˜∆ is of full rank due to the following claim:
Claim 2.
Given a linearly independent set { u i } i ∈ [ k ] , if { u i − v } i ∈
[ k ] is not linearly independent, then v = (cid:80) i ∈
[ k ] β i · u i where (cid:80) i ∈
[ k ] β i = 1 .
Proof.
Since { u i − v } i ∈
[ k ] is linearly dependent, we can write some u j − v as the linear combination of other { u i − v } i ∈
[ k ] ,i (cid:54) =
j .
Let’s take j = k wlog, and denote the coefﬁcients of the linear combination as { α i } i ∈ [ k − 1] .
Then
[Equation]
The right hand side is non-zero since { u i } i ∈ [ k ] are linearly independent by assumption, hence 1 − (cid:80) i ∈
[ k − 1]
α i (cid:54) = 0 , and we get
[Equation]
Note that (cid:80) i ∈
[ k ] β i = 1 , hence v is an afﬁne combination of { u i : i ∈ [ k ] } .
Since ˜∆ is full rank, it must be diag ( R i ) − R i R (cid:62) i = 0 , which implies R is a permutation matrix.
This is because for any non-
zero v s.t. diag ( v ) − vv (cid:62) = 0 , the entries of v satisfy v 2 i = 1 , v i v j = 0 for i (cid:54) =
j .
Hence v has exactly one non-zero entry which is ± 1 .
Since Rφ ( x ) = ˜
φ ( x ) where φ ( x ) , ˜ φ ( x ) are both probability vectors with non-negative entries, this non-zero entry has to be 1 (and not -1).
Since R is of rank- k by Assumption 4, this non-zero entry is at different positions for different R i , hence R is a permutation matrix.
A.1.2 Proof of Lemma
We show that if M, ˜ M parameterize φ, ˜ φ respectively and that φ = ˜
φ , then ˜ M must equal to either
M or a unique (and somewhat special) transformation of M :
Lemma (Lemma 1 restated) .
For d ≥ k ≥ 2 , then under Assumption 4, 5, φ = ˜ φ implies ˜ M = M or ˜ M
= HM , where H is a Householder transformation of the form H :
=
I d − 2ˆ v ˆ v (cid:62) ∈ R
d × d , with ˆ v :=
[Equation]
Proof.
Let’s start with d = k .
First, let’s check the conditions for φ = ˜ φ .
For any x ∈ R d , we have [
[Equation]
[˜ φ
[Equation]
x )]
i , ∀ i ∈ [ k ] ⇒ exp (cid:0) −
[Equation]
exp
[Equation]
(cid:1) ,
[Equation] [Equation]
] ⇒ (cid:107) x − µ i (cid:107) 2 − (cid:107) x − ˜ µ i (cid:107)
= (cid:107) x − µ j (cid:107) 2 − (cid:107) x − ˜ µ j (cid:107) 2 , ∀
i, j ∈ [ k ] ⇒ 2 (cid:0) (˜ µ i − µ i ) − (˜ µ j − µ j ) (cid:1) (cid:62) x = (cid:0) (cid:107) µ
[Equation]
(cid:107) 2 − (cid:107) ˜ µ
[Equation]
(cid:107)
[Equation]
(cid:1) − (cid:0) (cid:107) µ i (cid:107) 2 − (cid:107) ˜ µ i (cid:107) 2 (cid:1) , ∀ i, j ∈ [ k ] .
Since the left hand side is linear in x ∈ R d and the right hand side is a constant, it must be that both sides are 0.
That is, the necessary conditions for φ =
˜ φ are that for any i, j ∈ [ k ] , 1) ˜ µ i − µ i = ˜
µ j − µ j , and 2) (cid:107) µ i (cid:107) 2 −(cid:107) ˜ µ i (cid:107)
= (cid:107) µ j (cid:107) 2 −(cid:107) ˜ µ j (cid:107) 2 .
It can be checked that these two conditions are also sufﬁcient for φ = ˜ φ .
Denote v := µ i − ˜ µ i .
The norms of the means are known and equal by Assumption 5, which gives
[Equation]
The last equality in equation 18 holds for a non-zero v when the span of { 2 µ i − v : i ∈
[ k ] } is ( d − 1) -dimensional subspace.
On the other hand, the span of { 2 µ i − v : i ∈ [ k ] } is at least ( k − 1) by Assumption 4.
When d = k , it must be that the dimension is exactly ( d − 1) , which means v is an afﬁne combination of { 2 µ i : i ∈ [ k ] } by Claim 2.
Moreover, v has to be orthogonal to { 2 µ i − v : i ∈ [ k ] } , which leads to the unique choice of v that is the projection of the origin onto the ( d − 1) -dimensional subspace speciﬁed by the afﬁne combinations of { 2 µ i : i ∈ [ k ] } .
Claim 3.
v is the projection of the origin to the hyperplane deﬁned by { 2 µ i : i ∈ [ k ] } , and is the only solution to equation 18.
Proof.
It is clear that this choice of v satisﬁes (2 µ i − v ) (cid:62) v = 0 , ∀ i ∈ [ k ] .
To see that this is the unique choice, suppose there exists some v (cid:48) lying in the hyperplane of { 2 µ i } , and denote δ := v (cid:48) − v .
Note that δ (cid:62) v = 0: let the hyperplane speciﬁed by { 2 µ i } i ∈
[ k ] be speciﬁed as { x : (cid:104) u, x (cid:105) =
c } for some u ∈ R d and c ∈ R .
Then v , the projection of the origin, can be written as v =
c (cid:107) u (cid:107) · u (cid:107) u (cid:107) , i.e. v is proportional to the normal vector u .
For any v (cid:48) in the hyperplane, it satisfy (cid:104) u, v (cid:48) (cid:105) =
c , and
[Equation]
Then for any v (cid:48) satisfying equation 18,
[Equation]
Since { 2 µ i − δ } i ∈ [ k ] spans the ( k − 1) -dimensional
hyperplane and that δ lies in the hyperplane, it must be that δ = 0 , i.e. v (cid:48)
= v .
Note that this choice of v also satisﬁes (cid:107) µ i − v (cid:107) =
(cid:107) µ
i (cid:107) , since v and the origin are reﬂections w.r.t.
the hyperplane that is the afﬁne hull of { µ i : i ∈
[ k ] } .
In other words, { µ i − v } i ∈ [ k ] is related to { µ i } i ∈
[ k ] via the Householder transformation of the form H :
= I d − 2 vv (cid:62) (cid:107) v (cid:107) 2 , i.e. µ i − v =
Hµ i .
Denote ˆ v := v
[Equation]
An explicit formula for ˆ v is ˆ v :=
[Equation]
This ﬁnishes the proof for d = k .
For d > k , the above argument still applies and H remains the only indeterminacy (up to permutation), where H :
= I d − 2ˆ v ˆ v (cid:62) for ˆ v :=
[Equation]
The reason is that even though the ambient dimension d is larger, { µ i − v : i ∈
[ k ] } has to have the same span as { µ i : i ∈
[ k ] } , since having the same predictor requires the column space of M, ˜ M to match.
Hence we only need to consider v in the k -dimensional column space of M , which reduces to the case of d = k .
A.2 Identiﬁability of predicting x t 2 ⊗ x t 3 | x t 1 , G-HMM
Theorem 5 shows that triplet prediction tasks (i.e. predict 2 tokens given 1) sufﬁces for the identiﬁability of HMM, using tools from the uniqueness of tensor decomposition.
The next theorem shows that the same conclusion also applies for G-HMM:
Theorem 6 (Identiﬁability from masked prediction on three tokens, G-HMM) .
Let ( t 1 , t 2 , t 3 ) be any permutation of (1 , 2 , 3) , and consider the prediction task x t 2 ⊗ x t 3 | x t 1 .
Under Assumption 1, 4, 5, if the optimal predictors 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.
Proof.
Similar to the discrete case, we will prove x 2 ⊗ x 3 | x 1 and x 1 ⊗ x 3 | x 2 separately; the proof for x 1 ⊗ x 2 | x 3 is analogous to x 2 ⊗ x 3 | x 1 by symmetry and hence omitted.
The proofs also follow a similar strategy as in the proof for Theorem 5, that is, to construct a 3-tensor using the predictor, on which applying Kruskal’s theorem provides identiﬁability.
Case 1, x 2 ⊗ x 3 | x 1 : Let X :=
{ x ( i ) ∈ R
d : i ∈ [ k ] } be a linearly independent set, and consider the following 3-tensor:
[Equation]
[Equation]
The matrices formed by second and third components are both of rank- k by Assumption 4.
Hence in order to apply Kruskal’s theorem on W , it sufﬁces to show that there exists a choice of X such that the matrix A :
=
[ a 1 , ..., a k ] is of rank k .
One such choice is to let x ( i ) =
µ i , which gives
[Equation]
Since M , T are both of rank k by Assumption 4, we only need to argue that the matrix Φ :=
[ φ ( µ 1 ) , ..., φ ( µ k )]
∈
R k × k is of full rank.
Recall that for a mixture of k Gaussians with identify covariance and mean { µ i ∈ R d : i ∈ [ k ] } , the posterior function φ is deﬁned entrywise as
[Equation]
To show Φ is of full rank, we can equivalently show that a columnwise scaled version of Φ is full rank.
In particular, let’s look at the matrix ˆΦ ∈ R k × k , where ˆΦ ij = exp( − (cid:107) µ i − µ j (cid:107) 2) ; that is, each column of ˆΦ can be considered as a scaled version of the column in Φ without the normalization for a unit (cid:96) 1 norm.
It can be seen that ˆΦ is a Gaussian kernel matrix which is known to be full rank.
Therefore we have shown that each component of the tensor W := (cid:80) i ∈
[ k ] a i ⊗ M i ⊗ ( MT )
i has
Kruskal rank k , which allows to recover columns of M, MT up to permutation and scaling by
Kruskal’s theorem.
The indeterminacy in scaling is further removed since the norms of { M i } i ∈
[ d ] are known by Assumption 5.
On the other hand, for any ˜ M, ˜ T that form the same predictor as M, T , W can also be written as
[Equation]
Hence columns of M, ˜ M and MT, ˜ M ˜ T are both matched up to a shared permutation, which proves identiﬁability.
Case 2, E
[ x 1 ⊗ x 3 | x 2 ] :
For the task of predicting x 1 , x 3 given x 2 , the predictor takes the form
[Equation]
[Equation]
Let X := { µ i : i ∈ [ k ] } as in the previous case, and consider the 3-tensor
[Equation]
The ﬁrst component is of rank- k as shown in the proof for x 2 ⊗ x 3 | x 1 , and the other two components are of rank- k by Assumption 4.
Thus Kruskal’s theorem applies and the columns of MT, MT (cid:62) are recovered up to a shared permutation.
The ﬁrst component { a i } i ∈ [ k ] are also recovered, which means that if ˜ M, ˜ T form the same predictor as M, T , then for any linearly independent set X with k elements (not necessarily the previous choice of { µ i } i ∈
[ k ] ) such that X leads to a full rank A , we have A =
˜
A where ˜ A is parameterized by ˜ M, ˜ T .
For any such X = { x ( i ) : i ∈ [ k ] } , denote X :=
[ x (1) , ..., x ( k ) ] , then
[Equation]
Since X is of rank- k by the choice of X , this means
[Equation]
Moreover, for any valid choice of X , matrices deﬁned with points in sufﬁciently small neighborhoods of x ( i ) are still of full rank by the upper continuity of matrix rank.
Hence the equality in equation 28 holds for points in these neighborhoods, and thus the Jacobian on both sides should be equal.
Then, the exact same proof of Lemma 2 applies, and we get ˜ φ , φ are equal up to a permutation of coordinates.
Thus ˜ M must be equal to (up to permutation) either M or HM for a Householder reﬂection H by Lemma 1.
Finally, the solution of HM is eliminated since it would lead to a ˜ T that is not a valid stochastic matrix, as shown in the proof of Theorem 3.
A.3 Identiﬁability from pairwise conditional distribution
We show that matching the entire conditional distribution for G-HMM provides identiﬁability.
Though this is implied by Theorem 3, which states that matching the conditional expectation already sufﬁces, having access to the full conditional distribution allows an even simpler proof.
Theorem 7 (Identiﬁability of conditional distribution) .
Let M, T and ˜ M, ˜ T be two set of parameters satisfying Assumption 1 and 4.
If p ( x 2 | x 1 ; M, T ) = p ( x 2 | x 1 ; ˜ M, ˜ T ) , ∀ x 1 , x 2 ∈ R d , then M = ˜ M ,
T =
˜ T up to a permutation of labeling.
Proof.
First note that the conditional distribution of x 2 given x 1 is a mixture of Gaussian, with means { µ i } i ∈ [ k ] and mixture weights given by P ( h 2 | x 1 ) =
TP ( h 1 | x 1 ) , hence we can directly apply the identiﬁability of Gaussian mixtures to recover the means { µ i } i ∈ [ k ] :
Lemma 3 (Proposition 4.3 in Lindsay and Basak [1993]) .
Let Q k denote a Gaussian mixture with means { ξ j } j ∈ [ k ] ∈ R d .
Suppose ∃ l ∈ [ d ] such that the set { [ ξ j ] l } has distinct values, then one can recover { ξ j } j ∈ [ k ] from moments of Q k .
We note that the assumption on the existence of a coordinate l ∈ [ k ] is with out loss of generality, since we can ﬁrst rotate the means to a different coordinate system in which this condition holds, then rotation back the means.
Such rotation is guaranteed to exist since ﬁnding such rotation is equivalent to ﬁnding a vector v s.t. v (cid:62) ( µ i − µ j ) (cid:54) = 0 for every i, j ∈ [ k ] , for which the solution set is
R d \ ∪ i,j ∈ [ k ] { u : u (cid:62) ( µ i − µ j ) = 0 } (cid:54) = ∅ .
Recovering { µ i } i ∈
[ k ] means the scaled likelihood and the posterior both match, i.e. ψ = ˜ ψ , and φ ( x ) =
P ( h | x ) = ψ (cid:107) ψ (cid:107) 1 .
The conditional distribution is p
[Equation]
Choose a set X :=
{ x ( i ) } i ∈ [ k ] such that Ψ X :=
[ ψ ( x (1) ) , ..., ψ ( x ( k ) )]
∈
R k × k is full rank.
Φ X :=
[ φ ( x (1) ) , ..., φ ( x ( k ) )] ∈
R k × k is also full rank since its columns are nonzero scalings of columns of Ψ X .
Then we have
[Equation]
B Proof of Theorem 4: non-identiﬁability of HMM from multiple pairwise predictions
Theorem (Theorem 4 restated: nonidentiﬁability of HMM from multiple pairwise predictions) .
There exists a pair of HMM distributions with parameters ( O, T ) and ( ˜ O, ˜ T ) , each satisfying Assumptions 1, 2 and 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 optimal predictors are the same under each distribution.
Proof.
We provide an example to show the nonidentiﬁability result in Theorem 4.
The goal is to ﬁnd ˜ O (cid:54) =
O , ˜ T (cid:54) =
T that produce the same predictors for predicting both x 2 | x 1 and x 3 | x 1 .
We will choose T, ˜ T to be symmetric, so that O, T and ˜ O, ˜ T also form the same predictors for the reversed direction, i.e., for predicting x 1 given x 2 and x 1 given x 3 , since the reverse chain has transition matrix T (cid:62)
= T .
Let’s consider the case where the all row sums of O and ˜
O are k/d .
Consequently, the posterior function is simply φ ( x ) =
O (cid:62) x (cid:107) O (cid:62) x (cid:107) 1 = dk
O
(cid:62) x , and similarly we have ˜ φ ( x ) = dk
˜ O (cid:62) x .
The predictors are of the form:
[Equation]
Matching f 2 | 1 ( x ) =
˜
f
2 | 1 ( x ) on all x ∈ X := { e i } i ∈
[ d ] means
[Equation]
Similarly, matching f 3 | 1 = ˜ f 3 | 1 gives OT 2 O (cid:62) =
˜ O ˜ T 2 ˜ O
(cid:62) , hence
[Equation] [Equation]
where step ( i ) uses ˜ O ˜ O † O
= O , since ˜ O, O share the same column space.
Denote R :=
˜ O † O ; R is orthogonal by the last equality in equation 33.
To construct the desired example, consider k = 3 , and let R represent a rotation with axis of rotation 13 ( e 1 + e 2 + e 3 ) .
This axis is the direction pointing from the origin to the projection of the origin on the hyperplane
P c := { v
∈ R d : (cid:80) i ∈
[ d ] v i =
c } for any positive constant c (i.e. P c is parallel to the hyperplane in which probability vectors lie).
This means such rotation guarantees Rv ∈ P c , ∀ v ∈ P c , and has the following property:
Claim 4.
Each row and each column of R sums up to 1.
Deﬁne ˜ O := OR , ˜ T := R (cid:62) TR , Claim 4 ensures that row sum and column sum of ˜ O, ˜ T remain the same as those of O, T .
When the rotation angle represented by R is sufﬁciently small, entries ˜ O, ˜ T remain in [0 , 1] , hence such ˜ O, ˜ T form a valid example.
We will provide a concrete example in the subsequent subsection.
B.1 Example for Theorem
The intuition of the nonidentiﬁability result in Theorem 4 is related to the non-uniqueness of matrix factorization: while adding additional pairwise prediction tasks introduces more equations on the product of matrices, these equations can be highly dependent, and there are cases where different set of matrices can simultaneously satisfy all the equations.
We now provide a concrete example for the non-identiﬁability of predicting x 2 | x 1 , x 1 | x 2 , x 3 | x 1 , and x 1 | x 3 , by ﬁnding 2 set of O, T such that the corresponding predictors (of the form speciﬁed in equation 31) match.
Let d
= 4 , k = 3 ,
O =   0 .
[Table]

 , ˜ O =  
[Table]

 ,
T =  
[Table]
  , ˜ T =  
[Table]
  , det ( O ) = det ( ˜ O )
= 0 .
0110 , det ( T )
[Equation]
Note that T, ˜ T are both symmetric as desired by the proof of Theorem 4, which means this is also a valid counter example for learning to predict x 1 | x 2 and x 1 | x 3 , and hence for all of x 2 | x 1 , x 1 | x 2 , x 3 | x 1 , and x 1 | x 3 .
B.2 Proof of Claim
Proof.
We would like to show that each row and each column of R sums up to 1.
Denote the d -dimensional simplex by ∆ d , i.e. ∆ d := { x ∈ R d : (cid:80) i ∈ [ d ] x i = 1 } , and let P c := { v ∈ R
d : (cid:80) i ∈ [ d ] v i =
c } for some positive constant c denote a hyperplane parallel to the hyperplane in which probability vectors lie.
Let’s ﬁrst check that the columns of R sum up to 1.
Any v ∈ P c can be written as v = c · [ α 1 , α 2 , ..., α d − 1 , 1 − (cid:80) i ∈
[ d − 1] α i ] for some [ α 1 , ..., α d − 1 ]
∈ ∆ d − 1 .
Let r i denote the i th row of
R , then Rv ∈ P c means (cid:80) i ∈
[ d ] (cid:104) r i , v (cid:105) = (cid:104) (cid:80) i ∈ [ d ] r i , v (cid:105) =
c .
Let β j denote the j th coordinate of (cid:80) i ∈
[ d ] r i , then (cid:88)
[Equation]
⇒ (cid:88)
[Equation]
⇒ β
i = 1 , ∀ i ∈
[ d ] .
[Equation]
It then follows that R − 1 =
R (cid:62) also has columns summing up to 1, since
[Equation]
C Nonidentiﬁability from large time gaps
As noted earlier, there is an inherent obstacle when using prediction tasks on tokens that are more than 1 time gaps apart.
For instance, if we are predicting x t +1 given x 1 for some t > 1 with G-HMM, then we are still able to identify M from the posterior function, however it remains to to recover T from T t .
For general matrices, it is clear that matching a power of a matrix does not imply the matrix itself is matched.
For our case, even though requiring T to be stochastic adds additional constraints, matching the matrix power still does not sufﬁce to identify the underlying matrix, as formalized in the following claim.
Claim 5 (Nonidentiﬁability of matrix powers (Claim 1 restated)) .
For any positive integer t , there exist stochastic matrices T, ˜ T satisfying Assumption 1, 3, such that T (cid:54) =
˜ T and T t =
˜ T t .
Proof.
As in Theorem 4, the nonidentiﬁability comes from the non-uniqueness of matrix factorization.
Speciﬁcally for this case, we will set ˜ T to be equal to T up to a special rotation that gets composed when taking the matrix power.
That is, we want ˜ T =
RT =
TR for some matrix R that implicitly performs a rotation, so that ˜ T t = T t R t .
Since R corresonds to a rotation, we can choose the rotation angle properly so that R t
= I , and hence ˜ T t = T t but ˜ T (cid:54)
= T .
Precisely, using notations for the G-HMM setup, set a ∈ [0 , 1] , and let the parameters ( T, M ) be given by
T
[Equation]
Let θ be some rotation angle, and denote by R ( θ ) :=
[Equation]
cos( θ ) −
[Equation]
sin(
[Equation]
a rotation that acts on the ﬁrst two dimensions.
We will show that for any θ ∈ R , we have
[Equation]
Assuming equation 37, since R ( θ ) represents a rotation of angle θ , (cid:0) R ( θ ) (cid:1)
τ corresponds to a rotation of angle τθ for any integer τ ( τ could be negative).
Setting θ := 2 πt , we then have
[Equation]
For ˜ T to serve as a valid example for our theorem, it remains to check that for every t , there exists a choice of a such that ˜ T :
=
RT , where R :=
M − 1 (cid:0) R ( 2 πt ) (cid:1) − 1 M , is a valid stochastic matrix.
That is, ˜ T has 1) columns and rows each summing up to 1, and 2) entries bounded in [0 , 1] .
Let’s ﬁrst show that the columns and rows each sum up to 1.
Noting that M − 1 =
[Equation]
 , the column sums are
[Equation]
where step ( i ) uses equation 37.
Similarly, the row sums are
[Equation]
[Equation]
To show that there exists a choice of T such that entries of ˜ T are non-negative, we provide a concrete example where T is deﬁned with a = 12 .
It can be checked that ˜ T :=
M − 1 ( R ( 2 πt )) − 1 M has non-negative entries for t ∈ { 2 , 3 , 4 , ..., 10 } .
For larger t , let θ = 2 πt , then we have by the Taylor expansion of R ( 2 πt ) :
[Equation]
for some constants c 1 ∈ [ − 14! , 14! ] , c 2 ∈ [ − 12 , 12 ] .
Substituting this into ˜ T :
= M − 1 R ( θ ) − 1 M gives ˜ T =
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where the inequality ( i ) is taken entry-wise.
It can be checked that all entries are non-negative for θ ≤ 2 π 10 .
Proof of equation 37 Let’s conclude the proof by proving the commutativity in equation 37.
Denote
R 2 ( θ ) :=
[Equation]
Denote
[Equation]
[Equation] [Equation]
(cid:21) .
We can write
[Equation]
Let R 2 ( θ ) denote a clockwise rotation of angle θ , then
[Equation]
4 π 3 (cid:1) v 3 , v 3 ] , (44) where v 1 = (cid:20)
(cid:21) , v 2 = (cid:20) − 1 / 2 −√ 3 / 2 (cid:21) , v 3 =
(cid:20) − 1 / 2 √ 3 / 2 (cid:21) .
Denote α ij := v (cid:62) i R (cid:62) 2 v j for i, j ∈ [3] .
Noting
T = aI
+ (1 − a )
[Equation] [Equation]
(cid:35) :=
[Equation]
+ (1 − a ) P , we have
[Equation]
(45) where step ( i ) uses 11 (cid:62) T
= T 11 (cid:62) =
11 (cid:62) .
The equality ( ∗ ) is true due to equation 44.
D Simulation
We empirically verify the identiﬁability results for HMM (Theorem 5) and G-HMM (Theorem 3) on simulation data, by checking whether matching the optimal predictor implies matching the parameters of the data generative model.
In particular, we aim to recover O ∈ R d × k , T ∈ R k × k for HMM, and M ∈ R d × k , T ∈ R k × k for
G-HMM, with d = 10 and k ∈ { 2 , 4 } .
Given a batch of samples B := { x ( i ) 1 } i ∈|B| , the objective to minimize for HMM is
[Equation]
where f 2 , 3 | 1 O,T ( x 1 ) :
= E O,T [ x 2 ⊗ x 3 | x 1 ] ; the objective for the task of predicting x 1 ⊗ x 3 given x 2 is deﬁned analogously.
Note that though equation 46 differs from equation 1 by a constant, 8 it sufﬁces for verifying parameter identiﬁability since both losses are minimized at f O,T = f O ∗ ,T ∗ .
We choose to use the form in equation 46 since it is more stable to optimize for and that its minimal loss value is 0, making it easy to check for optimality.
Similarly, the objective to minimizer for G-HMM is
[Equation] [Equation]
where f 2 | 1 M,T :=
E M,T [ x 2 | x 1 ] .
For both HMM and G-HMM, we optimize for O (or M ) and T alternatingly in different epochs.
We found that it is usually helpful to use a larger learning rate for T than for O (or M ), and that normalized gradient descent helps speed up training
Figure 1 and 2 show the results for HMM and G-HMM.
It can be seen that as the objective value approaches the optimum, the parameter distances indeed go to zero, corroborating Theorem 5 and 3.