Optimal Clustering in Anisotropic Gaussian Mixture Models

We first consider a GMM where covariance matrices of different clusters are unknown but are assumed to be equal to each other. The data generating progress can be displayed as follow:

It is called Stretched Mixture Model in [23] as the density of $Y_{j}$ is elliptical. Throughout the paper, we call it Model 1 for simplicity and to distinguish it from a more complicated model that will be introduced in Section 3. The goal is to recover the underlying cluster assignment vector $z^{*}$ from $Y$.

We first establish the minimax lower bound for the clustering problem under Model 1.

Theorem 2.1 allows the cluster numbers $k$ to grow with $n$ and shows that $\textsf{SNR}\rightarrow\infty$ is a necessary condition to have a consistent clustering if $k$ is a constant. Theorem 2.1 holds for any arbitrary $\{\theta^{*}_{a}\}_{a\in[k]}$ and $\Sigma^{*}$, and the minimax lower bound depend on them through SNR. The parameter space is only for $z^{*}$ while $\{\theta^{*}_{a}\}_{a\in[k]}$ and $\Sigma^{*}$ are fixed. Hence, (6) can be interpreted as a pointwise result, and it captures precisely the explicit dependence of the minimaxity on $\{\theta^{*}_{a}\}_{a\in[k]}$ and $\Sigma^{*}$.

Theorem 2.1 is closely related to the Linear Discriminant Analysis (LDA). If there are only two clusters, and if the centers and the covariance matrix are known, then estimating each $z^{*}_{j}$ is exactly the task of LDA: we want to figure out which normal distribution an observation $Y_{j}$ is generated from, where the two normal distributions have different means but the same covariance matrix. In fact, this is also how Theorem 2.1 is proved: we will first reduce the estimation problem of $z^{*}$ into a two-point hypothesis testing problem for each individual $z^{*}_{j}$, the error of which is given in Lemma 2.1 by the analysis of LDA, and then aggregate all the testing errors together.

In the following lemma, we give a sharp and explicit formula for the testing error of the LDA. Here we have two normal distributions $\mathcal{N}(\theta^{*}_{1},\Sigma^{*})$ and $\mathcal{N}(\theta^{*}_{2},\Sigma^{*})$ and an observation $X$ that is generated from one of them. We are interested in estimating which distribution it is from. By Neyman-Pearson lemma, it is known that the likelihood ratio test ${\mathbb{I}\left\{{2(\theta^{*}_{2}-\theta^{*}_{1})^{T}(\Sigma^{*})^{-1}X\geq\theta_{2}^{*T}(\Sigma^{*})^{-1}\theta^{*}_{2}-\theta^{*T}_{1}(\Sigma^{*})^{-1}\theta^{*}_{1}}\right\}}$ is the optimal testing procedure. Then by using the Gaussian tail probability, we are able to obtain the optimal testing error, the lower bound of which is given in Lemma 2.1.

With the help of Lemma 2.1, we have a geometric interpretation of SNR. In the left panel of Figure 1, we have two normal distributions $\mathcal{N}(\theta^{*}_{1},\Sigma^{*})$ and $\mathcal{N}(\theta^{*}_{2},\Sigma^{*})$ for $X$ to be generated from. The black line represents the optimal testing procedure $\phi$ displayed in Lemma 2.1 that divides the space into two half-spaces. To calculate the testing error, we can make a transformation $X^{\prime}=(\Sigma^{*})^{-\frac{1}{2}}(X-\theta_{1}^{*})$ so that the two normal distributions become isotropic: $\mathcal{N}(0,I_{d})$ and $\mathcal{N}((\Sigma^{*})^{-\frac{1}{2}}(\theta^{*}_{2}-\theta^{*}_{1}),I_{d})$ as displayed in the right panel. Then the distance between the two centers are $\|(\Sigma^{*})^{-\frac{1}{2}}(\theta^{*}_{2}-\theta^{*}_{1})\|$, and the distance between a center and the black curve is half of it. Then $\mathbb{P}_{H_{0}}(\hat{\phi}=1)$ is the probability of $\mathcal{N}(0,I_{d})$ in the grayed area, which is equal to $\exp(-(1+o(1))\|(\Sigma^{*})^{-\frac{1}{2}}(\theta^{*}_{2}-\theta^{*}_{1})\|^{2}/8)$ by Gaussian tail probability. As a result, $\|(\Sigma^{*})^{-\frac{1}{2}}(\theta^{*}_{2}-\theta^{*}_{1})\|$ is the effective distance between the two centers of $\mathcal{N}(\theta^{*}_{1},\Sigma^{*})$ and $\mathcal{N}(\theta^{*}_{2},\Sigma^{*})$ for the clustering problem, considering the geometry of the covariance matrix. Since we have multiple clusters, SNR defined in (2) can be interpreted as the minimum effective distances among the centers $\{\theta^{*}_{a}\}_{a\in[k]}$ considering the anisotropic structure of $\Sigma^{*}$ and it captures the intrinsic difficulty of the clustering problem.

In this section, we will propose a computationally feasible and rate-optimal procedure for clustering under Model 1. Summarized in Algorithm 1, the proposed algorithm is a variant of the Lloyd algorithm. Starting from some initialization, it updates the estimation of the centers $\{\theta^{*}_{a}\}_{a\in[k]}$ (in (7)), the covariance matrix $\Sigma^{*}$ (in (8)), and the cluster assignment vector $z^{*}$ (in (9)) iteratively. It differs from the Lloyd’s algorithm in the sense that Lloyd’s algorithm is for isotropic GMMs without the covariance matrix update (8). In addition, in (9) it updates the estimation of $z^{*}_{j}$ by $\mathop{\rm argmin}_{a\in[k]}(Y_{j}-\theta_{a}^{(t)})^{T}(Y_{j}-\theta_{a}^{(t)})$ instead. To distinguish them from each other, we call the classical Lloyd’s algorithm as the vanilla Lloyd’s algorithm, and name Algorithm 1 as the adjusted Lloyd’s algorithm, as it is adjusted to the unknown and anisotropic covariance matrix.

Algorithm 1 can also be interpreted as a hard EM algorithm. If we apply the Expectation Maximization (EM) for Model 1, we will have an M step for estimating parameters $\{\theta^{*}_{a}\}_{a\in[k]}$ and $\Sigma^{*}$ and an E step for estimating $z^{*}$. It turns out the updates on the parameters (7) - (8) are exactly the same as the updates of EM (M step). However, the update on $z^{*}$ in Algorithm 1 is different from that in the EM. Instead of taking a conditional expectation (E step), we also take a maximization in (9). As a result, Algorithm 1 consists solely of M steps for both the parameters and $z^{*}$, which is known as a hard EM algorithm.

In Theorem 2.2, we give a computational and statistical guarantee of the proposed Algorithm 1. We show that starting from a decent initialization, within $\log n$ iterations, Algorithm 1 achieves an error rate $\exp\left(-(1+o(1))\textsf{SNR}^{2}/8\right)$ which matches with the minimax lower bound given in Theorem 2.1. As a result, Algorithm 1 is a rate-optimal procedure. In addition, the algorithm is fully adaptive to the unknown $\{\theta^{*}_{a}\}_{a\in[k]}$ and $\Sigma^{*}$. The only information assumed to be known is $k$ the number of clusters, which is commonly assumed to be known in clustering literature [15, 8, 14]. The theorem also shows that the number of iterations to achieve the optimal rate is at most $\log n$, which provides implementation guidance to practitioners.

We have remarks on the assumptions of Theorem 2.2. We allow the number of clusters $k$ to grow with $n$. When $k$ is a constant, the assumption on $\textsf{SNR}\rightarrow\infty$ is the necessary condition to have a consistent recovery of $z^{*}$ according to the minimax lower bound presented in Theorem 2.1. The assumption on $\Sigma^{*}$ is to make sure the covariance matrix is well-conditioned. The dimensionality $d$ is assumed to be at most $O(\sqrt{n})$, an assumption that is stronger than that in [15, 8, 14] which only needs $d=O(n)$. This is due to that compared to these papers, we need to estimate the covariance matrix $\Sigma^{*}$ and to have a control on the estimation error $\|\Sigma^{(t)}-\Sigma^{*}\|$.

The requirement for the initialization $\ell(z^{(0)},z^{*})=o(n/k)$ can be fulfilled by simple procedures. A popular choice is the vanilla Lloyd’s algorithm the performance of which is studied in [15, 8]. Since $\epsilon_{j}$ are sub-Gaussian random variables with proxy variance $\lambda_{\max}$, [8] implies the vanilla Lloyd’s algorithm output $\hat{z}$ satisfies $\ell(\hat{z},z^{*})\leq n\exp(-(1+o(1))\Delta^{2}/(8\lambda_{\max}))$ with probability at least $1-\exp(-\Delta)-n^{-1}$, under the assumption that $\textsf{SNR}/k\rightarrow\infty$. Note that [8] is for isotropic GMMs, but its results can be extended to sub-Gaussian mixture models with nearly identical proof. Then we have $\ell(\hat{z},z^{*})=o(n/k)$, as $\Delta^{2}/\lambda_{\max}$ and $\textsf{SNR}^{2}$ are both in the same order under the assumption $\textsf{SNR}/k\rightarrow\infty$. As a result, we immediately have the following corollary.

In this section, we study the GMM with covariance matrices from each cluster unknown and not necessarily equal to each other. The data generation process can be displayed as follow,

We call it Model 2 throughout the paper to distinguish it from Model 1 studied in Section 2. The difference between (10) and (1) is that we now have $\{\Sigma^{*}_{a}\}_{a\in[k]}$ instead of a shared $\Sigma^{*}$. We consider the same loss functions as in (4) and (5).

We first establish the minimax lower bound for the clustering problem under Model 2.

Despite that the statement of Theorem 3.1 looks similar to that of Theorem 2.1, the two minimax lower bounds are different from each other due to the discrepancy in the dependence of the centers and the covariance matrices in $\textsf{SNR}^{\prime}$ and SNR. By the same argument as in Section 2.2, the minimax lower bound established in Theorem 3.1 is closely related to the Quadratic Discriminant Analysis (QDA) between two normal distributions with different means and different covariance matrices.

From Lemma 3.1, we can have a geometric interpretation of $\textsf{SNR}^{\prime}$. In the left panel of Figure 2, we have two normal distributions $\mathcal{N}(\theta^{*}_{1},\Sigma_{1}^{*})$ and $\mathcal{N}(\theta^{*}_{2},\Sigma_{2}^{*})$ where $X$ can be generated from. The black curve represents the optimal testing procedure $\phi$ displayed in Lemma 3.1. Since $\Sigma_{1}^{*}$ is not necessarily equal to $\Sigma_{2}^{*}$, the black curve is not necessarily a straight line. If $\mathbb{H}_{0}$ is true, then the probability for $X$ to be incorrectly classified is when $X$ falls in the gray area, which is $\mathbb{P}_{H_{0}}(\hat{\phi}=1)$. To calculate it, we can make a transformation $X^{\prime}=(\Sigma^{*}_{1})^{-\frac{1}{2}}(X-\theta^{*}_{1})$. Then displayed in the right panel of Figure 2, the two distributions become $\mathcal{N}(0,I_{d})$ and $\mathcal{N}((\Sigma^{*}_{1})^{-\frac{1}{2}}(\theta^{*}_{2}-\theta^{*}_{1}),(\Sigma^{*}_{1})^{-\frac{1}{2}}\Sigma^{*}_{2}(\Sigma^{*}_{1})^{-\frac{1}{2}})$, and the optimal testing procedure ${\mathbb{I}\left\{{X^{\prime}\in\mathcal{B}_{1,2}}\right\}}$. As a result, in the right panel of Figure 2, $\mathcal{B}_{1,2}$ represents the space colored by gray, and the black curve is its boundary. Then $\mathbb{P}_{H_{0}}(\hat{\phi}=1)$ is equal to $\mathbb{P}\left(\mathcal{N}(0,I_{d})\in\mathcal{B}_{1,2}\right)$, which can be shown to be determined by the minimum distance between the center of $\mathcal{N}(0,I_{d})$ and the space $\mathcal{B}_{1,2}$. Denote the minimum distance by $\textsf{SNR}_{1,2}^{\prime}/2$, by Lemmas 7.10 and Lemma 7.11, we can show $\mathbb{P}\left(\mathcal{N}(0,I_{d})\in\mathcal{B}_{1,2}\right)=\exp(-(1+o(1))\textsf{SNR}_{1,2}^{{}^{\prime}2}/8)$. As a result, the $\textsf{SNR}^{\prime}$ can be interpreted as the minimum effective distance among the centers $\{\theta^{*}_{a}\}_{a\in[k]}$ considering the anisotropic and heterogeneous structure of $\{\Sigma^{*}_{a}\}_{a\in[k]}$ and it captures the intrinsic difficulty of the clustering problem under Model 2.

In this section, we will propose a computationally feasible and rate-optimal procedure for clustering under Model 2. Similar to Algorithm 1, the proposed Algorithm 2 can be seen as a variant of the Lloyd’s algorithm that is adjusted to the unknown and heterogeneous covariance matrices. It can also be interpreted as a hard EM algorithm under Model 2. Algorithm 2 differs from Algorithm 1 in (13) and (14), as now there are $k$ covariance matrices to be estimated.

In Theorem 3.2, we give a computational and statistical guarantee of the proposed Algorithm 2. We show that provided with some decent initialization, Algorithm 2 is able to achieve the minimax lower bound within $\log n$ iterations. The assumptions needed in Theorem 3.2 are similar to those in Theorem 3.2, except that we require stronger assumptions on $k$ and the dimensionality $d$ since now we have $k$ (instead of one) covariance matrices to be estimated. In addition, by $\max_{a,b\in[k]}\lambda_{d}(\Sigma^{*}_{a})/\lambda_{1}(\Sigma^{*}_{b})=O(1)$ we not only assume each of the $k$ covariance matrices is well-conditioned, but also assume they are comparable to each other.

The vanilla Lloyd’s algorithm can be used as the initialization for Algorithm 2. Under the assumption that $\lambda_{d}(\Sigma^{*})/\lambda_{1}(\Sigma^{*})=O(1)$, Model 2 is also a sub-Gaussian mixture model. By the same argument as in Section 2.3 we have the following corollary.