Near-Optimal Clustering in Mixture of Markov Chains

1. Problem Statement

The paper addresses the Mixture of Markov Chains (MMC) problem.

• Setup: We are given $T$ trajectories, each of length $H$. Each trajectory is generated by one of $K$ unknown, ergodic Markov chains defined over a finite state space of size $S$.
• Goal: The objective is to cluster the $T$ trajectories according to their generating Markov chain (i.e., recover the ground-truth decoding function $f: [T] \to [K]$) without prior knowledge of the transition kernels or the number of clusters (in the parameter-free setting).
• Context: Unlike static data clustering (e.g., Gaussian Mixture Models or Stochastic Block Models), MMC involves temporal data where the signal is distributed across the trajectory. The challenge lies in balancing the trajectory length $H$ and the number of trajectories $T$ to achieve accurate clustering.

2. Methodology: A Two-Stage Algorithm

The authors propose a computationally tractable, two-stage algorithm that achieves near-optimal clustering error rates without requiring prior knowledge of model parameters (like the pseudo-spectral gap or cluster sizes).

Stage I: Initial Spectral Clustering via L-Embedding

The core innovation of Stage I is a new injective Euclidean embedding for ergodic Markov chains, termed the L-embedding.

• Definition: For a Markov chain with transition matrix $P$ and stationary distribution $\pi$, the embedding maps the chain to a vector in $\mathbb{R}^{S^2}$ defined as:
  $$L(M) = \text{vec}\left( \text{diag}(\pi)^{1/2} P \right)$$
• Empirical Estimator: Given a trajectory, the empirical L-embedding $\hat{L}$ normalizes transition counts by the square root of the visitation counts of the source states. This specific normalization is crucial for variance stabilization.
• Spectral Clustering:
  1. Construct an empirical data matrix $\hat{W}$ by row-concatenating the L-embeddings of all $T$ trajectories.
  2. Perform Singular Value Decomposition (SVD) on $\hat{W}$.
  3. If $K$ is known, apply K-means to the top-$K$ singular vectors. If $K$ is unknown, use adaptive thresholding on singular values followed by a density-based clustering step to estimate $K$ and assign initial cluster labels.
• Key Property: The L-embedding allows for sharp concentration results. The error bound for the empirical embedding does not suffer from a detrimental dependence on the inverse of the minimum stationary probability ($\pi_{\min}^{-1}$), which is a common bottleneck in Markov chain estimation.

Stage II: Likelihood-Based Refinement

Stage I provides a "good enough" initial clustering. Stage II refines this using a trajectory-wise Maximum Likelihood Estimator (MLE).

• Process:
  1. Estimate the transition kernels $\hat{p}^{(k)}$ for each cluster $k$ using the trajectories assigned to it in Stage I.
  2. Reassign each trajectory to the cluster $k$ that maximizes the log-likelihood of the observed trajectory under the estimated kernel $\hat{p}^{(k)}$.
• Theoretical Basis: This step is essentially a hard Expectation-Maximization (EM) iteration. It leverages the Neyman-Pearson lemma to achieve optimal separation based on the Kullback-Leibler (KL) divergence.

3. Key Contributions

A. Instance-Dependent Lower Bound

The authors derive a high-probability lower bound on the clustering error rate.

• Metric: The difficulty of the problem is governed by a quantity $D$, defined as the minimum stationary-weighted KL divergence between the transition kernels of any two distinct chains:
  $$D = \min_{k \neq k'} \sum_{s \in S} \pi^{(k)}(s) \text{KL}\left(p^{(k)}(\cdot|s), p^{(k')}(\cdot|s)\right)$$
• Result: The lower bound shows that the misclassification rate scales as $\exp(-2HD)$. This establishes the fundamental limit: clustering is possible if the product of trajectory length and the KL gap is sufficiently large.

B. Near-Optimal Algorithm

The proposed two-stage algorithm achieves an error rate of:
$$\text{Error} \lesssim T \exp\left(-C \gamma_{ps} H D_{\pi}\right)$$
where $\gamma_{ps}$ is the pseudo-spectral gap and $D_{\pi}$ is the stationary divergence.

• Significance: The upper bound matches the lower bound up to a constant factor involving the spectral gap and regularity parameters. This proves the algorithm is near-optimal.
• Parameter-Free: The algorithm adapts to the problem difficulty without requiring prior knowledge of $K$, $\gamma_{ps}$, or cluster sizes.

C. Theoretical Analysis of Gaps

The paper rigorously analyzes the gaps between different separability metrics:

• $D_{\pi}$ (KL Gap): The information-theoretic optimal gap.
• $\Delta_W^2$ (L2 Gap): The geometric gap induced by the L-embedding.
• $\alpha \Delta^2$ (Kausik et al. Gap): The gap used in prior work, which relies on separation at a single "good" state.
• Finding: The authors show that for certain instances (specifically "Cyclic-Bump" chains), $D_{\pi}$ can be larger than $\Delta_W^2$ by a factor of $S$, and $\Delta_W^2$ can be larger than $\alpha \Delta^2$ by a factor of $S$. This explains why their method outperforms previous approaches that rely on single-state separation.

4. Results and Experiments

• Synthetic Data: Experiments on a synthetic "Cyclic-Bump" MMC instance demonstrate that the proposed algorithm significantly outperforms the state-of-the-art method by Kausik et al. (2023). The proposed method achieves near-zero error with shorter trajectories, whereas Kausik's method requires much longer trajectories or fails entirely depending on hyperparameter tuning.
• Real-World Data: The algorithm was tested on the Last.fm 1K dataset (music listening logs). It successfully clustered users based on their listening patterns, significantly outperforming the baseline method.
• Ablation Studies:
  • EM Iterations: While theory requires only one refinement step, empirical results show that multiple EM iterations further reduce error, especially in low-data regimes.
  • Unknown $K$: The algorithm successfully recovers the number of clusters when $K$ is unknown, though it is sensitive to the thresholding hyperparameters.
  • Spectral Gap: Experiments confirm that the convergence rate is indeed dependent on the pseudo-spectral gap $\gamma_{ps}$, validating the theoretical bounds.

5. Significance and Limitations

Significance:

• Theoretical Breakthrough: It provides the first instance-dependent lower bound and a matching near-optimal algorithm for MMC clustering.
• New Embedding: The L-embedding is a novel contribution that enables sharp concentration bounds for spectral clustering of Markov chains, potentially useful beyond this specific problem.
• Practical Utility: The method is robust, parameter-free, and effective on real-world temporal data.

Limitations:

• Computational Complexity: The SVD step in Stage I has a complexity of $O(T S^2 \min(T, S^2))$, which may be prohibitive for very large state spaces ($S$) or massive numbers of trajectories ($T$).
• Regularity Assumption: The theoretical guarantee for Stage II relies on an $\eta$-regularity assumption (Assumption 2), which bounds the ratio of transition probabilities between chains. While the authors argue this is rarely a practical issue (as extreme non-uniformity makes chains easier to distinguish), it is a theoretical constraint.
• Gap in Bounds: There remains a small gap between the upper and lower bounds regarding the dependence on the spectral gap $\gamma_{ps}$ in the exponent, which the authors attribute to the use of Markovian Bernstein inequalities.

In conclusion, this paper establishes a new standard for clustering trajectories in Markov mixtures, bridging the gap between information-theoretic limits and practical, parameter-free algorithms.