Learning Spectral Methods by Transformers

TL;DR

The work asks whether pre-trained Transformers can perform unsupervised learning tasks by learning spectral algorithms. It provides constructive proofs showing that multi-layer Transformers can simulate the Power Method to extract left singular vectors and principal components, enabling PCA, and designs spectral algorithms for clustering a Gaussian mixture model within the Transformer framework. Theoretical results establish ERM-based guarantees and explicit Transformer constructions with depth and width bounds, alongside an auxiliary design matrix to seed iterative steps. Empirically, the authors validate eigenvalue/eigenvector prediction and GMM clustering on synthetic data and real datasets (e.g., MNIST/FMNIST), demonstrating that pre-trained Transformers can autonomously perform spectral estimation and clustering tasks. The findings highlight a principled pathway for embedding classical unsupervised spectral methods into Transformer pretraining, with implications for unsupervised learning in large-scale models and beyond.

Abstract

Transformers demonstrate significant advantages as the building block of modern LLMs. In this work, we study the capacities of Transformers in performing unsupervised learning. We show that multi-layered Transformers, given a sufficiently large set of pre-training instances, are able to learn the algorithms themselves and perform statistical estimation tasks given new instances. This learning paradigm is distinct from the in-context learning setup and is similar to the learning procedure of human brains where skills are learned through past experience. Theoretically, we prove that pre-trained Transformers can learn the spectral methods and use the classification of bi-class Gaussian mixture model as an example. Our proof is constructive using algorithmic design techniques. Our results are built upon the similarities of multi-layered Transformer architecture with the iterative recovery algorithms used in practice. Empirically, we verify the strong capacity of the multi-layered (pre-trained) Transformer on unsupervised learning through the lens of both the PCA and the Clustering tasks performed on the synthetic and real-world datasets.

Figures (12)

• Figure 1: The Constructive Proof for the Approximation of PCA. The above diagram illustrates the construction of our Transformer model in the existence proof of theorem \ref{['thm3.1']}. There are three important sub-networks in the design: (1) The Symmetrization sub-network symmetrizes $\boldsymbol{X}$ and stamp $\boldsymbol{X}\boldsymbol{X}^\top$ in the output, which corresponds to the first step in the Power Method. (2) The Power Iterations sub-network performs in a total of $\tau$ iterations for each of the principal eigenvectors, corresponds to the iterative update step in the Power Method. (3) The Removal of principal Eigenvectors subnetwork performs the estimate of $\widehat{\lambda}_{\ell}$ and the update of matrix $\boldsymbol{A}_{\ell}$ in the Power Method. Finally, we apply $\tilde{\boldsymbol{W}}_0$ and $\tilde{\boldsymbol{W}}_1$ to adjust the dimension of the output. The different colors in the diagram correspond to the different types of layers: (1) Yellow Blocks denote the Attention layer with 2 heads. (2) Orange Blocks denote the multihead Tansformers with larger $M\gg 2$. (3) Pink Blocks denote the FC layer.
• Figure 2: Eigenvalue Prediction on Synthetic Data.(1) Left: Evidence of Transformer's Ability to Predict Multiple Eigenvalues. We use a small Transformer (layer $=3$, head $=2$, embedding $= 64$) to predict top $10$ eigenvalues with $d=20$ and $N=50$. (2) Middle: Predictions of Eigenvalues with Different Input Dimension $d$. We use a small Transformer and use $N=10$ in this experiment. (3) Right: Predictions of Eigenvalues with Different Number of Layers. We use the same input as the previous multiple eigenvalues predictions experiment, and use a small Transformer to predict top-$3$ eigenvalues.
• Figure 3: Eigenvector Prediction on Synthetic Data.(1) Left: Prediction of top-$1$ eigenvector with different input dimension $d$. We use a small Transformer and $N=10$. (2) Middle: Prediction of top-$1$ eigenvector with varying number of layers. We start from a small Transformer and use $N=10$ and $d=10$. (3) Right: Predictions of eigenvectors with different numbers of $k$. We use $N=10$ and $d=10$ in this experiment. The decrease in performance when the number of eigenvectors $k$ gets larger might be due to the number of layers being small.
• Figure 4: Eigenvalues Prediction on Real World Data.(1) Left: Predicting Top-$10$ Eigenvalues on the MNIST Dataset.(2) Right: Predicting Top-$10$ Eigenvalues on the FMNIST Dataset. Our results indicate that the pre-trained Transformer model can automatically extract principal components from the complicated, high-dimensional images.
• Figure 5: Principal Space Prediction. We train a large Transformer to predict multiple principal component, but replace cosine similarity loss with eigenspace distance loss. (1) Left: Eigenspace Prediction on the Synthetic Data.(2) Right: Eigenspace Prediction on the MNIST dataset. Each point in the two figures is obtained by training the large Transformer for $100k$ steps and evaluated on $2048$ testing data, and average over 10 runs. We note that the performance degradation at $k=1,2$ in the MNIST dataset might be caused by the spectral gap for the first few principal components being too large.

Theorems & Definitions (26)

• Definition 1: Attention Layer
• Remark 1
• Definition 2: FC Layer
• Definition 3: Transformer
• Remark 2
• Theorem 2.1: Transformer Approximation of the Power Method
• Remark 3
• Lemma 2.1
• Remark 4
• Proposition 1