Global Tangent Spaces

In GeometricOptimizers standard neural network optimizers are generalized to homogeneous spaces by leveraging the special structure of the tangent spaces of this class of manifolds. When we introduced homogeneous spaces we already talked about that every tangent space to a homogeneous space $T_Y\mathcal{M}$ is of the form:

\[ T_Y\mathcal{M} = \mathfrak{g} \cdot Y := \{AY: A\in{}\mathfrak{g}\}.\]

We then have a decomposition of $\mathfrak{g}$ into a vertical part $\mathfrak{g}^{\mathrm{ver}, Y}$ and a horizontal part $\mathfrak{g}^{\mathrm{hor}, Y}$ and the horizontal part is isomorphic to $T_Y\mathcal{M}$ via:

\[ \mathfrak{g}^{\mathrm{hor}, Y} = \{\Omega(\Delta): \Delta\in{}T_Y\mathcal{M} \}.\]

We now identify a special element $E \in \mathcal{M}$ and designate the horizontal component $\mathfrak{g}^{\mathrm{hor}, E}$ as our global tangent space. We will refer to this global tangent space by $\mathfrak{g}^\mathrm{hor}$. We can now find a transformation from any $\mathfrak{g}^{\mathrm{hor}, Y}$ to $\mathfrak{g}^\mathrm{hor}$ and vice-versa (these spaces are isomorphic).

We should note that we have written all Lie group and Lie algebra actions as simple matrix multiplications, like $AE = Y$. For some Lie groups and Lie algebras, as the Lie group of isomorphisms on some domain $\mathcal{D}$, this notation may not be appropriate [7]. These Lie groups are however not relevant for what we use in GeometricOptimizers and we will stick to regular matrix notation.

Global Sections

Note that the theorem above requires us to find an element $A\in{}G$ such that $AE = Y$. We will call such a mapping $\lambda:\mathcal{M}\to{}G$ a global section^[1].

Note that in general global sections are not unique because the rank of $G$ is in general greater than that of $\mathcal{M}$. We give an example of how to construct such a global section for the Stiefel and the Grassmann manifolds below.

The Global Tangent Space for the Stiefel Manifold

We now discuss the specific form of the global tangent space for the Stiefel manifold. We pick as distinct element $E$ (which build by calling StiefelProjection):

\[E = \begin{bmatrix} \mathbb{I}_n \\ \mathbb{O} \end{bmatrix}\in{}St(n, N).\]

Based on this, elements of the vector space $\mathfrak{g}^{\mathrm{hor}, E} =: \mathfrak{g}^{\mathrm{hor}}$ are:

\[\bar{B} = \begin{pmatrix} A & B^T \\ B & \mathbb{O} \end{pmatrix},\]

where $A$ is a skew-symmetric matrix of size $n\times{}n$ and $B$ is an arbitrary matrix of size $(N - n)\times{}n$. Arrays of type $\mathfrak{g}^{\mathrm{hor}, E} \equiv \mathfrak{g}^\mathrm{hor}$ are implemented in GeometricOptimizers under the name StiefelLieAlgHorMatrix.

We can call this with e.g. a skew-symmetric matrix $A$ and an arbitrary matrix $B$:

N, n = 5, 2

A = rand(SkewSymMatrix, n)

2×2 SkewSymMatrix{Float64, Vector{Float64}}:
 0.0       -0.585542
 0.585542   0.0

B = rand(N - n, n)

3×2 Matrix{Float64}:
 0.133535  0.986067
 0.996111  0.471151
 0.814833  0.652578

The constructor is then called as follows:

B̄ = StiefelLieAlgHorMatrix(A, B, N, n)

5×5 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}:
 0.0       -0.585542  -0.133535  -0.996111  -0.814833
 0.585542   0.0       -0.986067  -0.471151  -0.652578
 0.133535   0.986067   0.0        0.0        0.0
 0.996111   0.471151   0.0        0.0        0.0
 0.814833   0.652578   0.0        0.0        0.0

We can also call it with a matrix of shape $N\times{}N$:

B̄₂ = Matrix(B̄) # note that this does not have any special structure

StiefelLieAlgHorMatrix(B̄₂, n)

5×5 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, SubArray{Float64, 2, Matrix{Float64}, Tuple{UnitRange{Int64}, UnitRange{Int64}}, false}}:
 0.0       -0.585542  -0.133535  -0.996111  -0.814833
 0.585542   0.0       -0.986067  -0.471151  -0.652578
 0.133535   0.986067   0.0        0.0        0.0
 0.996111   0.471151   0.0        0.0        0.0
 0.814833   0.652578   0.0        0.0        0.0

Or we can call it on $T_E\mathcal{M}\subset\mathbb{R}^{N\times{}n},$ i.e. a matrix of shape $N\times{}n$:

E = StiefelProjection(N, n)

5×2 StiefelProjection{Float64, Matrix{Float64}}:
 1.0  0.0
 0.0  1.0
 0.0  0.0
 0.0  0.0
 0.0  0.0

B̄₃ = B̄ * E

StiefelLieAlgHorMatrix(B̄₃, n)

5×5 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, SubArray{Float64, 2, Matrix{Float64}, Tuple{UnitRange{Int64}, UnitRange{Int64}}, false}}:
 0.0       -0.585542  -0.133535  -0.996111  -0.814833
 0.585542   0.0       -0.986067  -0.471151  -0.652578
 0.133535   0.986067   0.0        0.0        0.0
 0.996111   0.471151   0.0        0.0        0.0
 0.814833   0.652578   0.0        0.0        0.0

We now demonstrate how to map from an element of $\mathfrak{g}^{\mathrm{hor}, Y}$ to an element of $\mathfrak{g}^\mathrm{hor}$:

using GeometricOptimizers: Ω

Y = rand(StiefelManifold, N, n)
Δ = rgrad(Y, rand(N, n))
ΩΔ = Ω(Y, Δ)
λY = GlobalSection(Y)

λY_mat = Matrix(λY)

round.(λY_mat' * ΩΔ * λY_mat; digits = 3)

5×5 Matrix{Float64}:
 -0.0    -1.36    0.942   0.269  -0.371
  1.36   -0.0     0.196   0.613   0.383
 -0.942  -0.196  -0.0    -0.0    -0.0
 -0.269  -0.613   0.0     0.0    -0.0
  0.371  -0.383   0.0     0.0    -0.0

Performing this computation directly is computationally very inefficient however and the user is strongly discouraged to call Matrix on an instance of GlobalSection. The better option is calling global_rep:

_round(global_rep(λY, Δ); digits = 3)

5×5 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}:
  0.0    -1.36   0.942  0.269  -0.371
  1.36    0.0    0.196  0.613   0.383
 -0.942  -0.196  0.0    0.0     0.0
 -0.269  -0.613  0.0    0.0     0.0
  0.371  -0.383  0.0    0.0     0.0

Internally GlobalSection calls the function GeometricOptimizers.global_section which does the following for the Stiefel manifold:

A = randn(N, N - n) # or the gpu equivalent
A = A - Y * (Y' * A)
Y⟂ = qr(A).Q[1:N, 1:(N - n)]

So we draw $(N - n)$ new columns randomly, subtract the part that is spanned by the columns of $Y$ and then perform a $QR$ composition on the resulting matrix. The $Q$ part of the decomposition is a matrix of $(N - n)$ columns that is orthogonal to $Y$ and is typically referred to as $Y_\perp$ [6, 8, 9]. We can easily check that this $Y_\perp$ is indeed orthogonal to $Y$.

The function global_rep furthermore makes use of the following:

\[ \mathtt{global\_rep}(Y) = \lambda(Y)^T\Omega(Y,\Delta)\lambda(Y) = EY^T\Delta{}E^T + \begin{bmatrix} \mathbb{O} \\ \bar{\lambda}^T\Delta{}E^T \end{bmatrix} - \begin{bmatrix} \mathbb{O} & E\Delta^T\bar{\lambda} \end{bmatrix},\]

where $\lambda(Y) = [Y, \bar{\lambda}].$

This expression of global_rep means we only need $Y^T\Delta$ and $\bar{\lambda}^T\Delta$ and this is what is used internally.

We now discuss the global tangent space for the Grassmann manifold. This is similar to the Stiefel case.

Global Tangent Space for the Grassmann Manifold

In the case of the Grassmann manifold we construct the global tangent space with respect to the distinct element $\mathcal{E}=\mathrm{span}(E)\in{}Gr(n,N)$, where $E$ is again the same matrix.

The tangent tangent space $T_\mathcal{E}Gr(n,N)$ can be represented through matrices:

\[\begin{pmatrix} 0 & \cdots & 0 \\ \cdots & \cdots & \cdots \\ 0 & \cdots & 0 \\ b_{11} & \cdots & b_{1n} \\ \cdots & \cdots & \cdots \\ b_{(N-n)1} & \cdots & b_{(N-n)n} \end{pmatrix}.\]

This representation is based on the identification $T_\mathcal{E}Gr(n,N)\to{}T_E\mathcal{S}_E$ that was discussed in the section on the Grassmann manifold^[2]. We use the following notation:

\[\mathfrak{g}^\mathrm{hor} = \mathfrak{g}^{\mathrm{hor},\mathcal{E}} = \left\{\begin{pmatrix} 0 & -B^T \\ B & 0 \end{pmatrix}: \text{$B\in\mathbb{R}^{(N-n)\times{}n}$ is arbitrary}\right\}.\]

This is equivalent to the horizontal component of $\mathfrak{g}$ for the Stiefel manifold for the case when $A$ is zero. This is a reflection of the rotational invariance of the Grassmann manifold: the skew-symmetric matrices $A$ are connected to the group of rotations $O(n)$ which is factored out in the Grassmann manifold $Gr(n,N)\simeq{}St(n,N)/O(n)$. In GeometricOptimizers we thus treat the Grassmann manifold as being embedded in the Stiefel manifold. In [9] viewing the Grassmann manifold as a quotient space of the Stiefel manifold is important for "feasibility" in "practical computations".

Library functions

AbstractLieAlgHorMatrix, StiefelLieAlgHorMatrix, GrassmannLieAlgHorMatrix, StiefelProjection, GlobalSection, global_section, global_rep and apply_section. Their docstrings are on the reference page, where every docstring in the package is rendered once; the names above link to them.

GeometricMachineLearning re-exports all of these, so a neural network built there is optimized with exactly the machinery described above.

References