Structured matrices in ML for science: circulants and FNO, spectral graphs and GNNs, equivariance, matrix manifolds

Before this lesson: lesson 02 (the incidence matrix), lesson 04 (the DFT is unitary), lesson 06 (spectral theorem, Rayleigh quotient) and lesson 08 (SVD and polar decomposition). Time: about 90 minutes, plus the notebook.

---

Why this matters #

Physical laws do not change when you shift, rotate or relabel the system. Such a symmetry is a group of matrices acting on the data, and a layer that respects it is equivariant. One fact connects the main architectures of ML for science: matrices that commute share eigenvectors. For shifts these are Fourier modes: CNNs and Fourier neural operators (FNOs). For graphs they are Laplacian eigenvectors: spectral GNNs such as GCN. Rotations give E(3)-equivariant interatomic potentials (NequIP, MACE). Orthogonality constraints and subspaces live on the Stiefel and Grassmann manifolds: PCA, orthogonal layers, Muon.

1. Predict first #

2. The picture #

The simplest case, $n = 3$. The cyclic shift $P$ maps $(x_0, x_1, x_2)$ to $(x_2, x_0, x_1)$, and $P^3 = I$. Take a "wave" $f = (1, z, z^2)$ with $z^3 = 1$. Then $Pf = (z^2, 1, z) = z^{-1}f$: shifting a wave only multiplies it by a phase. There are three such waves, for the three cube roots of unity, with three different eigenvalues $z^{-1}$.

Now take any linear filter $C$ that does not care where the signal sits: $CP = PC$. Then $P(Cf) = C(Pf) = z^{-1}(Cf)$, so $Cf$ is again an eigenvector of $P$ for the same eigenvalue. Each eigenspace of $P$ is a line, so $Cf$ must be a multiple of $f$. A shift-invariant filter can only scale each wave: it is diagonal in the Fourier basis. This is the convolution theorem.

A computation by hand, $n = 4$. Take the filter with first column $c = (2, 1, 0, 1)$: each output is twice the sample plus its two neighbors, with wrap-around. Its matrix has entry $c_{j-l}$ in row $j$, column $l$ (indices mod 4), so each row is the row above, shifted one place to the right:

For $n = 4$, $\omega = e^{2\pi i/4} = i$, and the first Fourier wave is $f_1 = (1, \omega, \omega^2, \omega^3) = (1, i, -1, -i)$. Multiply row by row:

1. Row 0: $2\cdot 1 + 1\cdot i + 0\cdot(-1) + 1\cdot(-i) = 2$.
2. Row 1: $1\cdot 1 + 2\cdot i + 1\cdot(-1) + 0\cdot(-i) = 2i$.
3. Row 2: $0\cdot 1 + 1\cdot i + 2\cdot(-1) + 1\cdot(-i) = -2$.
4. Row 3: $1\cdot 1 + 0\cdot i + 1\cdot(-1) + 2\cdot(-i) = -2i$.

So $Cf_1 = (2, 2i, -2, -2i) = 2\,f_1$: the same wave, scaled by 2. Theorem 1 below predicts the scale without any multiplication: $\hat c_1 = \sum_m c_m\omega^{-m} = 2 + 1\cdot(-i) + 0 + 1\cdot(i) = 2$. All four eigenvalues are $\hat c = \mathrm{fft}(c) = (4, 2, 0, 2)$.

Interactive figure: loading…

Graphs. Replace the cycle by any graph, and the shift by the graph Laplacian $L$. The eigenvectors of $L$ play the role of waves: small eigenvalues belong to signals that are smooth over the graph, large ones to signals that oscillate between neighbors.

Interactive figure: loading…

Equivariance is the same commuting property for any symmetry: transform the input and then apply the layer, or apply the layer and then transform the output; both give the same result.

Interactive figure: loading…

Manifolds. A unit vector in $\mathbb{R}^n$ lives on a sphere, a curved set. A frame of $p$ orthonormal vectors lives on the Stiefel manifold. If only their span matters, it lives on the Grassmann manifold. Section 3 shows how to take a gradient step on a sphere and stay on it.

3. Definitions and main results #

Circulant matrices and the DFT #

Indices run from $0$ to $n-1$ and are taken mod $n$. The cyclic shift is $(Px)_j = x_{j-1}$. The circulant with first column $c = (c_0, \dots, c_{n-1})$ is

the circular convolution of $c$ and $x$. For $n = 3$: $C = \begin{pmatrix} c_0 & c_2 & c_1 \\ c_1 & c_0 & c_2 \\ c_2 & c_1 & c_0 \end{pmatrix}$. The DFT (lesson 04) uses $\omega = e^{2\pi i/n}$ and $F_{kj} = \omega^{-jk}$, so $\hat c = Fc$ and $F^{-1} = F^*/n$. The Fourier vector $f_k$ has entries $(f_k)_j = \omega^{jk}$; it is column $k$ of $F^*$.

Theorem 1 (circulants are diagonalized by the DFT). $C(c)f_k = \hat c_k f_k$ for $k = 0, \dots, n-1$. Hence $C(c) = F^{-1}\operatorname{diag}(\hat c)F$ and $c \circledast x = F^{-1}(\hat c \odot Fx)$, which costs $O(n\log n)$ with the FFT.

Proof.

1. $(Pf_k)_j = (f_k)_{j-1} = \omega^{(j-1)k} = \omega^{-k}(f_k)_j$. The wrap-around index is fine because $\omega^{nk} = 1$. So $Pf_k = \omega^{-k}f_k$.
2. Repeating, $P^m f_k = \omega^{-mk} f_k$.
3. $C f_k = \sum_m c_m P^m f_k = \big(\sum_m c_m\omega^{-mk}\big) f_k = \hat c_k f_k$.
4. Put the $n$ equations side by side as columns: $CF^* = F^*\operatorname{diag}(\hat c)$. Since $(F^*)^{-1} = F/n$, $C = F^*\operatorname{diag}(\hat c)F/n = F^{-1}\operatorname{diag}(\hat c)F$. $\blacksquare$

Theorem 2 (shift-equivariant linear maps are circulants). $CP = PC$ if and only if $C$ is circulant.

Proof.

1. $(PC)_{jl} = C_{j-1,l}$, because $P$ moves row $j-1$ to row $j$.
2. $(CP)_{jl} = \sum_q C_{jq}P_{ql} = C_{j,l+1}$, because $P_{ql} = 1$ exactly when $q = l+1$.
3. So $CP = PC$ means $C_{j-1,l} = C_{j,l+1}$ for all $j, l$, that is, $C_{j,l} = C_{j+1,l+1}$: the entries are constant along the wrapped diagonals. Then $C_{jl}$ depends only on $j - l$: $C$ is circulant.
4. Conversely, a circulant is a polynomial in $P$, so it commutes with $P$. $\blacksquare$

Consequences:

• The eigenvalues $\hat c_k$ are the frequency response of the filter. For real $c$, $\hat c_{n-k} = \overline{\hat c_k}$; for symmetric $c$ ($c_m = c_{n-m}$) they are real.
• All circulants share one eigenbasis, so they commute with each other.
• A 2-D periodic convolution is a block-circulant matrix with circulant blocks. The 2-D DFT, $F\otimes F$ (lesson 16), diagonalizes it.
• Deep-learning "convolution" is cross-correlation (a flipped kernel, still circulant). Zero padding gives a Toeplitz matrix, not a circulant, so shift equivariance fails near the borders. Strides and pooling also break exact equivariance.

The Fourier layer of an FNO #

The FNO Fourier layer (Li et al. 2021). A function on a periodic grid with $n$ points and $d_v$ channels is $v \in \mathbb{R}^{n\times d_v}$. One layer computes $v \mapsto \sigma\big(vW^\top + \mathcal{K}v\big)$, with $W \in \mathbb{R}^{d_v\times d_v}$ acting pointwise and

where $(Fv)_k \in \mathbb{C}^{d_v}$ is mode $k$ of all channels and $R_k \in \mathbb{C}^{d_v\times d_v}$ is learned. Read it with Theorem 1:

• For one channel ($d_v = 1$), $\mathcal{K}$ is exactly a circulant with eigenvalue $R_k$ on the kept modes (and $\overline{R_k}$ on the mirror modes $n - k$, so a real input gives a real output; rfft and irfft do this for you) and $0$ on the others: a learned global convolution that is also a low-pass filter. Its kernel $F^{-1}R$ spreads over the whole domain, so one layer has a global receptive field.
• With $d_v$ channels, $\mathcal{K}$ is a $d_v\times d_v$ block matrix of circulants, one per pair of channels.
• Parameters: $k_{\max}d_v^2$ complex numbers per layer in 1-D (one $R_k$ per kept mode), independent of $n$. The paper uses 4 layers, $k_{\max} = 16$, $d_v = 64$ in 1-D. Cost: the FFTs, $O(d_v\, n\log n)$, plus $O(k_{\max}d_v^2)$ for the mixing.
• Discretization invariance: $R$ lives on frequencies, not on grid points. On any uniform grid with $n/2 \ge k_{\max}$ that resolves the input, the layer computes the same function. This is why an FNO trained at one resolution can be evaluated at another.
• Shift equivariance: $\mathcal{K}$, $W$ and $\sigma$ all commute with cyclic shifts of the grid. The reference code appends the grid coordinates as extra input channels, which breaks this on purpose.

Interactive figure: loading…

Why the grid does not matter, and when it does (aliasing)

Sample a periodic function $u$ on $n$ uniform points $x_j = j/n$. The forward DFT of the samples grows like $n$ (it is a sum of $n$ terms) and the inverse DFT divides by $n$, so the factors cancel: the layer works with the Fourier coefficients $a_k \approx \frac1n\sum_j u(x_j)e^{-2\pi ikj/n}$ of the function, not of the grid. If $u$ has no modes above $n/2$, these sums are exact, and every such grid gives the same $a_0, \dots, a_{k_{\max}-1}$ and the same output.

If $u$ has a mode $m > n/2$, the samples cannot tell it apart from the mode $m - n$ (or $n - m$): $e^{2\pi imj/n} = e^{2\pi i(m-n)j/n}$ at every grid point. This is aliasing. The high mode is added to a low coefficient, the kept coefficients change a little, and the agreement between grids becomes approximate. FNOs on non-uniform meshes need another transform (for example the geometry-aware variants), because the FFT needs a uniform grid.

The graph Laplacian #

A graph has $n$ nodes and symmetric weights $A_{ij} = A_{ji} \ge 0$ ($A_{ij} > 0$ exactly when $i$ and $j$ are joined). The degrees are $d_i = \sum_j A_{ij}$, $D = \operatorname{diag}(d)$, and the Laplacian is $L = D - A$. For an unweighted graph, $L = B^\top B$ with the incidence matrix $B$ of lesson 02. Its quadratic form measures how much a signal $x \in \mathbb{R}^n$ varies across edges:

Theorem 3 (what the eigenvalues say). Sort the eigenvalues upward, $\lambda_1 \le \dots \le \lambda_n$ (lesson 06 called them $\mu$). Then: (a) $L \succeq 0$ and $L\mathbf 1 = 0$, so $\lambda_1 = 0$; (b) the multiplicity of $0$ equals the number of connected components; (c) $\lambda_2 = \min_{x\perp\mathbf 1,\,x\ne0} x^\top Lx/x^\top x$, and $\lambda_2 > 0$ exactly when the graph is connected; (d) $\lambda_n \le 2d_{\max}$.

Proof.

1. (a) The quadratic form is a sum of non-negative terms. $(L\mathbf 1)_i = d_i - \sum_j A_{ij} = 0$.
2. (b) Since $L \succeq 0$, $Lx = 0 \iff x^\top Lx = 0 \iff x_i = x_j$ on every edge $\iff$ $x$ is constant on each component. The indicator vectors of the components are independent and span this null space.
3. (c) Courant–Fischer (lesson 06) with the first eigenvector $\mathbf 1$. By (b), $\lambda_2 = 0$ exactly when there are at least two components.
4. (d) $(x_i - x_j)^2 \le 2x_i^2 + 2x_j^2$. Summing over the edges, $x^\top Lx \le 2\sum_i d_ix_i^2 \le 2d_{\max}\|x\|^2$. $\blacksquare$

How to read the spectrum:

• The eigenvector $u_k$ has variation $u_k^\top Lu_k = \lambda_k$: small $\lambda$ = smooth over the graph, large $\lambda$ = oscillating. $\hat x = U^\top x$ is the graph Fourier transform.
• On the cycle graph, $L$ is the circulant with $c = (2, -1, 0, \dots, 0, -1)$: its eigenvalues are $2 - 2\cos(2\pi k/n)$ and its eigenvectors are Fourier modes. On a cycle, graph Fourier = DFT.
• The minimizer in (c) is the Fiedler vector. Its signs split the graph in two (spectral bisection); the first $k$ eigenvectors give spectral clustering. A small $\lambda_2$ means a bottleneck. Cheeger's inequality makes this precise for the normalized Laplacian (next item): $\lambda_2/2 \le \phi(G) \le \sqrt{2\lambda_2}$ with $\lambda_2 = \lambda_2(L_{\text{sym}})$, where the conductance $\phi(G)$ is the smallest ratio of cut weight to the degree sum of the smaller side.
• The normalized Laplacian $L_{\text{sym}} = I - D^{-1/2}AD^{-1/2}$ (all $d_i > 0$) has eigenvalues in $[0, 2]$, with $2$ exactly when a component is bipartite.

From a spectral filter to the GCN layer #

Theorem 4 (the GCN propagation rule; Kipf and Welling 2017, §2). Let $L_{\text{sym}} = U\Lambda U^\top$. The GCN layer is the end of six steps. Each step makes one approximation or one change; the "Changes" line says what it buys and what it costs.

1. Spectral filter. Filter a signal $x$ by $g(L_{\text{sym}})x = Ug(\Lambda)U^\top x$. Changes: nothing yet. This is the exact graph analogue of a convolution, but it needs the eigendecomposition ($O(n^3)$) and dense products ($O(n^2)$), and the filter is not local.
2. Polynomial filter (ChebNet). Take $g(\Lambda) = \sum_{k=0}^K\theta_kT_k(\tilde\Lambda)$ with Chebyshev polynomials $T_k$ and $\tilde\Lambda = 2\Lambda/\lambda_{\max} - I$. Because $U\Lambda^kU^\top = L_{\text{sym}}^k$, this is $\sum_k\theta_kT_k(\tilde L)x$ with $\tilde L = 2L_{\text{sym}}/\lambda_{\max} - I$. Changes: no eigenvectors are needed; only nodes at most $K$ hops apart interact; the cost is $O(K|E|)$; the filter has $K+1$ parameters.
3. First order. Take $K = 1$ and $\lambda_{\max} \approx 2$ (its upper bound). Then $\tilde L = L_{\text{sym}} - I = -D^{-1/2}AD^{-1/2}$ and $g \star x \approx \theta_0x - \theta_1D^{-1/2}AD^{-1/2}x$. Changes: each layer sees only direct neighbors; a wider view comes from stacking layers.
4. One parameter. Set $\theta = \theta_0 = -\theta_1$: $g\star x \approx \theta(I + D^{-1/2}AD^{-1/2})x = \theta(2I - L_{\text{sym}})x$. Changes: fewer parameters (less overfitting). But $2I - L_{\text{sym}}$ has eigenvalues in $[0, 2]$, so stacking layers can multiply a signal by up to $2^{\text{depth}}$: signals and gradients explode or vanish.
5. Renormalization trick. Replace $I + D^{-1/2}AD^{-1/2}$ by $\hat A = \tilde D^{-1/2}\tilde A\tilde D^{-1/2}$ with $\tilde A = A + I$ (self-loops) and $\tilde D = \operatorname{diag}(\tilde A\mathbf 1)$. Changes: the eigenvalues move into $(-1, 1]$ (proof below), so repeated layers are stable.
6. Channels. For node features $H \in \mathbb{R}^{n\times C}$ and weights $W \in \mathbb{R}^{C\times F}$: $H' = \sigma(\hat AHW)$, at cost $O(|E|CF)$. Changes: $C$ input channels and $F$ filters; the same $W$ is shared by all nodes.

Interactive figure: loading…

A 3-node example. The path 1–2–3 has degrees $d = (1, 2, 1)$.

• $L_{\text{sym}}$ has eigenvalues $0, 1, 2$ (the path is bipartite, so 2 appears).
• Step 4: $2I - L_{\text{sym}}$ has eigenvalues $2, 1, 0$. Ten layers on $x = e_1$ give $(I + D^{-1/2}AD^{-1/2})^{10}e_1 \approx (256.5, 362.0, 255.5)$: the signal exploded.
• Step 5: with self-loops, $\tilde d = (2, 3, 2)$ and $\hat A_{ij} = \tilde A_{ij}/\sqrt{\tilde d_i\tilde d_j}$:

• One layer on $e_1$: $\hat Ae_1 = (0.5, 0.408, 0)$. The feature moved one hop.
• Twenty layers: $\hat A^{20}e_1 \approx (0.286, 0.350, 0.286)$. This is the projection of $e_1$ onto the top eigenvector $\sqrt{\tilde d}/\|\sqrt{\tilde d}\| = (\sqrt2, \sqrt3, \sqrt2)/\sqrt7$, which is $(2/7, \sqrt6/7, 2/7)$. The other parts decayed like $(1/2)^{20}$ and $(1/6)^{20}$.

Why step 5 works: the eigenvalues of $\hat A$ lie in $(-1, 1]$.

1. $\tilde D - \tilde A = D - A = L$, so $I - \hat A = \tilde D^{-1/2}L\tilde D^{-1/2} \succeq 0$: every eigenvalue is $\le 1$.
2. $L\mathbf 1 = 0$ gives $\hat A\,\tilde D^{1/2}\mathbf 1 = \tilde D^{1/2}\mathbf 1$: the eigenvalue 1 is attained.
3. $I + \hat A = \tilde D^{-1/2}(D + A + 2I)\tilde D^{-1/2} \succ 0$, because $x^\top(D + A)x = \sum_{\text{edges}}A_{ij}(x_i + x_j)^2 \ge 0$: every eigenvalue is $> -1$. $\blacksquare$

Over-smoothing. Applying $\hat A$ many times is power iteration (lesson 05). On a connected graph, $\hat A^kH$ converges to the projection of $H$ onto $\tilde D^{1/2}\mathbf 1$: every node's feature becomes proportional to $\sqrt{\tilde d_i}$, and deep GCNs without residual connections lose the difference between nodes.

Permutation equivariance. Relabeling the nodes with a permutation matrix $\Pi$ maps $A \to \Pi A\Pi^\top$ and $H \to \Pi H$. Then $\tilde D \to \Pi\tilde D\Pi^\top$, $\hat A \to \Pi\hat A\Pi^\top$, and $\sigma(\Pi\hat A\Pi^\top\Pi HW) = \Pi\,\sigma(\hat AHW)$, because $\Pi^\top\Pi = I$ and $\sigma$ acts entrywise.

Equivariance #

A group $G$ is a set of transformations that is closed under composition and inverses: shifts, permutations, rotations. A representation $\rho$ gives each $g \in G$ an invertible matrix $\rho(g)$ with $\rho(gh) = \rho(g)\rho(h)$.

Definition. A map $f$ is equivariant if $f(\rho(g)x) = \rho'(g)f(x)$ for all $g \in G$ and all inputs $x$, where $\rho$ acts on inputs and $\rho'$ on outputs. It is invariant if $\rho'(g) = I$.

Data	Group	$\rho(g)$ on the input	$\rho'(g)$ on the output	Equivariant layers
periodic signal $x \in \mathbb{R}^n$	cyclic shifts	$P^s$	$P^s$	circulant convolution, FNO layer
graph $(A, H)$, $H \in \mathbb{R}^{n\times C}$	permutations $S_n$	$(\Pi A\Pi^\top, \Pi H)$	$\Pi H$ (node outputs); $I$ (graph outputs)	GCN, message passing, DeepSets
atoms $X \in \mathbb{R}^{N\times3}$ (one row per atom)	E(3): $Q \in O(3)$, $t \in \mathbb{R}^3$	$X \mapsto XQ^\top + \mathbf 1t^\top$	forces $F \mapsto FQ^\top$; energy unchanged	EGNN, NequIP, MACE
wind field on the sphere	rotations SO(3)	rotate the positions and the vectors	the same	spherical CNNs, SFNO

Theorem 5 (equivariant linear maps form a subspace). A linear map $W$ is equivariant if and only if $W\rho(g) = \rho'(g)W$ for all $g$. It suffices to check a set of generators of $G$. With vec and Kronecker products (lesson 16), the conditions read $(\rho(g)^\top\otimes I - I\otimes\rho'(g))\operatorname{vec}(W) = 0$: the equivariant layers are a null space.

Proof of Theorem 5

1. Equivariance as a matrix equation. For $f(x) = Wx$, equivariance says $W\rho(g)x = \rho'(g)Wx$ for all $x$. Two matrices that agree on every $x$ are equal: $W\rho(g) = \rho'(g)W$. The converse is the same line read backwards.
2. Products. If the equation holds for $g$ and $h$, then $W\rho(gh) = W\rho(g)\rho(h) = \rho'(g)W\rho(h) = \rho'(g)\rho'(h)W = \rho'(gh)W$.
3. Inverses. Multiply $W\rho(g) = \rho'(g)W$ by $\rho'(g)^{-1}$ on the left and by $\rho(g)^{-1}$ on the right: $\rho'(g)^{-1}W = W\rho(g)^{-1}$, that is, $W\rho(g^{-1}) = \rho'(g^{-1})W$.
4. Every element of $G$ is a product of generators and their inverses, so by steps 2 and 3 the generators suffice.
5. Vec form. The identity $\operatorname{vec}(AXB) = (B^\top\otimes A)\operatorname{vec}(X)$ (lesson 16) gives $\operatorname{vec}(W\rho) = (\rho^\top\otimes I)\operatorname{vec}(W)$ and $\operatorname{vec}(\rho'W) = (I\otimes\rho')\operatorname{vec}(W)$. So each condition is one homogeneous linear equation in $\operatorname{vec}(W)$, and the solutions of a set of such equations form a null space. $\blacksquare$

Example (DeepSets, Zaheer et al. 2017). For permutations acting on $\mathbb{R}^n$, the equivariant linear maps $\mathbb{R}^n \to \mathbb{R}^n$ are exactly $W = aI + b\mathbf 1\mathbf 1^\top$.

1. $\Pi W\Pi^\top = W$ for every permutation $\pi$ means $W_{\pi(i)\pi(j)} = W_{ij}$ for all $i, j$.
2. A permutation can move any diagonal position $(i, i)$ to any other $(k, k)$: all diagonal entries are equal.
3. It can move any off-diagonal position $(i, j)$, $i \ne j$, to any other $(k, l)$, $k \ne l$: all off-diagonal entries are equal.
4. So $W = aI + b\mathbf 1\mathbf 1^\top$; conversely, $\Pi\mathbf 1 = \mathbf 1$ shows that it commutes with every $\Pi$. $\blacksquare$

Two parameters instead of $n^2$. With $C$ input and $F$ output channels, $H \mapsto H\Lambda + \mathbf 1\mathbf 1^\top H\Gamma$ ($\Lambda, \Gamma \in \mathbb{R}^{C\times F}$): $2CF$ parameters.

An invariant energy gives equivariant forces. If $E(XQ^\top + \mathbf 1t^\top) = E(X)$ for all $Q \in O(3)$ and $t$, then the forces $F(X) = -\nabla E(X) \in \mathbb{R}^{N\times3}$ satisfy $F(XQ^\top + \mathbf 1t^\top) = F(X)Q^\top$.

1. Let $Y = XQ^\top + \mathbf 1t^\top$, so $dY = dX\,Q^\top$.
2. Chain rule: $\nabla_X[E(Y)] = \nabla E(Y)\,Q$.
3. The left side equals $\nabla E(X)$, since $E(Y) = E(X)$ for all $X$. Multiply by $Q^\top$ on the right: $\nabla E(Y) = \nabla E(X)Q^\top$. $\blacksquare$

This is why machine-learned interatomic potentials predict an invariant energy and get the forces by autodiff.

Building E(3)-equivariant layers. Distances $r_{ij} = \|x_i - x_j\|$ are invariant. Relative vectors $x_i - x_j$ rotate with $Q$ and ignore $t$. So $f_i(X) = \sum_{j\ne i}\phi(r_{ij})(x_i - x_j)$ is equivariant for any scalar function $\phi$; this is the coordinate update of EGNN (Satorras et al. 2021). Richer features (spherical harmonics of degree $\ell$, the irreducible representations of SO(3), combined with Clebsch–Gordan coefficients) are the language of e3nn, NequIP and MACE.

How to test equivariance numerically.

1. Draw random inputs $x$ and random group elements $g$: shifts; permutations; $Q \in O(3)$ (for example scipy.stats.ortho_group) and $t \in \mathbb{R}^3$.
2. Compute the relative error $e = \|f(\rho(g)x) - \rho'(g)f(x)\|/\|f(x)\|$.
3. An exactly equivariant layer gives $e$ at round-off level ($\approx10^{-15}$ in float64, $\approx10^{-6}$ in float32) in every trial. Cover the generators, and for O(3) both $\det Q = +1$ and $-1$.
4. Run the same test on a layer that is not equivariant and check that $e = O(1)$: the test must be able to fail.
5. An approximately equivariant layer (zero padding under shifts, a CNN under 45° rotations of a pixel grid) gives a small $e$ that is far above round-off. Report it; do not call it equivariant.

Matrix manifolds: Stiefel and Grassmann #

The sphere first ($p = 1$). The set of unit vectors $\{x \in \mathbb{R}^n : x^\top x = 1\}$ is a sphere. Suppose we want to minimize a function $f$ over it, for example $f(x) = -\tfrac12 x^\top Sx$ (its minimizer is the top eigenvector of $S$, lesson 06). A plain gradient step $x - \eta\nabla f(x)$ leaves the sphere. Riemannian gradient descent fixes this in three steps:

1. Tangent plane. Differentiate $x(t)^\top x(t) = 1$ at $t = 0$: $2x^\top x'(0) = 0$. So the directions in which we can move without leaving the sphere (to first order) are $T_x = \{v : x^\top v = 0\}$, the plane perpendicular to $x$.
2. Project the gradient. Remove the part of $G = \nabla f(x)$ along $x$: $P_x(G) = G - x(x^\top G)$. Then $x^\top P_x(G) = x^\top G - (x^\top x)(x^\top G) = 0$.
3. Step, then retract. The step $v = -\eta P_x(G)$ is perpendicular to $x$, so $\|x + v\|^2 = 1 + \|v\|^2 > 1$: the point $x + v$ is off the sphere. Divide by its length: $R_x(v) = (x + v)/\|x + v\|$.

Interactive figure: loading…

The Stiefel manifold. The same three steps work for $p$ orthonormal columns. The Stiefel manifold is $\operatorname{St}(n, p) = \{X \in \mathbb{R}^{n\times p}: X^\top X = I_p\}$, $n \ge p$. For $p = 1$ it is the unit sphere; for $p = n$ it is the orthogonal group $O(n)$.

• Dimension $np - p(p+1)/2$: $np$ entries, minus $p(p+1)/2$ independent equations in the symmetric matrix equation $X^\top X = I$. For $p = 1$: $n - 1$, the dimension of a sphere in $\mathbb{R}^n$.
• Tangent space. Differentiate $X(t)^\top X(t) = I$: $T_X = \{V: X^\top V + V^\top X = 0\}$, so $X^\top V$ is skew-symmetric.
• Riemannian gradient. Project the Euclidean gradient $G$ onto $T_X$: $P_X(G) = G - X\operatorname{sym}(X^\top G)$, with $\operatorname{sym}(M) = (M + M^\top)/2$. For $p = 1$, $X^\top G$ is a number and this is the formula of step 2.
• Retraction (a smooth way back to the manifold). QR retraction: the Q factor of $X + V$ (with a positive diagonal in R). Polar retraction: $UW^\top$, where $X + V = U\Sigma W^\top$ is a thin SVD; it is the closest point of $\operatorname{St}(n, p)$ to $X + V$ in the Frobenius norm. For $p = 1$ both are "divide by the length".
• One step of Riemannian gradient descent: $X \leftarrow R_X\big(-\eta\,P_X(\nabla f(X))\big)$.

For $p = n$, the polar factor $UW^\top$ is orthogonal but not always a rotation: it contains a reflection when $\det(X + V) < 0$ (see the box "orthogonal does not mean rotation" in lesson 04).

The Grassmann manifold $\operatorname{Gr}(n, p)$ is the set of $p$-dimensional subspaces of $\mathbb{R}^n$. The frames $X$ and $XQ$ ($Q \in O(p)$) span the same subspace, so $\operatorname{Gr}(n, p) = \operatorname{St}(n, p)/O(p)$.

• A unique representative of the subspace is the orthogonal projector $XX^\top$.
• Dimension $p(n - p)$: the Stiefel dimension minus $\dim O(p) = p(p-1)/2$.
• $f$ is a function on $\operatorname{Gr}(n, p)$ when $f(XQ) = f(X)$ for all $Q \in O(p)$.
• Two subspaces are compared through their principal angles (lesson 17), not through their bases.

One ML use of each:

• Stiefel: orthonormal weights. All singular values equal 1, so the layer keeps norms and is 1-Lipschitz (orthogonal RNNs, orthogonal initialization).
• Grassmann: PCA. The objective $\operatorname{tr}(X^\top SX)$ over $X \in \operatorname{St}(n, k)$ does not change under $X \to XQ$, so PCA is a problem on $\operatorname{Gr}(n, k)$: the objective finds the top-$k$ subspace. The individual principal directions inside it come from the eigenvectors of $S$, and they are unique (up to sign) only when the eigenvalues are distinct (Boumal §2.4).

Grassmann distances, LoRA subspaces and Muon

• Distances. For orthonormal bases $X, Y$ of two subspaces, the principal angles $\theta_1, \dots, \theta_p$ satisfy $\cos\theta_i = \sigma_i(X^\top Y)$. The geodesic distance on $\operatorname{Gr}(n, p)$ is $\|\theta\|_2$; the projection distance is $\|XX^\top - YY^\top\|_F = \sqrt2\,\|\sin\theta\|_2$, the $\sin\Theta$ of Davis–Kahan (lesson 17).
• Comparing learned subspaces. The LoRA paper (§7.2) compares learned update subspaces with the Grassmann-based similarity $\|U_i^\top U_j\|_F^2/\min(i, j)$; climate scientists compare the EOF subspaces of two models by principal angles.
• Muon (2024) replaces each momentum update $G = U\Sigma V^\top$ by an approximation of its polar factor $UV^\top$ (a few Newton–Schulz steps): the nearest point of the Stiefel manifold to $G$, so every direction of the update gets the same size. "Manifold Muon" (Bernstein, 2025) keeps the weights themselves on the Stiefel manifold.

Answers to section 1. (1) Yes: the Fourier vectors $f_k$. A circularly padded convolution commutes with the shift, so it is a circulant (Theorem 2), and Theorem 1 applies; the kernel only decides the eigenvalues $\hat c_k$. (2) Three, by Theorem 3(b).

4. Worked example #

(a) The 4-cycle is a circulant. Nodes 1–2–3–4–1. Its Laplacian $L_{C_4} = \begin{pmatrix} 2&-1&0&-1\\-1&2&-1&0\\0&-1&2&-1\\-1&0&-1&2\end{pmatrix}$ has first column $c = (2, -1, 0, -1)$. With $\omega = i$: $\hat c_k = 2 - \omega^{-k} - \omega^{-3k} = 2 - 2\cos(\pi k/2)$, so for $k = 0, 1, 2, 3$ the eigenvalues are $0, 2, 4, 2$.

• $\lambda = 0$ with $f_0 = (1,1,1,1)$: one component.
• $\lambda = 4 = 2d_{\max}$ with $f_2 = (1,-1,1,-1)$: the bound of Theorem 3(d) is reached (the cycle is bipartite).
• $\lambda = 2$ twice, with $f_{1,3} = (1, \pm i, -1, \mp i)$. Real versions: $(1,0,-1,0)$ and $(0,1,0,-1)$, the slowest waves around the cycle. Any combination is a Fiedler vector: it is not unique.

(b) Cut one edge. Remove the edge 4–1 to get the path 1–2–3–4. It is no longer a circulant. Its eigenvalues are $2 - 2\cos(\pi k/4)$: $0$, $2 - \sqrt2 \approx 0.586$, $2$, $2 + \sqrt2 \approx 3.414$. The Fiedler vector $\approx (0.653, 0.271, -0.271, -0.653)$ splits $\{1, 2\} \mid \{3, 4\}$, the cut in the middle. $\lambda_2$ drops from 2 to 0.586: the path is easier to cut. Since $L_{C_4} = L_{\text{path}} + bb^\top$ with $b = e_1 - e_4$, we have $L_{\text{path}} \preceq L_{C_4}$, so no eigenvalue can go up (lesson 07): $0 \le 0$, $0.586 \le 2$, $2 \le 2$, $3.414 \le 4$. ✓

(c) One GCN step on the path. With self-loops, $\tilde d = (2, 3, 3, 2)$ and $\hat A_{ij} = \tilde A_{ij}/\sqrt{\tilde d_i\tilde d_j}$:

• A feature on node 1 only, $x = e_1$: $\hat Ax = (0.5, 0.408, 0, 0)$. It moved one hop.
• Eigenvalues: $1$, $\tfrac14 + \tfrac{\sqrt{33}}{12} \approx 0.729$, $\tfrac16$, $\tfrac14 - \tfrac{\sqrt{33}}{12} \approx -0.229$, all in $(-1, 1]$. Without the trick, $I + D^{-1/2}AD^{-1/2}$ has eigenvalues $0, 0.5, 1.5, 2$.
• Over-smoothing: $\hat A^{20}e_1 \approx (0.201, 0.245, 0.245, 0.199)$. This is almost the projection of $e_1$ onto $\sqrt{\tilde d} = (\sqrt2, \sqrt3, \sqrt3, \sqrt2)$, which is $(0.2, 0.245, 0.245, 0.2)$. The rest decays like $0.729^{20} \approx 0.002$. After 20 layers the network cannot tell where the feature started.

(d) Stiefel and Grassmann. $X = [e_1\; e_2]$ and $Y = [e_1\; (e_2 + e_3)/\sqrt2]$ in $\mathbb{R}^3$: $X^\top Y = \operatorname{diag}(1, 1/\sqrt2)$, so the principal angles are $0°$ and $45°$, and $\|XX^\top - YY^\top\|_F = \sqrt2\sin45° = 1$. $\operatorname{St}(3, 2)$ has dimension $6 - 3 = 3$ and $\operatorname{Gr}(3, 2)$ has dimension $2\cdot1 = 2$ (a plane through 0 in $\mathbb{R}^3$ is fixed by its normal direction).

5. In code #

NumPy and SciPy check the worked example and the main theorems. The notebook for this lesson has these cells and the exercises.

import numpy as np
from scipy.linalg import circulant, subspace_angles
from scipy.sparse.csgraph import laplacian
from scipy.stats import ortho_group

rng = np.random.default_rng(0)

# 1) A circulant matrix is diagonalized by the DFT; its eigenvalues are fft(c).
c = np.array([2.0, -1.0, 0.0, -1.0])            # first column: Laplacian of the 4-cycle
C = circulant(c)
F = np.fft.fft(np.eye(4))                        # DFT matrix, F[k, j] = exp(-2 pi i jk / 4)
print(np.fft.fft(c).real)                        # [0. 2. 4. 2.]
print(np.allclose(C, np.linalg.inv(F) @ np.diag(np.fft.fft(c)) @ F))   # True
x = rng.standard_normal(4)
print(np.allclose(C @ x, np.fft.ifft(np.fft.fft(c) * np.fft.fft(x)).real))  # True
P = circulant([0.0, 1.0, 0.0, 0.0])              # cyclic shift, (Px)_j = x_{j-1}
print(np.allclose(C @ P, P @ C))                 # True

# 2) Laplacians of the 4-cycle and of the path (the cycle without edge 4-1).
A_cyc = circulant([0.0, 1.0, 0.0, 1.0])
A_path = A_cyc.copy()
A_path[0, 3] = A_path[3, 0] = 0.0
print(np.linalg.eigvalsh(laplacian(A_cyc)).round(3) + 0.0)               # [0. 2. 2. 4.]
print(np.linalg.eigvalsh(laplacian(A_cyc, normed=True)).round(3) + 0.0)  # [0. 1. 1. 2.]
lam, U = np.linalg.eigh(laplacian(A_path))
print(lam.round(3) + 0.0, U[:, 1].round(3))      # [0. 0.586 2. 3.414], Fiedler (up to sign)
Pi = np.eye(4)[[2, 0, 3, 1]]                     # relabel the nodes
print(np.allclose(laplacian(Pi @ A_path @ Pi.T), Pi @ laplacian(A_path) @ Pi.T))  # True

# 3) Equivariance test under E(3): X (N atoms x 3) -> X Q^T + t.
def rel_err(f, X, act_in, act_out):
    return np.linalg.norm(f(act_in(X)) - act_out(f(X))) / np.linalg.norm(f(X))

def vector_layer(X):                             # f_i = sum_j exp(-r_ij^2) (x_i - x_j)
    D = X[:, None, :] - X[None, :, :]
    r2 = (D**2).sum(axis=-1, keepdims=True)
    return (np.exp(-r2) * D).sum(axis=1)

W1, W2 = rng.standard_normal((3, 16)), rng.standard_normal((16, 3))
mlp = lambda X: np.tanh(X @ W1) @ W2             # a plain MLP on raw coordinates
X = rng.standard_normal((10, 3))
for _ in range(3):
    Q = ortho_group.rvs(3, random_state=rng)     # random orthogonal: det +1 or -1
    t = rng.standard_normal(3)
    print(round(np.linalg.det(Q)),
          rel_err(vector_layer, X, lambda Y: Y @ Q.T + t, lambda Y: Y @ Q.T),  # ~5e-16
          rel_err(mlp, X, lambda Y: Y @ Q.T, lambda Y: Y @ Q.T))               # ~1: fails

# 4) Shift test of a 1-D convolution: circular padding is exact, zero padding is not.
k = rng.standard_normal(5)                       # kernel taps for offsets -2..2
u = rng.standard_normal(64)
c64 = np.zeros(64); c64[:3] = k[2:]; c64[-2:] = k[:2]   # the same kernel as a circulant
conv_circular = lambda x: np.fft.ifft(np.fft.fft(c64) * np.fft.fft(x)).real
conv_zeros = lambda x: np.convolve(x, k, mode="same")   # zero padding at both ends
for name, conv in [("circular", conv_circular), ("zeros", conv_zeros)]:
    e = np.linalg.norm(conv(np.roll(u, 3)) - np.roll(conv(u), 3)) / np.linalg.norm(conv(u))
    print(name, e)                               # circular ~2e-16, zeros ~0.16

# 5) Stiefel St(6, 2): tangent projection and polar retraction; Grassmann angles.
X = np.linalg.qr(rng.standard_normal((6, 2)))[0]
G = rng.standard_normal((6, 2))                  # a Euclidean gradient
xi = G - X @ (X.T @ G + G.T @ X) / 2             # projection onto the tangent space
print(np.abs(X.T @ xi + xi.T @ X).max())         # ~3e-16: xi is tangent
Ur, _, Vrt = np.linalg.svd(X - 0.1 * xi, full_matrices=False)
Y = Ur @ Vrt                                     # polar retraction of a gradient step
print(np.abs(Y.T @ Y - np.eye(2)).max())         # ~4e-16: back on St(6, 2)
Xg = np.eye(3)[:, :2]
Yg = np.column_stack([[1, 0, 0], np.array([0, 1, 1]) / np.sqrt(2)])
print(np.degrees(subspace_angles(Xg, Yg)))       # [45.  0.]
print(np.linalg.norm(Xg @ Xg.T - Yg @ Yg.T))     # 1.0 = sqrt(2) * sin(45 deg)

The equivariance test in part 3 passes for $\det Q = -1$ too, because distances and relative vectors do not care about orientation. In part 4 the zero-padded layer misses by about 16%: the error comes only from the two borders, where zero padding replaces the wrapped values.

6. Where it shows up #

• Weather and climate. FourCastNet (Pathak et al. 2022) mixes tokens with learned Fourier-domain filters (AFNO). Spherical FNOs (Bonev et al. 2023) replace the DFT by the spherical harmonic transform: spherical harmonics are the eigenfunctions of the Laplacian on the sphere. GraphCast (Lam et al. 2023) is a message-passing GNN on an icosahedral mesh.
• PDE surrogates. FNOs for Navier–Stokes and Darcy flow; MeshGraphNets on simulation meshes.
• Sensor and station networks. Spectral clustering with the first eigenvectors of $L$; smoothness penalties $x^\top Lx$; graph diffusion $e^{-tL}x$ (lesson 16).
• Molecules and materials. NequIP and MACE predict an invariant energy and take forces as $-\nabla E$.

7. Common mistakes #

• Mixing DFT conventions. In NumPy's convention the eigenvalues are fft of the first column. The first row gives $\hat c_{n-k}$ instead (conjugates for real $c$).
• Calling a zero-padded or strided CNN exactly shift equivariant. Zero padding gives a Toeplitz matrix, and the error at the borders is real. Strides and pooling break exact equivariance too.
• Thinking an FNO works on any grid. It needs a uniform grid with at least the $k_{\max}$ modes, and fine enough for the input; on a coarser grid the input aliases.
• Mixing Laplacians. $D - A$ has eigenvalues up to $2d_{\max}$; $L_{\text{sym}}$ has them in $[0, 2]$.
• Expecting a unique Fiedler vector. With a repeated $\lambda_2$ (the 4-cycle) every vector in the eigenspace is a minimizer.
• Weak equivariance tests: one input only, only $\det Q = +1$, no failing baseline.
• Taking a plain gradient step on a manifold. $X - \eta\nabla f$ leaves $\operatorname{St}(n, p)$: project the gradient onto the tangent space and retract.
• Comparing subspaces through their bases. $X$ and $XQ$ are different Stiefel points but the same subspace: compare $XX^\top$ or the principal angles.

8. Check yourself #

1. Give the eigenvalues and the determinant of the circulant with first column $c = (1, 2, 0)$.
   Answer

   $C = \begin{pmatrix}1&0&2\\2&1&0\\0&2&1\end{pmatrix}$. $\hat c_k = 1 + 2\omega^{-k}$ with $\omega = e^{2\pi i/3}$: $3$, $1 + 2(-\tfrac12 - \tfrac{\sqrt3}{2}i) = -i\sqrt3$, and $i\sqrt3$. $\det C = 3\cdot(-i\sqrt3)(i\sqrt3) = 9$.
2. A 1-D FNO layer has $k_{\max} = 16$ and $d_v = 64$. How many parameters are in $R$? Does the number change between a grid of 1024 and 8192 points?
   Answer

   $16\cdot64^2 = 65{,}536$ complex numbers (131,072 real), plus $64^2$ for $W$. No: the parameters belong to Fourier modes, not to grid points.
3. An unweighted graph on 5 nodes has Laplacian eigenvalues $0, 0, 1, 2, 3$. How many components? How many edges? What can you say about the largest degree?
   Answer

   Two zero eigenvalues: two components, so it is not connected ($\lambda_2 = 0$). The eigenvalues sum to $\operatorname{tr}L = \sum_i d_i = 2|E| = 6$: three edges. $3 = \lambda_n \le 2d_{\max}$ gives $d_{\max} \ge 1.5$, so some node has degree at least 2. (Example: a path on 3 nodes, eigenvalues $0, 1, 3$, plus a separate edge, eigenvalues $0, 2$.)
4. How many parameters does the most general permutation-equivariant linear layer $\mathbb{R}^n \to \mathbb{R}^n$ have? Why?
   Answer

   Two: $W = aI + b\mathbf 1\mathbf 1^\top$. $\Pi W\Pi^\top = W$ for all permutations forces all diagonal entries to be equal and all off-diagonal entries to be equal. With $C$ input and $F$ output channels: $2CF$.
5. Your force model passes the rotation test ($e \approx 10^{-7}$ in float32) but gives $e \approx 0.5$ for $Q$ with $\det Q = -1$. What does this mean?
   Answer

   It is SO(3)- but not O(3)-equivariant: it uses orientation (for example a cross product or chirality-aware features). That is correct if mirror images may behave differently (chiral molecules), and a bug if the target is mirror symmetric.
6. Give $\dim\operatorname{St}(1000, 10)$ and $\dim\operatorname{Gr}(1000, 10)$. Why is $\operatorname{tr}(X^\top SX)$ a function on the Grassmann manifold?
   Answer

   $10{,}000 - 55 = 9945$ and $10\cdot990 = 9900$. $\operatorname{tr}((XQ)^\top S(XQ)) = \operatorname{tr}(Q^\top X^\top SXQ) = \operatorname{tr}(X^\top SX)$ for $Q \in O(10)$ (cyclic trace): it depends only on the subspace.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Give the eigenvalues and eigenvectors of the circulant with first column $c = (2, 1, 0, 1)$.

   Show answer

2. Which linear maps commute with the cyclic shift $P$?

   Show answer

3. What does one FNO Fourier layer compute (one channel)?

   Show answer

4. Define the graph Laplacian and its quadratic form.

   Show answer

5. What does the number of zero eigenvalues of $L = D - A$ tell you, and why?

   Show answer

6. What is the Fiedler vector, and what is it used for?

   Show answer

7. Write the GCN propagation rule.

   Show answer

8. Define an equivariant map, and an invariant one.

   Show answer

9. What is the Stiefel manifold $\operatorname{St}(n, p)$, and what is its dimension?

   Show answer

10. How is the Grassmann manifold built from the Stiefel manifold, and what is its dimension?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why is convolution diagonal in Fourier space?" — It commutes with the shift; the shift has the Fourier vectors as eigenvectors with distinct eigenvalues; commuting matrices share them.
• "Derive the GCN layer." — First-order Chebyshev filter, $\lambda_{\max} \approx 2$, one shared parameter, self-loops and symmetric normalization (renormalization trick): $H' = \sigma(\hat AHW)$.
• "What is over-smoothing, in linear-algebra terms?" — $\hat A^k$ is power iteration; features converge to the top eigenvector $\tilde D^{1/2}\mathbf 1$. Fixes: residual connections, normalization, fewer layers.
• "Why can an FNO trained on 64 points run on 256?" — Its weights belong to Fourier modes; every grid that resolves those modes gives the same operator.
• "How do you take a gradient step with an orthogonality constraint?" — Project the gradient onto $\{V: X^\top V \text{ skew}\}$, step, and retract with QR (positive diagonal in R) or the polar factor.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Write circular convolution with the FFT and the eigenvalues of a circulant, and check $C = F^{-1}\operatorname{diag}(\hat c)F$ against scipy.linalg.circulant.
2. Write a one-channel FNO Fourier layer with np.fft.rfft, and show that it commutes with shifts and gives the same function on grids of 64 and 256 points.
3. Build the Laplacian of a graph, count its components from the spectrum, and split a graph with two planted communities by the signs of the Fiedler vector.
4. Build the GCN matrix $\hat A$, check that its eigenvalues lie in $(-1, 1]$, test permutation equivariance, and watch over-smoothing.
5. Take Riemannian gradient steps on the Stiefel manifold (tangent projection + QR retraction) to find the top-2 eigenvector subspace of a symmetric matrix, and measure the principal angles to the true subspace.

11. Watch and read #

31. Eigenvectors of Circulant Matrices: Fourier MatrixGilbert Strang · MIT 18.065, lecture 31 · 53 min

Strang derives Theorem 1: the Fourier vectors are the eigenvectors of every circulant.

Fourier Neural Operator (FNO) [Physics Informed Machine Learning]Steve Brunton · Physics Informed Machine Learning · 18 min

The FNO layer as FFT, truncation to the low modes, a learned multiplication and an inverse FFT.

Daniel Spielman “Miracles of Algebraic Graph Theory”Daniel Spielman · Joint Mathematics Meetings · 53 min

Laplacian eigenvectors as graph drawings, and what the eigenvalues say about cuts (Theorem 3).

ICLR 2021 Keynote - "Geometric Deep Learning: The Erlangen Programme of ML" - M BronsteinMichael Bronstein · ICLR 2021 keynote · 38 min

Symmetry groups, equivariance and why CNNs, GNNs and transformers fit one framework.

• Gilbert Strang, MIT 18.065 (2018), lecture 32 "ImageNet is a convolutional neural network (CNN), the convolution rule" and lecture 35 "Finding clusters in graphs": course page with all videos.
• Bronstein, Bruna, Cohen, Veličković, Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges (2021): §3.1 (symmetries, representations, invariance and equivariance), §4.1 (graphs and sets, permutation equivariance), §4.2 (grids: circulant matrices and the Fourier transform), §5.3 (GNNs).
• Li et al., "Fourier Neural Operator for Parametric Partial Differential Equations" (ICLR 2021): §3 (Definitions 1–2, the neural operator), §4 (Definition 3 and Eq. (4)–(5), the Fourier layer; "Invariance to discretization").
• Kipf and Welling, "Semi-Supervised Classification with Graph Convolutional Networks" (ICLR 2017): §2.1–2.2, Eq. (2)–(8), the derivation of Theorem 4.
• Nicolas Boumal, An Introduction to Optimization on Smooth Manifolds (2023, free PDF): §2.4 (PCA on Stiefel and Grassmann), §7.3 (the Stiefel manifold, Eq. (7.13)–(7.28)), §9.16 (Grassmann summary).
• Sanchez-Lengeling et al., "A Gentle Introduction to Graph Neural Networks" (Distill, 2021): message passing with interactive figures.