Uncovering Structure in the MNIST Dataset with Principal Component Analysis
The MNIST dataset[1] is considered one of the foundational datasets in machine learning, particularly in the domain of computer vision. Consisting of 70,000 grayscale images of handwritten digits (60,000 for training and 10,000 for testing), each 28×28 pixels, MNIST has served as a benchmark and a starting point for developing image classification algorithms for over two decades.
While MNIST can now be classified with over 99% accuracy using convolutional neural networks[2], it remains a useful dataset for understanding fundamental machine learning concepts. In this post, we‘ll dive deep into one such fundamental technique: Principal Component Analysis (PCA), and see how it can be used to uncover latent structure in the MNIST dataset.
Feature Extraction and Dimensionality Reduction
In machine learning, we often work with high-dimensional data, where each sample is represented by a large number of features. In the case of MNIST, each image is a 28×28 grid of pixels, which when flattened becomes a 784-dimensional vector. While neural networks can work directly with raw pixel values, classic machine learning algorithms often benefit from first transforming the data into a lower-dimensional representation that captures the essence of the original data.
This is where dimensionality reduction techniques like PCA come in. PCA is an unsupervised learning method that seeks to find a lower-dimensional projection of the data that preserves as much of the original variance as possible. Intuitively, it tries to find a small set of latent variables that can be linearly combined to reconstruct the original features.
PCA can be thought of as a feature extraction method, as the learned projection can be used to transform the original data into a new feature space. These extracted features often correspond to meaningful latent factors in the data. In the case of images, the principal components might correspond to fundamental shapes or patterns.
It‘s worth noting that PCA is a linear dimensionality reduction method, meaning it can only capture linear structure in the data. For highly nonlinear datasets, more advanced methods like autoencoders[3] or t-SNE[4] may be more appropriate. However, PCA remains widely used due to its simplicity, interpretability, and efficiency.
The Mathematics of PCA
At its core, PCA relies on the eigendecomposition of the data covariance matrix. Let‘s break this down step by step. Assume we have a data matrix $X$ of shape $(n, d)$, where $n$ is the number of samples and $d$ is the number of features.
-
First, we center the data by subtracting the mean of each feature:
$X_c = X – \frac{1}{n} \mathbf{1}_n \mathbf{1}_d^T X$
where $\mathbf{1}_n$ is an $n$-dimensional vector of ones. This ensures that the principal components pass through the origin.
-
Next, we compute the covariance matrix:
$\Sigma = \frac{1}{n-1} X_c^T X_c$
The covariance matrix is symmetric and positive semidefinite, with shape $(d, d)$.
-
We then perform eigendecomposition on the covariance matrix:
$\Sigma = U \Lambda U^T$
where $U$ is a $(d, d)$ orthogonal matrix whose columns are the eigenvectors of $\Sigma$, and $\Lambda$ is a diagonal matrix of the corresponding eigenvalues.
-
The eigenvectors in $U$ are ordered by their eigenvalues in descending order. To reduce to $k$ dimensions, we take the first $k$ columns of $U$, denoted $U_k$.
-
Finally, we project the centered data onto the top-$k$ eigenvectors:
$Z = X_c U_k$
The resulting $Z$ is an $(n, k)$ matrix representing the data in the reduced $k$-dimensional space.
The eigenvalues in $\Lambda$ tell us how much variance is explained by each principal component. The fraction of variance explained by the $i$-th component is given by:
$\frac{\lambdai}{\sum{j=1}^d \lambda_j}$
By plotting the cumulative explained variance as a function of the number of components, we can determine how many dimensions are needed to retain a desired percentage of the original variance.
PCA on the MNIST Dataset
Let‘s see how to apply PCA to the MNIST dataset using Python and the scikit-learn library[5]. First, we load the data:
from sklearn.datasets import fetch_openml
mnist = fetch_openml(‘mnist_784‘)
X, y = mnist.data, mnist.target The X matrix has shape (70000, 784), while the y vector contains the corresponding labels.
To apply PCA, we first center the data and then fit the PCA object:
from sklearn.decomposition import PCA
pca = PCA()
X_centered = X - X.mean(axis=0)
pca.fit(X_centered)We can then examine the explained variance ratio of each component:
plt.plot(np.cumsum(pca.explained_variance_ratio_))
plt.xlabel(‘Number of Components‘)
plt.ylabel(‘Cumulative Explained Variance‘)
We see that with just 50 components, we can explain over 90% of the total variance in the data! This suggests that there is significant redundancy in the pixel features that can be compressed.
To visualize the learned principal components, we can reshape the eigenvectors to the original 28×28 image shape:
fig, axes = plt.subplots(3, 5, figsize=(10, 6), subplot_kw={‘xticks‘:[], ‘yticks‘:[]})
for i, ax in enumerate(axes.flat):
ax.imshow(pca.components_[i].reshape(28, 28), cmap=‘gray‘)
These "eigendigits" represent the fundamental patterns that can be linearly combined to reconstruct the original digit images. The first component captures the overall brightness, while subsequent components represent increasingly finer details.
We can also visualize how the data is distributed in the reduced space by projecting it onto the first two principal components:
X_pca = pca.transform(X_centered)[:,:2]
plt.figure(figsize=(10, 8))
plt.scatter(X_pca[:,0], X_pca[:,1], c=y, cmap=‘Spectral‘, alpha=0.5)
plt.xlabel(‘First Principal Component‘)
plt.ylabel(‘Second Principal Component‘)
plt.colorbar(label=‘Digit‘)
While there is some separation between the digits, there is still significant overlap, especially among the ones (red) and sevens (yellow). This suggests that linear methods like PCA may not be sufficient to fully disentangle the classes.
Using PCA for Dimensionality Reduction and Feature Extraction
PCA is commonly used in machine learning pipelines for both dimensionality reduction and feature extraction prior to training a model. By reducing the dimensionality, we can often train models faster and with less risk of overfitting. The extracted features can also lead to improved model performance compared to using the raw features.
For example, we can train a simple logistic regression model on the MNIST dataset both with and without PCA preprocessing:
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
lr_raw = LogisticRegression(max_iter=500)
lr_raw.fit(X_train, y_train)
print(f"Test accuracy (raw features): {lr_raw.score(X_test, y_test):.3f}")
pca = PCA(n_components=50)
X_train_pca = pca.fit_transform(X_train)
X_test_pca = pca.transform(X_test)
lr_pca = LogisticRegression(max_iter=500)
lr_pca.fit(X_train_pca, y_train)
print(f"Test accuracy (PCA features): {lr_pca.score(X_test_pca, y_test):.3f}")Test accuracy (raw features): 0.916
Test accuracy (PCA features): 0.924We see that using the first 50 principal components as features leads to a modest improvement in test accuracy compared to using the raw pixel values. More sophisticated techniques building on PCA like sparse PCA[6] have achieved even better results.
Strengths and Limitations of PCA
PCA has several notable strengths:
- It is a non-parametric technique that makes no assumptions about the underlying distribution of the data.
- It reduces dimensionality in a way that maximally preserves the variance in the data.
- The learned components are interpretable and can reveal meaningful structure in the data.
- It is computationally efficient and can be applied to very high-dimensional datasets.
However, PCA also has some important limitations:
- It only captures linear structure and may miss important nonlinear relationships in the data.
- It is sensitive to the scaling of the features and requires the data to be centered.
- The choice of the number of components to retain is somewhat arbitrary and must be determined empirically.
- It does not consider class labels and therefore may not yield the most discriminative features for classification tasks.
For these reasons, PCA is often used as a preprocessing step rather than an end-to-end solution. More advanced nonlinear dimensionality reduction techniques like t-SNE, UMAP[7], and autoencoders can capture more complex structure but are also more computationally intensive and less interpretable.
Conclusion and Future Directions
In this post, we took a deep dive into Principal Component Analysis and its application to the MNIST dataset. We saw how PCA can uncover latent linear structure in high-dimensional data and compress it into a lower-dimensional representation that preserves most of the original variance. By visualizing the learned components and the projected data, we gained insight into the inherent patterns and variability in the MNIST digits.
Looking ahead, there are many promising directions for building on the foundation of PCA:
- Kernel PCA[8] extends PCA to capture nonlinear structure by first applying a kernel function to the data.
- Sparse PCA[6] yields components that are combinations of only a few original features, improving interpretability.
- Supervised PCA[9] incorporates label information to learn components that are most predictive of an outcome.
- Deep learning-based methods like variational autoencoders[10] learn nonlinear embeddings that can capture more complex structure than PCA.
Ultimately, the choice of dimensionality reduction technique depends on the specific characteristics of the data and the goals of the analysis. For the MNIST dataset, PCA provides a strong baseline and starting point for more advanced methods. Its simplicity and interpretability make it a valuable tool in any machine learning practitioner‘s toolkit.
