Machine Learning on Stiefel Manifold

The Stiefel manifold is a useful geometric object in machine learning because there are applications where the target transformation must preserve the norm of input vectors.

For example, in recurrent neural networks (RNNs), the transformation formula is defined as:

\[h_t = \sigma(W_{rec}h_{t-1} + W_{in}x_{t-1} + b)\]

Here, \(W_{rec}\) is a linear transformation acting on the previous system state, \(W_{in}\) is a linear transformation for the input data from the current time step, \(b\) is a bias, and \(\sigma\) is a non-linearity. This is a very powerful idea that gives a model a notion of memory through its previous states. However, it also introduces a significant problem: vanishing and exploding gradients. Let's examine how the current state depends on previous ones, assuming an identity transformation for the activation function (\(\sigma = \text{id}\)):

\[\frac{\partial h_t}{\partial h_{t-1}} = \frac{\partial}{\partial h_{t-1}}(W_{rec}h_{t-1} + W_{in}x_{t-1} + b) = W_{rec}\]

Using the chain rule, let's find the expression for step \(t-2\):

\[\frac{\partial h_t}{\partial h_{t-2}} = \frac{\partial h_t}{\partial h_{t-1}} \frac{\partial h_{t-1}}{\partial h_{t-2}} = W_{rec}W_{rec} = W_{rec}^2\]

In general, for a time step \(k \leq t\), the expression is as follows:

\[\frac{\partial h_t}{\partial h_k} = W_{rec}^{t-k}\]

It is now clear that for matrices \(W_{rec}\) with eigenvalues (or singular values for non-square matrices) \(|\lambda| > 1\), the gradients will grow exponentially as \(t\) increases. Conversely, for \(|\lambda| < 1\), the gradients will decay exponentially. To see this, let's use the eigendecomposition of the matrix \(W_{rec}\):

\[W_{rec} = Q \Lambda Q^{-1}\]

Raising it to the \(t\)-th power:

\[W_{rec}^t = Q \Lambda Q^{-1} Q \Lambda Q^{-1} \dots Q \Lambda Q^{-1} = Q \Lambda^t Q^{-1}\]

The traditional approach to solving this problem is to use LSTM or GRU architectures, which replace the multiplication operations with addition. Alternatively, we can use a more geometric approach by applying a constraint on the weights of \(W_{rec}\) such that their eigenvalues (or singular values) remain equal to 1. Essentially, we want to ensure that our transformation does not impact the norm of the vector, but only introduces a rotation. This can be achieved using the Stiefel manifold, which can be defined with respect to the Euclidean topology as:

\[V_k(\mathbb{R}^n) := \{ X \mid X^\top X = \mathbb{I}_k, k \leq n \}\]

Keeping the weights of \(W_{rec}\) on this manifold ensures that gradients will not vanish or explode after an arbitrary number of time steps \(t\). However, this restricts the class of transformations to rotations only.

Comments