Is it possible to use `parametrizations.orthogonal` inside a parametrization?

dherrera1911: <p>I want to do learning constrained to symmetric positive definite (SPD) matrices. For that I want to use <a href="https://pytorch.org/tutorials/intermediate/parametrizations.html" rel="noopener nofollow ugc">pytorch parametrizations</a>.</p>
<p>I am currently using the spectral decomposition of the SPD matrices, with the following implementation:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class SPD(nn.Module):
    """Constrains the matrix to be symmetric positive definite."""

    def __init__(self):
        super().__init__()

    def forward(self, X, Y):
        # Turn vector X into orthogonal matrix
        eigvec = vec_2_orthogonal_matrix(X)
        # Turn vector Y into positive vector
        eigvals = softmax(Y)
        # Generate SPD matrix
        SPD = torch.einsum('ij,j,jk-&gt;ik', Q, eigvals, Q.t())
        return SPD

    def right_inverse(self, SPD):
        # Take spectral decomposition of matrix
        eigvals, Q = torch.linalg.eigh(SPD)
        # Convert orthogonal matrix to vector
        X = triu_2_vec(
          orthogonal_2_skew(Q)
        )
        # Convert positive vector to vector
        Y = inv_softmax(eigvals)
        return X, Y
</code></pre>
<p>What this code does is that it turns an unconstrained vector <code>X</code> into an orthonormal matrix called <code>eigvecs</code> of size <code>(n, n)</code> using the matrix exponential, and an unconstrained vector <code>Y</code> into a vector called <code>eigvals</code> of size <code>(n)</code> with all positive elements. <code>eigvecs</code> and <code>eigvals</code> are then used to put together the SPD matrix.</p>
<p>When I define my model class, I then use <code>parametrize.register_parametrization(self, "covariance", SPD())</code>, and the parameter <code>covariance</code> is constrained to be SPD.</p>
<p>This code currently works, but I would like to make use of PyTorch’s <a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrizations.orthogonal.html" rel="noopener nofollow ugc">already implemented</a> <code>torch.nn.utils.parametrizations.orthogonal()</code> to make <code>eigvec</code> orthogonal, rather than my implementation. This is because of my understanding that <code>parametrizations.orthogonal()</code> uses trivializations which help with convergence.</p>
<p>However, it is unclear to me how I could use <code>parametrizations.orthogonal()</code> inside another parametrization. This is because <code>parametrizations.orthogonal()</code> is something that is used to constrain an <code>nn.Parameter</code> to be orthogonal, and I don’t have any <code>nn.Parameter</code> in my implementation above.</p>
<p>How could I use <code>parametrizations.orthogonal()</code> in the example above?</p>

dherrera1911: <p>It was possible to do what I wanted. For this, we just need to define the orthogonal matrix as an nn.Parameter inside of the parametrization class. The following code works for this:</p>
<pre><code class="lang-auto">class SPD(nn.Module):
    """Constrains the matrix to be symmetric positive definite."""

    def __init__(self, n_dim):
        super().__init__()
        self.ortho = nn.Parameter(torch.eye(n_dim))
        orthogonal(self, "ortho")
        self.ortho = torch.eye(n_dim)

    def forward(self, Y):
        """ Make SPD matrix."""
        # Turn vector Y into positive vector
        eigvals = softmax(Y)
        # Generate SPD matrix
        SPD = torch.einsum(
          'ij,j,jk-&gt;ik', self.ortho, eigvals, self.ortho.t()
        )
        return SPD

    def right_inverse(self, SPD):
        # Take spectral decomposition of matrix
        eigvals, ortho = torch.linalg.eigh(SPD)
        # Update orthogonal matrix
        self.ortho = ortho
        # Convert positive vector to vector
        Y = inv_softmax(eigvals)
        return Y
</code></pre>
<p>Inside our model, we can then call <code>parametrize.register_parametrization(self, "spd", SPD(n_dim))</code> and the parameter <code>self.spd</code> will be constrained to be SPD (note that we need to pass <code>n_dim</code> when we call <code>parametrize.register_parametrization</code>).</p>