PCA on embeddings — Linear algebra for LLMs

The goal

You have $n$ embeddings in $\mathbb{R}^{4096}$ and want to see them on a screen. Principal component analysis finds the directions along which the points spread out the most, so projecting onto the top two loses as little as possible.

Write

PCA: find the directions of largest spread; project onto the top

k

.

Step 1: center

Stack the points as rows and subtract the mean of each column. These four points already have mean $[0, 0]$ . Without centering, the top direction would just point at the mean.

Write

Center:

X \leftarrow X - \bar{x}

(subtract each column’s mean).

Step 2: the directions

X^\top X =

\displaystyle \lambda = 32 \text{ on } \tfrac{[1, 1]}{\sqrt2}, \qquad \lambda = 8 \text{ on } \tfrac{[1, -1]}{\sqrt2}

The principal directions are the eigenvectors of $X^\top X$ (proportional to the covariance), which are the right singular vectors $V$ of $X$ . Their eigenvalues are $\sigma_i^2$ and measure the spread along each direction.

Write

Directions: columns of

V

in

X = U\Sigma V^\top

(eigenvectors of

X^\top X

).

Step 3: project

The new coordinates of each point are its dot products with the principal directions: $XV$ . Keep the first $k$ columns. PC1 ranges over $\pm 2.83$ , PC2 over $\pm 1.41$ .

Explained variance

\displaystyle \frac{\sigma_1^2}{\sum_i \sigma_i^2} = \frac{32}{32 + 8} = 80\%

The first direction captures 80% of the spread. For real embeddings, the top two components often capture only a small share of the variance, which is a reason to be careful about reading too much into a 2D plot.

Write

Explained variance of PC

i

:

\sigma_i^2 / \sum_j \sigma_j^2

.

✏️ Your turn

Project a point

x = [2, 0], \quad v_1 = \tfrac{[1, 1]}{\sqrt 2}, \quad v_2 = \tfrac{[1, -1]}{\sqrt 2}

Write $x$ ’s two principal coordinates.

Work it out on paper, then go to the next card to check.

Answer

\displaystyle x \cdot v_1 = \tfrac{2}{\sqrt2} = \sqrt 2 \approx 1.41, \qquad x \cdot v_2 = \sqrt2 \approx 1.41

In practice

Xc = X - X.mean(dim=0)
U, S, Vh = torch.linalg.svd(Xc, full_matrices=False)
coords = Xc @ Vh[:2].T                  # (n, 2): the 2D plot
explained = S**2 / (S**2).sum()

In a transformer

Transformer hidden states often have a few dimensions with very large values (“outlier” or “massive” activations). They can dominate the top principal components, so the plot shows those dimensions rather than meaning. Normalizing first, or dropping them, helps.

📌 Recap

PCA in one card

Center the data, take the SVD, project: coordinates $= X_c V_{:, :k}$ .
Principal directions are eigenvectors of $X^\top X$ (right singular vectors).
Explained variance: $\sigma_i^2 / \sum_j \sigma_j^2$ .
Outlier dimensions can dominate; check before you interpret.

Write

The four lines above.