Optimal prediction of Markov chains with and without spectral gap

We study the following learning problem with dependent data : Observing a trajectory of length n from a stationary Markov chain with k states , the goal is to predict the next state . For 3 ≤ k ≤ O ( √ n ) , using techniques from universal compression , the optimal prediction risk in Kullback-Leibler divergence is shown to be Θ ( k 2 n log n k2 ) , in contrast to the optimal rate of Θ ( log logn n ) for k = 2 previ- ously shown in [ FOPS16 ] . These rates , slower than the parametric rate of O ( k 2 n ) , can be attributed to the memory in the data , as the spectral gap of the Markov chain can be arbitrarily small . To quantify the memory effect , we study irreducible reversible chains with a prescribed spectral gap . In addition to characterizing the optimal prediction risk for two states , we show that , as long as the spectral gap is not excessively small , the prediction risk in the Markov model is O ( k 2 n ) , which coincides with that of an iid model with the same number of parameters . 1 Introduction . Learning distributions from samples is a central question in statistics and machine learning . While significant progress has been achieved in property testing and estimation based on independent and identically distributed ( iid ) data , for many applications , most notably natural language processing , two new challenges arise : ( a ) Modeling data as independent observations fails to capture their temporal dependency ; ( b ) Distributions are commonly supported on a large domain whose cardinality is comparable to or even exceeds the sample size . Continuing the progress made in [ FOPS16 , HOP18 ] , in this paper we study the following prediction problem with dependent data modeled as Markov chains . Suppose X1 , X2 , . . . is a stationary first-order Markov chain on state space [ k ] ≜ { 1 , . . . , k } with unknown statistics . Observing a trajectory Xn ≜ ( X1 , . . . , Xn ) , the goal is to predict the next state Xn+1 by estimating its distribution conditioned on the present data . We use the KullbackLeibler ( KL ) divergence as the loss function : For distributions P = [ p1 , . . . , pk ] , Q = [ q1 , . . . , qk ] , D ( P∥Q ) = ∑k i=1 pi log pi qi if pi = 0 whenever qi = 0 and D ( P∥Q ) = ∞ otherwise . The minimax 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . prediction risk is given by Riskk , n ≜ inf M̂ sup π , M E [ D ( M ( ·|Xn ) ∥M̂ ( ·|Xn ) ) ] = inf M̂ sup π , M k∑ i=1 E [ D ( M ( ·|i ) ∥M̂ ( ·|i ) ) 1 { Xn=i } ] ( 1 ) where the supremum is taken over all stationary distributions π and transition matrices M ( rowstochastic ) such that πM = π , the infimum is taken over all estimators M̂ = M̂ ( X1 , . . . , Xn ) that are proper Markov kernels ( i.e . rows sum to 1 ) , and M ( ·|i ) denotes the ith row of M . Our main objective is to characterize this minimax risk within universal constant factors as a function of n and k. The prediction problem ( 1 ) is distinct from the parameter estimation problem such as estimating the transition matrix [ Bar51 , AG57 , Bil61 , WK19 ] or its properties [ CS00 , KV16 , HJL+18 , HKL+19 ] in that the quantity to be estimated ( conditional distribution of the next state ) depends on the sample path itself . This is precisely what renders the prediction problem closely relevant to natural applications such as autocomplete and text generation . In addition , this formulation allows more flexibility with far less assumptions compared to the estimation framework . For example , if certain state has very small probability under the stationary distribution , consistent estimation of the transition matrix with respect to usual loss function , e.g . squared risk , may not be possible , whereas the prediction problem is unencumbered by such rare states . In the special case of iid data , the prediction problem reduces to estimating the distribution in KL divergence . In this setting the optimal risk is well understood , which is known to be k−12n ( 1 + o ( 1 ) ) when k is fixed and n → ∞ [ BFSS02 ] and Θ ( kn ) for k = O ( n ) [ Pan04 , KOPS15 ] . 1 Typical in parametric models , this rate kn is commonly referred to the “ parametric rate ” , which leads to a sample complexity that scales proportionally to the number of parameters and inverse proportionally to the desired accuracy . In the setting of Markov chains , however , the prediction problem is much less understood especially for large state space . Recently the seminal work [ FOPS16 ] showed the surprising result that for stationary Markov chains on two states , the optimal prediction risk satisfies Risk2 , n = Θ ( log log n n ) , ( 2 ) which has a nonparametric rate even when the problem has only two parameters . The follow-up work [ HOP18 ] studied general k-state chains and showed a lower bound of Ω ( k log lognn ) for uniform ( not necessarily stationary ) initial distribution ; however , the upper bound O ( k 2 log logn n ) in [ HOP18 ] relies on implicit assumptions on mixing time such as spectral gap conditions : the proof of the upper bound for prediction ( Lemma 7 in the supplement ) and for estimation ( Lemma 17 of the supplement ) is based on Berstein-type concentration results of the empirical transition counts , which depend on spectral gap . The following theorem resolves the optimal risk for k-state Markov chains : Theorem 1 ( Optimal rates without spectral gap ) . There exists a universal constant C > 0 such that for all 3 ≤ k ≤ √ n/C , k2 Cn log ( n k2 ) ≤ Riskk , n ≤ Ck2 n log ( n k2 ) . ( 3 ) Furthermore , the lower bound continues to hold even if the Markov chain is restricted to be irreducible and reversible . Remark 1 . The optimal prediction risk of O ( k 2 n log n k2 ) can be achieved by an average version of the add-one estimator ( i.e . Laplace ’ s rule of succession ) . Given a trajectory xn = ( x1 , . . . , xn ) of length n , denote the transition counts ( with the convention Ni ≡ Nij ≡ 0 if n = 0 , 1 ) Ni = n−1∑ ℓ=1 1 { xℓ=i } , Nij = n−1∑ ℓ=1 1 { xℓ=i , xℓ+1=j } . ( 4 ) 1Here and below ≍ , ≲ , ≳ or Θ ( · ) , O ( · ) , Ω ( · ) denote equality and inequalities up to universal multiplicative constants . The add-one estimator for the transition probability M ( j|i ) is given by M̂+1xn ( j|i ) ≜ Nij + 1 Ni + k , ( 5 ) which is an additively smoothed version of the empirical frequency . Finally , the optimal rate in ( 3 ) can be achieved by the following estimator M̂ defined as an average of add-one estimators over different sample sizes : M̂xn ( xn+1|xn ) ≜ 1 n n∑ t=1 M̂+1xnn−t+1 ( xn+1|xn ) . ( 6 ) In other words , we apply the add-one estimator to the most recent t observations ( Xn−t+1 , . . . , Xn ) to predict the next Xn+1 , then average over t = 1 , . . . , n. Such Cesàro-mean-type estimators have been introduced before in the density estimation literature ( see , e.g. , [ YB99 ] ) . It remains open whether the usual add-one estimator ( namely , the last term in ( 6 ) which uses all the data ) or any add-c estimator for constant c achieves the optimal rate . In contrast , for two-state chains the optimal risk ( 2 ) is attained by a hybrid strategy [ FOPS16 ] , applying add-c estimator for c = 1logn for trajectories with at most one transition and c = 1 otherwise . Also note that the estimator in ( 6 ) can be computed in O ( nk ) time . To derive this first note that given any j ∈ [ k ] calculating M̂+1 xn−11 ( j|xn−1 ) takes O ( n ) time and given any M+1 xn−1n−t+1 ( j|xn−1 ) we need O ( 1 ) time to calculate M̂+1xn−1n−t+2 ( j|xn−1 ) . Summing over all j we get the algorithmic complexity upper bound . Theorem 1 shows that the departure from the parametric rate of k 2 n , first discovered in [ FOPS16 , HOP18 ] for binary chains , is even more pronounced for larger state space . As will become clear in the proof , there is some fundamental difference between two-state and three-state chains , resulting in Risk3 , n = Θ ( logn n ) ≫ Risk2 , n = Θ ( log logn n ) . It is instructive to compare the sample complexity for prediction in the iid and Markov model . Denote by d the number of parameters , which is k − 1 for the iid case and k ( k − 1 ) for Markov chains . Define the sample complexity n∗ ( d , ϵ ) as the smallest sample size n in order to achieve a prescribed prediction risk ϵ . For ϵ = O ( 1 ) , we have n∗ ( d , ϵ ) ≍ d ϵ iid d ϵ log log 1 ϵ Markov with 2 states d ϵ log 1 ϵ Markov with k ≥ 3 states . ( 7 ) At a high level , the nonparametric rates in the Markov model can be attributed to the memory in the data . On the one hand , Theorem 1 as well as ( 2 ) affirm that one can obtain meaningful prediction without imposing any mixing conditions ; 2 such decoupling between learning and mixing has also been observed in other problems such as learning linear dynamics [ SMT+18 , DMM+19 ] . On the other hand , the dependency in the data does lead to a strictly higher sample complexity than that of the iid case ; in fact , the lower bound in Theorem 1 is proved by constructing chains with spectral gap as small as O ( 1n ) ( see Section 3 ) . Thus , it is conceivable that with sufficiently favorable mixing conditions , the prediction risk improves over that of the worst case and , at some point , reaches the parametric rate . To make this precise , we focus on Markov chains with a prescribed spectral gap . It is well-known that for an irreducible and reversible chain , the transition matrix M has k real eigenvalues satisfying 1 = λ1 ≥ λ2 ≥ . . . λk ≥ −1 . The absolute spectral gap of M , defined as γ∗ ≜ 1−max { |λi| : i ̸= 1 } , ( 8 ) quantifies the memory of the Markov chain . For example , the mixing time is determined by 1/γ∗ ( relaxation time ) up to logarithmic factors . As extreme cases , the chain which does not move ( M is identity ) and which is iid ( M is rank-one ) have spectral gap equal to 0 and 1 , respectively . We refer the reader to [ LP17 ] for more background . Note that the definition of absolute spectral gap requires irreducibility and reversibility , thus we restrict ourselves to this class of Markov chains ( it is possible to use more general notions such as pseudo spectral gap to quantify the memory of the 2To see this , it is helpful to consider the extreme case where the chain does not move at all or is periodic , in which case predicting the next state is in fact easy . process , which is beyond the scope of the current paper ) . Given γ0 ∈ ( 0 , 1 ) , define Mk ( γ0 ) as the set of transition matrices corresponding to irreducible and reversible chains whose absolute spectral gap exceeds γ0 . Restricting ( 1 ) to this subcollection and noticing the stationary distribution here is uniquely determined by M , we define the corresponding minimax risk : Riskk , n ( γ0 ) ≜ inf M̂ sup M∈Mk ( γ0 ) E [ D ( M ( ·|Xn ) ∥M̂ ( ·|Xn ) ) ] ( 9 ) Extending the result ( 2 ) of [ FOPS16 ] , the following theorem characterizes the optimal prediction risk for two-state chains with prescribed spectral gaps ( the case γ0 = 0 correspond to the minimax rate in [ FOPS16 ] over all binary Markov chains ) : Theorem 2 ( Spectral gap dependent rates for binary chain ) . For any γ0 ∈ ( 0 , 1 ) Risk2 , n ( γ0 ) ≍ 1 n max { 1 , log log ( min { n , 1 γ0 } ) } . Theorem 2 shows that for binary chains , parametric rate O ( 1n ) is achievable if and only if the spectral gap is nonvanishing . While this holds for bounded state space ( see Corollary 4 below ) , for large state space , it turns out that much weaker conditions on the absolute spectral gap suffice to guarantee the parametric rate O ( k 2 n ) , achieved by the add-one estimator applied to the entire trajectory . In other words , as long as the spectral gap is not excessively small , the prediction risk in the Markov model behaves in the same way as that of an iid model with equal number of parameters . Similar conclusion has been established previously for the sample complexity of estimating the entropy rate of Markov chains in [ HJL+18 , Theorem 1 ] . Theorem 3 . The add-one estimator in ( 5 ) achieves the following risk bound . ( i ) For any k ≥ 2 , Riskk , n ( γ0 ) ≲ k 2 n provided that γ0 ≳ ( log k k ) 1/4 . ( ii ) In addition , for k ≳ ( log n ) 6 , Riskk , n ( γ0 ) ≲ k 2 n provided that γ0 ≳ ( log ( n+k ) ) 2 k . Corollary 4 . For any fixed k ≥ 2 , Riskk , n ( γ0 ) = O ( 1n ) if and only if γ0 = Ω ( 1 ) .

Summary: From a single trajectory of $n$ observations drawn from a Markov chain over $k$ different states, we are interested in predicting the next one. For some notion of risk they define, the authors prove that when $k \geq 3$, the rate is of the order $\Theta((k^2/n) \log (n/k))$. In particular, this shows the surprising fact that this rate has a different (worse) dependence in $n$ than chains over $k = 2$ states, known to be $\Theta((1/n)\log \log n)$. Furthermore, they analyze the minimax problem over the set of chains whose spectral gap is controlled from below, and show that under favorable (quantified) mixing conditions, one even recovers the "iid parametric" rate $\Theta(k^2/n)$.