Identifying Sparse Low-Dimensional Structures in Markov Chains: A Nonnegative Matrix Factorization Approach
Mahsa Ghasemi Abolfazl Hashemi Haris Vikalo Ufuk Topcu

    
A weakly supervised adaptive triplet loss for deep metric learning
    Xiaonan Zhao1, Huan Qi1, Rui Luo1 and Larry Davis1,2



1Amazon.com Services, Inc., Seattle, USA 2University of Maryland, College Park, USA

{xiaonzha,qihuan,luorui,lrrydav}@amazon.com

    December 30, 2023
=============================================================================================================================================================================================




We consider the problem of learning low-dimensional representations for large-scale Markov chains. We formulate the task of representation learning as that of mapping the state space of the model to a low-dimensional state space, called the kernel space. The kernel space contains a set of meta states which are desired to be representative of only a small subset of original states. To promote this  structural property, we constrain the number of nonzero entries of the mappings between the state space and the kernel space. By imposing the desired characteristics of the representation, we cast the problem as a constrained nonnegative matrix factorization. To compute the solution, we propose an efficient block coordinate gradient descent and theoretically analyze its convergence properties.





§ INTRODUCTION


A variety of queries about stochastic systems boil down to study  of Markov chains and their properties. They have been widely used as a modeling tool in applications including control <cit.>, machine learning <cit.>, and computational biology <cit.>. Moreover, they create the foundation for more complex probabilistic graphical models including hidden Markov models and Markov decision processes <cit.>.


In many practical settings, a system modeled as a Markov chain, has a large state space. For instance, fine discretization of a zero-input dynamical model with continuous space may lead to a Markov chain with a huge discrete state space. The fact that analyzing such large-scale models may be intractable has motivated significant research on model reduction algorithms. These algorithms attempt to create compressed abstractions that enable efficient downstream analysis without compromising the performance.


A key enabling factor in abstracting Markov chains is the existence of certain structural properties of the characterizing transition probabilities. For instance, the transition probabilities, captured by a stochastic matrix, may be low-rank or sparse. Therefore, one can exploit these structural properties to construct abstractions that accurately approximate the original model. 
The present work is motivated by the state aggregation framework for reducing the complexity of
reinforcement learning and control systems <cit.>. State aggregation schemes  attempt to group similar states into a small number of meta states, which are typically handpicked based on domain-specific knowledge <cit.>, or based on a given similarity metric or feature function <cit.>.


In this paper, we propose an algorithm to find a surrogate representation that approximates the original Markov chain while having a low-dimensional state space. The proposed framework aims to learn a bidirectional mapping between the original high-dimensional state space and the low-dimensional state space, referred to as the kernel space. Additionally, to improve representativeness of the states in the kernel space, we constrain the mappings to be sparse. We model the task of learning the mappings and the kernel transition as a combinatorial optimization problem. Then, we relax this formulation and establish a sparsity-promoting constrained nonnegative matrix factorization problem.
In order to solve this factorization, we propose an efficient block coordinate gradient descent algorithm that starting from an initial guess, learns the bidirectional mappings and the kernel transition in an iterative fashion. We further prove that under certain conditions on the step sizes, the algorithm converges to a stationery point of the proposed optimization problem. We complement our methodology with extensive simulation results where we demonstrate efficacy of the proposed algorithm in terms of the quality of the low-dimensional representation as well as its computational cost.




 §.§ Related Work



In the analysis of dynamical systems, different model-reduction techniques have been designed, such as approximating the transfer
operators <cit.>, dynamic mode decomposition <cit.>, and data-driven approximations <cit.>.
In control theory and reinforcement learning, state abstraction has been widely studied as a way of reducing the complexity of computing the optimal controller or optimal value function <cit.>.
Additionally, a related direction of research, called representation learning, tries to construct basis functions for representing high-dimensional value functions. Different Laplacian-based methods have been proposed to generate a surrogate for the exact transition operator <cit.>.


Matrix factorization is an optimization framework that decomposes a matrix into a product of two or more matrices <cit.>. In general matrix factorization framework, as opposed to spectral decomposition, one can impose additional desired structural properties. Common structural properties are those of being low-rank or sparse that can be promoted by nuclear norm regularization and ℓ_1-norm regularization, respectively. Furthermore, matrix factorization is typically amenable to efficient gradient descent solutions.
The problem of decomposing a matrix into a product of factors arises in different applications such as learning Markov models <cit.> and bioinformatics <cit.>. What makes matrix factorization methods appealing is the fact that their solution complexity depends on the rank of the factors which is typically much smaller than the input matrix. 


Recovery of a low-rank probability transition matrix has been considered in <cit.>. The majority of the proposed methods are based on spectral decomposition framework. 
In <cit.>, the authors use the notion of anchor states to enhance the interpretability of the meta states and further, develop a method of learning a soft-aggregation model from a set of system's trajectories.
In contrast, we propose an optimization formulation for the task of learning a low-dimensional representation of a known Markov chain. This formulation can easily induce different desired structural properties, such as sparsity, by imposing additional constraints. Furthermore, it enables an efficient solution that relies on block coordinate gradient descent.



§ PROBLEM FORMULATION

In this section, we provide the outline of the related concepts and definitions, and formally state the problem of learning representations for Markov chains.




 §.§ Preliminaries



In this paper, we focus on time-homogeneous discrete-time Markov chains with finite states, as formally defined below.


	A Markov chain is a random process such that its evolution is characterized by a tuple = (, , ), where S is a finite set of states with cardinality |S| = n,  is an initial distribution over the states, and : ×→ [0, 1] ⊆ is a probabilistic transition function such that for all s ∈, ∑_s'∈(s,s')=1.

A finite path in  is a realization of a finite-length sequence X_0,X_1,… of states, denoted by = x_0 x_1 x_2 …, such that x_0 is in the support of  and ∀ i ∈ℤ: (x_i,x_i+1) > 0. Using the Markovian property, the probability of sampling = x_0 x_1 x_2 … x_T is

    () = (x_0) ∑_t=1^T(x_t-1,x_t),

and the probability of going from state s_i at time step t to state s_j, in m time steps, is (X_t+m=s_j | X_t=s_i) = p_ij^(m), where p_ij^(m) = [P^m]_ij is an entry of the m-step transition matrix.


Runnenburg <cit.> introduced the notion of Markov chains with small rank as a type of dependence that is close to independence. Hoekstra <cit.> has further analyzed properties of Markov chains with small rank. Next, we provide the formal definition of the nonnegative rank of a Markov chain that will later motivate the proposed factorization.


	Let  denote the transition matrix of a Markov chain . The nonnegative rank of  is the smallest k ∈ for which the following decomposition exists:
	
    (X_t+1 | X_t) = ∑_l=1^k f_l(X_t) g_l(X_t+1),

	where f_1,f_2,…,f_k and g_1,g_2,…,g_k are real-valued functions mapping  to _+.

In particular, f_1,f_2,…,f_k denote the left Markov features and g_1,g_2,…,g_k denote the right Markov features. Without loss of generality, one can assume that the left and right Markov features are probability mass functions. Notice that if the nonnegative rank of  is k, it holds that () ≤ k <cit.>.
In other words, the nonnegative rank of a Markov chain upperbounds the rank of its transition matrix.

The next proposition states a critical property of nonnegative rank of Markov chains that we will later exploit.

	The nonnegative rank of a Markov chain is k if and only if there exists ∈_+^n × k, ∈_+^k × k, and ∈_+^k × n such that =, where U, P, and  are stochastic matrices, i.e., =, =, and =.

In Proposition <ref>,  resembles a surrogate lower-dimensional model for the transition matrix  if  admits a low nonnegative rank. We refer to the state space and the transitions in this abstract model as kernel space and kernel transition, respectively. 
In Proposition <ref>,  maps the states of the original Markov chain into the kernel space,  indicates the kernel transition, and  maps the kernel states back to the original states. 
In control and reinforcement learning, rows of  correspond to aggregation distributions while rows of  correspond to disaggregation distributions <cit.>.

The strength of the model abstraction setting introduced in Proposition <ref> is that it is independent of the downstream analysis to be performed on the Markov chain. Therefore, once such abstraction is found, it can be used for accelerating different types of analyses. For instance, in the next proposition, we show how the m-step transition matrix can be computed more efficiently by using the factorized model.

	Given a Markov chain = (, , ), assume that a perfect low-rank decomposition of the transition matrix exists such that =, ∈_+^k × k. Let K =. Then, the m-step transition matrix of  can be computed by
	
    (X_t+m | X_t) = ∑_l_1=1^k∑_l_2=1^k_X_t,l_1 [ K^m-1]_l_1 l_2_l_2,X_t+m.


Hence, one can reduce the complexity of computing the m-step transition matrix from 𝒪 (mn^2) to 𝒪 (mk^2). 




 §.§ Problem Statement


In many applications, often the Markov chain model of a system has underlying structural properties, including possession of a low-rank or sparse transition function. Motivated by this fact, we seek an abstraction of a Markov chain in a low-dimensional kernel space. To that end, we need to find the mapping from the original state space to the state space of the abstracted (surrogate) model as well as the inverse mapping from the state space of the surrogate model to the original state space. Figure <ref> demonstrates a pictorial overview of the mapping between the spaces. Essentially, a transition in the original model can be represented through three steps: step 1 maps a state in the original model to meta states in the surrogate model; step 2 is a transition inside the surrogate model; and step 3 is a mapping from a meta state back to original states. Additionally, we would like the meta states to be representative of a small subset of states. This property means that each meta state should be connected to as few states as possible and hence, we impose that by looking for sparse mappings between the spaces. 


	Given a Markov chain = (, , ), we aim to find a kernel space and kernel transition, denoted by (, ), with lower dimensionality, i.e., || ≪ ||. Further, we look for a sparse bidirectional mapping (,) where U represents the mapping from  to  while  represents the mapping from  to . The surrogate model (, ) along the bidirectional mapping (,) must be such that the decomposition property = holds.





§ APPROACH

Let n = || to be the size of the high-dimensional state space and k denote the nonnegative rank of the . Let 𝒟: ^n× n×^n× n→_+ denote a metric on the space of n× n matrices. As we discussed in Section <ref>, we seek to promote sparsity patterns in the rows of the bidirectional mapping (,). Therefore, in order to find the factorization in Problem 1, we propose the following optimization task

    ≥ 0,≥ 0,≥ 0min  𝒟(,)
       s.t.  ∑_j = 1^k _ij = 1,   u_i_0 ≤ s_i^(u), ∀ i ∈ [n], 
           ∑_j = 1^k _ℓ j = 1,   ∀ℓ∈ [k],
           ∑_j = 1^n _ℓ j = 1,    v_ℓ_0 ≤ s_ℓ^(v), ∀ℓ∈ [k],

where ._0 is the so-called ℓ_0-norm and returns the number of nonzero entries of its argument, and {s_i^(u)} and {s_ℓ^(v)} are positive integers that determine the extent of the desired sparsity structure in the rows of  and . Notice that because of the  ℓ_0-norm constraints, (<ref>) is a combinatorial optimization problem and generally NP-hard to solve. Therefore, we propose to relax these constraints by using the ℓ_1-norm which is the convex envelope of the ℓ_0-norm and is known to promote sparsity in the solution of an optimization problem. Following this idea and by specifying 𝒟(X,Y) = 1/2X-Y^2_F as the metric, we consider the relaxed and regularized problem

    ≥ 0,≥ 0,≥ 0min  1/2-^2_F+λ_u_1+λ_v_1
       s.t.   = ,    = ,    = .

Here, λ_u>0 and λ_v>0 are the regularization parameters that determine the sparsity level of the rows of the bidirectional mapping matrices  and . Recall that sparsity of  means that each state is mapped to one (or few) meta state(s). On the other hand, sparsity of  asks for mapping of a meta state to a few states. This sparsity promoting terms ensure that the meta states are a good representative of their corresponding states. Additionally, the constraints of the optimization make sure that each of matrices are a stochastic matrix, i.e., can be interpreted as a transition matrix.








The objective function in (<ref>) consists of a convex function and hence 
is convex in each of the matrices when the other matrices are fixed. However, due to the fact that the first term in the objective function contains a product of the unknowns, (<ref>) is generally a nonconvex program. 
Notice that even if P is perfectly decomposable into the corresponding factors, any permutation of the low-rank abstraction  is also a solution. Therefore, in general, Problem <ref> has at least k! global optima.

To facilitate a computationally efficient search for the solution of (<ref>), we rely on a modified gradient search algorithm which exploits the special structures of , , and . The algorithm (summarized as Algorithm <ref>) is essentially a block coordinate gradient descent (BCGD) method that alternatively updates matrices , , and  in an iterative fashion starting from an initial point (_0,_0,_0). That is, in (t+1)st iteration (t = 0,…,T-1 where T is the total number of iterations), given (_t,_t,_t) we optimize with respect to  to find _t+1. Similarly, we find updated _t+1 and _t+1 using the values (_t+1,_t,_t) and (_t+1,_t+1,_t), respectively. 

In Algorithm <ref>, 𝒯_η(x) = (|x|-η)_+sgn(x) is the so-called
shrinkage-thresholding operator that acts on each element of the given matrix, and Π__d(.) denotes the projection operator that projects each row of its argument onto the probability simplex in ^d. This projection can be efficiently computed by the method of <cit.> that we summarize in Algorithm <ref> for completeness.











 §.§ Convergence Analysis of BCGD


In this section, we analyze the convergence properties of BCGD. Specifically, in Theorem <ref> we
establish that given judicious choices of the step sizes, the value of the objective function in (<ref>) decreases as one alternates between updating the factor matrices  which in turn implies that the BCGD algorithm converges to a stationary point of the nonconvex optimization task in (<ref>).

Assume the step sizes of Algorithm 1 satisfy

    α_t=C_1(_t)_F^2/(_t)_t_t_F^2,


    β_t=C_2(_t)_F^2/_t+1(_t)_t_F^2,


    γ_t=C_3(_t)_F^2/_t+1_t+1(_t)_F^2,

where C_1,C_2,C_3 ∈(0,2). Then, the solution (^∗,^∗,^∗) found by the BCGD scheme	is a stationary point of (<ref>). 


Let f(,,)= 1/2-^2_F.
For the proposed algorithm to converge, it must hold that 

    f(_t+1,_t+1,_t+1) ≤ f(_t,_t,_t),

for all t. Note that since 𝒯_η(x) and Π__d are projections onto convex sets of constraints (the former being the projection operator onto the ℓ_1 ball), following a similar analysis as those in the proofs of the projected gradient descent and the iterative shrinkage-thresholding algorithms (ISTA) <cit.> 

    f(_t+1,_t,_t) ≤ f(_t+1/2,_t,_t),
 
    f(_t+1,_t+1,_t) ≤ f(_t+1,_t+1/2,_t),
 
    f(_t+1,_t+1,_t+1) ≤ f(_t+1,_t+1,_t+1/2).
 Thus, it suffices to show 

    f(_t+1/2,_t,_t) ≤ f(_t,_t,_t),


    f(_t+1,_t+1/2,_t) ≤ f(_t+1,_t,_t),


    f(_t+1,_t+1,_t+1/2) ≤ f(_t+1,_t+1,_t).

We now show that under (<ref>), the sufficient condition (<ref>) holds. Proof of the parts that (<ref>) and (<ref>) hold under (<ref>) and (<ref>), respectively, follows a similar argument.

Note that

    f(_t+1/2,_t,_t) - f(_t,_t,_t)
      =1/2-_t_t + α_t(_t)_t_t^2_F
    111111111111111-1/2-_t_t_t^2_F
      = α_tTr((-_t_t)^⊤(_t)_t_t)
    
    +α_t^2/2(_t)_t_t^2_F.

Now, consider the first term in the last line of (<ref>). Following straightforward linear algebra, we obtain

    Tr((-_t_t)^⊤(_t)_t_t)
       =-Tr((-_t_t)^⊤(-_t_t_t)_t^⊤_t^⊤_t_t)
       =-Tr(_t_t(-_t_t)^⊤(-_t_t_t)_t^⊤_t^⊤)
       =-(-_t_t_t)_t^⊤_t^⊤_F^2=-(_t)^⊤_F^2.

Therefore,

    f(_t+1/2,_t,_t) - f(_t,_t,_t)
        =α_t^2/2(_t)_t_t^2_F
    -α_t(_t)^⊤_F^2
    =(C_1^2/2-C_1)(_t)^⊤_F^4/(_t)_t_t_F^2,

where the last equality follows according to the definition of α_t in (<ref>). It is now clear that if C_1 ∈ (0,2) it must be the case that (<ref>) holds, which in turn implies convergence of Algorithm 1.





 §.§ Computational Complexity of BCGD


The computational complexity of the proposed BCGD algorithm is analyzed next. Note that the determining factor for cost per iteration of Algorithm <ref> is computation of the gradients. Finding (_t) incurs 𝒪(nk) as it contains matrix products between k× k and k× n matrices. Similarly, (_t) and (_t) require 𝒪(nk) computational costs. Thus, Algorithm <ref> incurs a linear complexity of 𝒪(nkT).



§ SIMULATION RESULTS

We implemented the proposed BCGD algorithm for nonnegative matrix factorization in Python.[The code is available at https://github.com/MahsaGhasemi/state-abstraction]
We evaluated the performance of the proposed abstraction solution for a variety of parameters. In particular, we investigated the effect of sparsity-promoting term on the approximation quality of the factorized transition matrix. We also ran the algorithm for different values of step size and compared the convergence of BCGD. 
For these simulations, we generated a transition matrix  of size 100 × 100 for the original Markov chain. The transition matrix has a rank of 25 and is constructed by multiplying three stochastic matrices. Each stochastic matrix is generated by independently sampling its rows by a uniform sampling from the simplex of proper size. We set the number of iterations of BCGD to 1000. If the difference between two consecutive instances of a factor, i.e., _t, _t, or _t falls below a threshold 10^-8 of their magnitude, the algorithm terminates, where the magnitudes are measured by Frobenius norm. Similarly, if the difference between two consecutive values of the objective function falls below 10^-8, the algorithm terminates. We run the simulation with each set of parameters for 10 independent instances and report the average values along the standard deviations.
All simulations were run on a machine with 2.0 GHz Intel Core i7-4510U CPU and with 8.00 GB RAM.




 §.§ Effect of Regularization Parameter on Performance


One of the key differences of the proposed abstraction formulation is the integration of a ℓ_1-norm regularization term in the optimization objective. This term promotes sparsity for the bidirectional mapping between the original and the kernel space. Intuitively, this sparsity ensures that the meta states in the kernel space are representative of a small number of original states. The sparsity level of the mappings depends on regularization parameters λ_u and λ_v. 

In Figure <ref>, we demonstrate the effect of these parameters on the quality of the solutions. We compare four different values, more specifically, λ = λ_u = λ_v ∈{0,0.001,0.005,0.01}. The results in Figure <ref>(a) show that a careful selection of λ, while increasing sparsity, does not affect the quality of the decomposition in terms of the approximation error, computed by -^2_F. However, a large λ leads to high approximation error. Therefore, one has to find the right trade-off between lower approximation error and higher sparsity of the mappings. Furthermore, we can see that error significantly reduces once the size of the kernel transition is set to values higher than 20. Above that value, the approximation error only slightly changes. Therefore, the proposed algorithm is successful in identifying a low-rank representation.

Figure <ref>(b) depicts the running time of the BCGD algorithm. As derived in Section IV-B, the running time is linear with respect to the kernel size and the results reflect that. Furthermore, the addition of ℓ_1-norm regularization term has negligible effect on the running time. Note that for λ = 0.01, the algorithm terminates early as it cannot improve the objective value sufficiently. Furthermore, for this choice of λ, we can observe high variance in the running time due to the varying sparsity degree of the randomly generated matrices.




 §.§ Effect of Step Size on Convergence


In Section IV-A, we derived necessary conditions on the step sizes α, β, and γ for the convergence of the BCGD algorithm. In this section, we show the sensitivity of the convergence to different values of the step size. To that end, we ran the algorithm for different values of step size that we kept constant throughout the run. In particular, we ran the algorithm for α = β = γ∈{0.002,0.02,0.2}. Figure <ref> depicts the evolution of error over the course of 500 iterations of BCGD. While the algorithm converges for all three values, the smaller step sizes achieve lower approximation error at the end. We also observed that the algorithm would often diverge for step sizes over 0.2.





§ CONCLUSION
 
We studied the problem of approximating a large-scale Markov chain by a surrogate model in a low-dimensional state space, called kernel space. We proposed a nonnegative matrix factorization formulation that learns a low-rank kernel transition as well as a pair of forward and backward mappings while promoting a sparse connection between the high and low-dimensional states. We showed that the formulated optimization is amenable to an efficient iterative solution that converges to a stationary solution under a judicious schedule of step sizes. 
As part of future work, we aim to extend the proposed matrix factorization formulation to model reduction of Markov decision processes. Furthermore, we would like to evaluate the abstracted Markov decision process in different analyses, including model checking and value function approximation.


ieeetr


