How to use orthogonal parameterization with GRU?

user0: <p>I am attempting to create a GRU that uses orthogonal parameterization for the weight matrices stored in <code>weight_hh_l0</code>. From the <a href="https://pytorch.org/docs/stable/generated/torch.nn.GRU.html" rel="noopener nofollow ugc">GRU documentation</a>, I know that <code>weight_hh_l0</code> contains three weight matrices <code>W_hr</code>,<code>W_hz</code>, and <code>W_hn</code> concatenated together. The problem that I’m facing is that because the GRU module stores the matrices in a concatenated format, applying the orthogonal parameterization via <code>torch.nn.utils.parametrizations.orthogonal</code> yields a matrix that is entirely 0 for the bottom two-thirds of the rows.</p>
<p>My question is, <strong>what is the correct way to enforce orthogonality constraints on GRU weight matrices so that <em>each</em> of the submatrices of the GRU weights are orthogonal?</strong> I want <code>W_hr</code>, <code>W_hz</code>, and <code>W_hn</code> to all be orthogonal, instead of <code>W_hr</code> orthogonal with the other matrices set to the zero matrix.</p>
<p>This is a minimal example demonstrating what I’m observing.</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.nn.utils.parametrizations import orthogonal
model_dim = 2
n_batch = 4
max_time = 20

input_tensor = torch.randn((max_time, n_batch, model_dim))

net = nn.GRU(model_dim, model_dim, num_layers=1, batch_first=False)
net = orthogonal(net, name=f"weight_hh_l0")
net.forward(input_tensor)

for name, param in net.named_parameters():
    print(name)
    print(param)
</code></pre>
<p>The result is that the weight matrix <code>weight_hh_l0</code> is only orthogonal in the first third (corresponding to <code>W_hr</code>), but the latter two-thirds (corresponding to <code>W_hz</code>, and <code>W_hn</code>) are all zeros. <strong>This is not what I want, because it does not contain three orthogonal sub-matrices</strong>. In other words, it seems that <code>orthogonal</code> is naïve about <code>weight_hh_l0</code> and treats it as a single weight matrix, instead of three independent matrices concatenated together, each of which has an orthogonality constraint.</p>
<pre><code class="lang-auto">parametrizations.weight_hh_l0.original
Parameter containing:
tensor([[-1.,  0.],
        [ 0., -1.],
        [ 0.,  0.],
        [ 0.,  0.],
        [ 0.,  0.],
        [ 0.,  0.]], requires_grad=True)
</code></pre>
<p>Is there a way that we can enforce orthogonality of all three matrices <code>W_hr</code>, <code>W_hz</code>, and <code>W_hn</code> using <code>torch.nn.GRU</code> and <code>torch.nn.utils.parametrizations.orthogonal</code>? How so?</p>

KFrank: <p>Hi user0!</p>
<aside class="quote no-group" data-username="user0" data-post="1" data-topic="180487" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/u/94ad74/48.png" class="avatar"> user0:</div>
<blockquote>
<p><strong>what is the correct way to enforce orthogonality constraints on GRU weight matrices so that <em>each</em> of the submatrices of the GRU weights are orthogonal?</strong></p>
</blockquote>
</aside>
<p>Probably the simplest way would be to implement your own <a href="https://pytorch.org/tutorials/intermediate/parametrizations.html" rel="noopener nofollow ugc">parametrization</a><br>
that splits <code>weight_hh_l0</code> into the three actual weight matrices, orthogonalizes<br>
them individually, and cats the orthogonal matrices back together.</p>
<p>A sound way to orthogonalize a matrix is to use its <a href="https://pytorch.org/docs/2.0/generated/torch.linalg.svd.html" rel="noopener nofollow ugc">singular-value decomposition</a>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.0.0'
&gt;&gt;&gt; thin = torch.randn (5, 3)
&gt;&gt;&gt; ortho_thin = torch.linalg.svd (thin, full_matrices = False)[0]
&gt;&gt;&gt; ortho_thin.shape
torch.Size([5, 3])
&gt;&gt;&gt; ortho_thin.T @ ortho_thin
tensor([[ 1.0000e+00, -1.0617e-07, -7.4506e-08],
        [-1.0617e-07,  1.0000e+00,  5.4017e-08],
        [-7.4506e-08,  5.4017e-08,  1.0000e+00]])
</code></pre>
<p>(This might not be as good / fast an algorithm as the one under the hood of<br>
<code>torch.nn.utils.parametrizations.orthogonal()</code>, but <code>orthogonal()</code><br>
doesn’t directly fit the use case you describe.)</p>
<aside class="quote no-group">
<blockquote>
<p>Is there a way that we can enforce orthogonality of all three matrices <code>W_hr</code>, <code>W_hz</code>, and <code>W_hn</code> using <code>torch.nn.GRU</code> and <code>torch.nn.utils.parametrizations.orthogonal</code>?</p>
</blockquote>
</aside>
<p>Not that I am aware of.</p>
<p>Best.</p>
<p>K. Frank</p>

user0: <p>Hi, K. Frank!</p>
<p>Thanks for the note. I understand what you’re saying about creating a custom <code>parameterization</code> to work around the way that the GRU weight matrices are concatenated. I was hoping that there was a simpler way, but I suppose there isn’t.</p>
<p>But I don’t understand when you say that I would have to use SVD in place of <code>orthogonal</code>. Is this because <code>orthogonal</code> only works on <code>nn.Module</code>, vice <code>torch.Tensor</code>? Or something else?</p>
<p>Moreover, doesn’t SVD suffer from sign ambiguity? If we’re computing an SVD factorization at each iteration, then we would need to somehow persist sign information so that we’re not using <code>U</code> as one iteration and <code>-U</code> at the next iteration.</p>
<p>Thanks</p>

Lezcano: <p>You’d need to implement that on your own, as all the weights are packed together in the GRU implementation. You can take inspiration from the implementations in <a href="https://github.com/lezcano/geotorch" class="inline-onebox" rel="noopener nofollow ugc">GitHub - lezcano/geotorch: Constrained optimization toolkit for PyTorch</a></p>
<p>The way I’d do it is to go from 3 different matrices into a larger one by putting a orthogonality constraint on the first one and no constraints on the last 2 and then concatenate the three to form the larger tensor. You can even keep the 2/3’s merged into one tensor, which may be slightly more efficient.</p>

user0: <p>Hi Lezcano,</p>
<p>Thanks for the reply. My goal is to have orthogonality constraints on all 3 matrices. Why do you suggest only constraining the first one? That doesn’t seem to address my question.</p>

Lezcano: <p>I misunderstood what you wanted to do. Then do exactly the same thing but with three matrices. You may need to have a look at how <code>P.orthogonal</code> is implemented, and take inspiration on it to implement your version. It shouldn’t be too difficult though.</p>

KFrank: <p>Hi user0!</p>
<aside class="quote no-group" data-username="user0" data-post="3" data-topic="180487" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/u/94ad74/48.png" class="avatar"> user0:</div>
<blockquote>
<p>Is this because <code>orthogonal</code> only works on <code>nn.Module</code>, vice <code>torch.Tensor</code>?</p>
</blockquote>
</aside>
<p>Three things are conspiring to make this more involved than it might have been.</p>
<p>First, <code>orthogonal()</code> doesn’t know anything about <code>GRU</code> (nor is there really any<br>
way for it to).  When you call <code>orthogonal(net, name=f"weight_hh_l0")</code>,<br>
you’re simply telling <code>orthogonal()</code> to act on the <code>"weight_hh_l0"</code> parameter<br>
of your <code>net</code> – <code>orthogonal()</code> doesn’t know that <code>weight_hh_l0</code> should be<br>
interpreted as consisting of three separate matrices that should be individually<br>
orthogonalized.</p>
<p>Second, pytorch does not provide a functional version of <code>GRU</code>.  Consider<br>
<code>Linear</code>, which does have a functional version:  You could (but you don’t need<br>
to do it this way) store <code>weight</code> as a free-standing trainable parameter, apply<br>
<code>orthogonal()</code> to your model and <code>weight</code>, and in your model’s forward function,<br>
pass the orthogonalized <code>weight</code> to <code>torch.nn.functional.linear()</code>.  You<br>
can’t use this scheme with <code>GRU</code>, as it lacks a functional version.</p>
<p>Third (and probably poor design in my opinion), <code>orthogonal()</code> doesn’t expose<br>
its orthogonalization algorithm to the outside world.  You can see from the <a href="https://pytorch.org/docs/2.0/_modules/torch/nn/utils/parametrizations.html#orthogonal" rel="noopener nofollow ugc">code</a><br>
for <code>orthogonal()</code> that its orthogonalization algorithm is intertwined with its<br>
implementation as a <code>parametrization</code>.  This makes it impractical to reuse<br>
<code>orthogonal()</code>'s internal algorithm in a custom <code>parametrization</code> you write<br>
that understands the structure of <code>GRU</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>Moreover, doesn’t SVD suffer from sign ambiguity?</p>
</blockquote>
</aside>
<p>I don’t know enough about pytorch’s <code>svd()</code> and your use case to know whether<br>
this would be a problem in practice, but it’s certainly a theoretical possibility.<br>
Provided that the relevant rows of <code>weight_hh_l0</code> don’t become (nearly)<br>
collinear (which would be a problem for any orthogonalization algorithm), you<br>
should be able to resolve any such sign ambiguity by flipping (if necessary) the<br>
rows of the orthogonalized matrices so that they have positive inner product<br>
with the corresponding rows of the original matrices.</p>
<p>Best.</p>
<p>K. Frank</p>

user0: <p>I’ve implemented a simple series of functions to create parameterizations of GRU weight matrices that are constrained to be orthogonal. Obviously, there is plenty of room for improvement and abstraction to alternative sequential networks (RNN, LSTM), but this is a dirt-simple starting point for that effort.</p>
<p>The overall strategy is to unstack the weight matrices, use Cayley transformations on the individual parts, and then concatenate them together.</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.nn.utils import weight_norm
import torch.nn.utils.parametrize as parametrize


def unpack_gru_weights(X):
    nrow, ncol = X.shape
    return torch.split(X, nrow//3)


def skew_symmetric(X):
    nrow, ncol = X.shape
    assert nrow == ncol  # this function assumes a square matrix
    A = torch.triu(X, diagonal=1)
    # exclude diagonal means A[i,i] = 0 for all i
    # main diagonal and all diagonals below are 0, so we can just add.
    return A - A.transpose(0, 1)  # subtract s.t. A.t() = - A


def caley_map(X):
    A = skew_symmetric(X)
    eye = torch.eye(A.shape[0])
    return torch.linalg.solve(A=eye + A, B=eye - A)


class GruOrthogonal(nn.Module):
    def forward(self, X):
        # The GRU weight matrices are weight_hh_l{k} = [W_hr|W_hz|W_hn]
        #                         and weight_ih_l{k} = [W_ir|W_iz|W_in]
        # We have no need to distinguish the 2 cases, so we name them W_xr, W_xz, W_xn
        w_xr, w_xz, w_xn = unpack_gru_weights(X)
        w_xr = caley_map(w_xr)
        w_xz = caley_map(w_xz)
        w_xn = caley_map(w_xn)
        Xq = torch.cat([w_xr, w_xz, w_xn])
        assert X.shape == Xq.shape
        return Xq


if __name__ == "__main__":
    model_dim = 2
    n_batch = 4
    max_time = 20
    norm_dim = 1

    input_tensor = torch.randn((max_time, n_batch, model_dim))

    net = nn.GRU(model_dim, model_dim, num_layers=1, batch_first=False)
    parametrize.register_parametrization(net, "weight_hh_l0", GruOrthogonal())

    net.forward(input_tensor)

    Q = net.all_weights[0][1]
    print(Q @ Q.t())
</code></pre>
<p>The print gives us a quick visual confirmation that this works: the 2 x 2 sub-matrices along the main diagonal shows that they are 2 x 2 identity matrices, as expected. All other sub-matrices can be ignored, because they are products of sub-matrices that are irrelevant for our goals here. (Numerical imprecision might mean that some of the elements are not precisely 0.)</p>
<p>I elected not to use SVD because (1) I didn’t want to worry about correcting sign indeterminism and (2) I suspect SVD requires the same or greater computational cost, compared to solving the linear system.</p>

Lezcano: <p>That’s exactly what I had in mind. I really like the very clean implementation! The only stylistic point I would make is: You may want to have a look at <code>.mT</code> which does the same thing as your <code>transpose</code> and <code>t</code> but in a more concise way (the <code>m</code> is for “matrix transpose”).</p>

user0: <p>Ah, I see – since I posted this code, I traded <code>.transpose</code> for <code>.T</code>. I’ve read the documentation for <code>.mT</code> and <code>.T</code> and I don’t think <code>.mT</code> is a good choice. The documentation states</p>
<blockquote>
<p>The use of <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.T" rel="noopener nofollow ugc"><code>Tensor.T()</code></a> on tensors of dimension other than 2 to reverse their shape is deprecated and it will throw an error in a future release. Consider <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.mT" rel="noopener nofollow ugc"><code>mT</code></a> to transpose batches of matrices or <code>x.permute(*torch.arange(x.ndim - 1, -1, -1))</code> to reverse the dimensions of a tensor.</p>
</blockquote>
<p>In this usage, we only have 2-tensors, so <code>.T</code> is doing exactly what we want. In a scenario where the user supplies a 3-tensor or higher-order tensor, that usage is unanticipated by these methods, so raising an error is the appropriate outcome.</p>
<p>Indeed, overly-general code is how I got here in the first place: <code>torch.nn.utils.parametrizations.orthogonal</code> didn’t object when I applied it to the weight tensors of a GRU, so I had to do some very annoying tracing to discover that the out-of-the-box orthogonal parameterization is too generalized &amp; too naïve to be used with GRUs.</p>
<p>Also, I’ve realized that <code>unpack_gru_weights</code> is redundant. Instead, I should just do <code>w_xr, w_xz, w_xn = torch.split(W, W.shape[0] // 3)</code>. I would edit that post, but for some reason there’s no edit button.</p>

Lezcano: <p>The point is that <code>.mT</code> is almost always what you want on higher order tensors, and <code>.T</code> is almost never what you want on higher order. But of course, to each their own <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Note that <code>parametrizations.orthogonal</code> is not overly general. It just has no way to know that what you want is not to make the whole matrix orthogonal but split it in three and then make each of them orthogonal.</p>

user0: <aside class="quote no-group" data-username="Lezcano" data-post="11" data-topic="180487">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lezcano/48/40650_2.png" class="avatar"> Lezcano:</div>
<blockquote>
<p>The point is that <code>.mT</code> is almost always what you want on higher order tensors, and <code>.T</code> is almost never what you want on higher order. But of course, to each their own <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>I think if you re-read my post, you’ll see that is exactly what I’m saying. The weight in <code>nn.GRU</code> are 2-tensors, so splitting them also produces 2-tensors, which is why using <code>T</code> is preferable to <code>mT</code>.</p>
<aside class="quote no-group" data-username="Lezcano" data-post="11" data-topic="180487">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lezcano/48/40650_2.png" class="avatar"> Lezcano:</div>
<blockquote>
<p>Note that <code>parametrizations.orthogonal</code> is not overly general. It just has no way to know that what you want is not to make the whole matrix orthogonal but split it in three and then make each of them orthogonal.</p>
</blockquote>
</aside>
<p>Well, it’s definitely over-general because part of the mathematical definition of an orthogonal matrix is that an orthogonal matrix is <em>square</em>. Users can apply <code>torch.nn.utils.parametrizations.orthogonal</code> to non-square matrices; indeed, this is the origination of the unexpected behavior outlined in my first post.</p>
<p>Your reply is consistent with what I mean by <em>naïve</em>. If there’s no intention to make the out-of-the-box function <code>torch.nn.utils.parametrizations.orthogonal</code> compatible with the out-of-the-box classes <code>nn.GRU</code> and <code>nn.LSTM</code>, then there should be a note to that effect somewhere in the docs. Or it should raise an error when it is called on a non-square matrix, which would be consistent with the mathematical definition of orthogonal.</p>
<p>Even better, an alternative way to make <code>parametrizations.orthogonal</code> compatible with native pytorch functionality would be to simply not store the weight matrices in this concatenated format! Then it would work as advertised, without users being compelled to write their own parameterization code.</p>

Lezcano: <p>If you think that any of the implementations in core can be improved, we always accept PRs. Now, I don’t think that what you propose can be achieved without incurring on a perf hit on the most common use case for LSTM and GRU, that is, when they do not have any parametrisation registered.</p>

user0: <p>I created an issue on the PyTorch github last month. It outlines several alternative paths to make the orthogonal parameterization clearer in the case of non-square matrices.  <a href="https://github.com/pytorch/pytorch/issues/102740" class="inline-onebox" rel="noopener nofollow ugc">parameterizations.orthogonal does not work as intended with nn.GRU or nn.LSTM · Issue #102740 · pytorch/pytorch · GitHub</a></p>