Numerical stability of torch.svd

geoffwoollard: <p>I am working on a project where my model outputs a vector in R9, and then I project onto the nearest orthogonal matrix with the function below, and then have a mean squared error loss between the projected prediction and it’s ground truth value.</p>
<p>The docs mention a gradient stability issue when the singular values get too close to each other.<br>
<a href="https://pytorch.org/docs/stable/generated/torch.linalg.svd.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.linalg.svd.html</a></p>
<p>The singular values of a rotation matrix are (1,1,1), so as the predicted rotation is getting closer to being on SO(3), the singular values will become (1,1,1), and this would cause numerical issues, as per the docs.</p>
<p>Any comments? The only thing I can think of at this point, is to change the encoding of the rotation matrix so that I can project it onto the group without the svd. For example by using a unit quaternion… (I can just norm the 4-vector).</p>
<p>But there are some reasons I’d like the prediction to come directly as a vector in R9, so if I can get it working this way I’d like to.</p>
<pre><code class="lang-auto">def project_onto_so3(rotation_out_group):
    """
    Project a rotation matrix onto SO(3) by finding the closest rotation matrix in Frobenius norm.
    rotation_out_group.shape = (n_batch, 3, 3)
    """
    u, _, v = torch.svd(rotation_out_group, some=False, compute_uv=True)
    ones_3 = torch.ones(len(rotation_out_group), 3, device=rotation_out_group.device)  # TODO: precompute?
    ones_3[:, -1] = torch.linalg.det(u @ v)
    rotation_in_group = u @ torch.diag_embed(ones_3) @ v
    return rotation_in_group
</code></pre>

KFrank: <p>Hi Geoff!</p>
<aside class="quote no-group" data-username="geoffwoollard" data-post="1" data-topic="193136" data-full="true">
<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/geoffwoollard/48/66014_2.png" class="avatar"> geoffwoollard:</div>
<blockquote>
<p>The singular values of a rotation matrix are (1,1,1), so as the predicted rotation is getting closer to being on SO(3), the singular values will become (1,1,1), and this would cause numerical issues, as per the docs.</p>
</blockquote>
</aside>
<p>I think your concern is valid.  (For a generic matrix that isn’t orthogonal,<br>
your singular-value-decomposition approach is a fine way to orthogonalize<br>
it, but, for the reason you discussed, if the matrix starts out being close to<br>
orthogonal, you won’t be able to backpropagate through the orthogonalization<br>
in a numerically stable way.)</p>
<aside class="quote no-group">
<blockquote>
<p>I’d like the prediction to come directly as a vector in R9, so if I can get it working this way I’d like to.</p>
</blockquote>
</aside>
<p>A standard iterative orthogonalization algorithm that should meet your<br>
needs is illustrated in this script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

torch.random.manual_seed (2023)

m = torch.randn (3, 3)

def orthoA (a):
    u, s, vh = torch.linalg.svd (a)
    o = u @ vh
    return o

def orthoB (a):   # iterative method that avoids svd
    nIter = 5     # this is naive -- could use a convergence criterion
    o = a
    for  i in range (nIter):
        o = (o + torch.linalg.inv (o.T)) / 2
    
    return o

m = torch.randn (3, 3)

print ('m = ...')
print (m)

oA = orthoA (m)
oB = orthoB (m)

print ('oA = ...')
print (oA)

print ('torch.allclose (oA, oB, atol = 1.e-7) =', torch.allclose (oA, oB, atol = 1.e-7))
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">m = ...
tensor([[-0.0794, -1.0846, -1.5421],
        [ 0.9377, -0.9787,  2.0930],
        [ 1.0231,  0.5431,  0.6514]])
oA = ...
tensor([[ 0.2961, -0.6718, -0.6790],
        [ 0.2847, -0.6165,  0.7341],
        [ 0.9118,  0.4106, -0.0088]])
torch.allclose (oA, oB, atol = 1.e-7) = True
</code></pre>
<p>Pytorch has no problem backpropagating through such an iterative<br>
computation (although for algorithms that converge more slowly, many<br>
iterations can be necessary and the computation graph can become large).</p>
<p>This algorithm fails if <code>m</code> is singular, but that is to be expected.</p>
<p>You can look at this iterative algorithm as generating the orthogonal <code>Q</code><br>
term of the <code>QP</code> <em>polar decomposition</em> of <code>m</code>.</p>
<p>As an aside, you should switch from <code>torch.svd()</code> to <code>torch.linalg.svd()</code><br>
in any other code you have as the former has been deprecated.  Note that<br>
the two versions differ slightly in syntax.</p>
<p>Best.</p>
<p>K. Frank</p>