Help with Projecting Gradients onto a Hypersphere's Tangent Plane

jumdc: <p>Hi everyone,</p>
<p>I am training a neural network with a loss function denoted as L. I aim to constrain the network’s outputs (z) such that they lie on a hypersphere of dimension d.<br>
To maintain a valid training framework, I need the backpropagated gradients to respect this constraint, meaning they should be tangent to the hypersphere at the corresponding points.<br>
To achieve this, I plan to project the gradients ∂L/∂z, where z is the network’s output, onto the tangent plane of the hypersphere. Based on my understanding, this projection can be computed as:</p>
<p>g_tangent=g−(g⋅z_normalized) z_normalized,</p>
<p>where g=∂L/∂z and z_normalized is the normalized version of z.</p>
<p>However, I’m having difficulty accessing these gradients, modifying them, and ensuring the loss is properly backpropagated after applying this projection.</p>
<p>I’d greatly appreciate any feedback, suggestions, or guidance on implementing this process. Thank you so much for your help!</p>

soulitzer: <p>Registering a backward hook would allow you to access gradients during backward and modify them. <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html#torch.Tensor.register_hook" class="inline-onebox" rel="noopener nofollow ugc">torch.Tensor.register_hook — PyTorch 2.5 documentation</a></p>
<p>Does this sound useful for your case?</p>

jumdc: <p>Hi, thank you for your answer!<br>
Just looking into this function, and it looks like it could work <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>

KFrank: <p>Hi jumdc!</p>
<aside class="quote no-group" data-username="jumdc" data-post="1" data-topic="214694" 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/j/82dd89/48.png" class="avatar"> jumdc:</div>
<blockquote>
<p>I aim to constrain the network’s outputs (z) such that they lie on a hypersphere of dimension d.<br>
To maintain a valid training framework, I need the backpropagated gradients to respect this constraint, meaning they should be tangent to the hypersphere at the corresponding points.</p>
</blockquote>
</aside>
<p>I’m not sure I follow what you are asking, but if I understand your use case, it should<br>
suffice to project your <code>z</code> onto the hypersphere and then compute your loss function<br>
using the values of the projected <code>z</code>.  The gradients with respect to the unprojected <code>z</code><br>
will then naturally be tangent to the hypersphere (because the objective function will<br>
now be independent of any changes in <code>z</code> normal to the hypersphere, as any such<br>
changes will have been projected away).</p>
<p>Here is an illustrative script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

t = torch.tensor ([2.0, 2.0])                            # two-dimensional target point
x = torch.randn (5, 2, requires_grad = True)             # batch of five two-dimensional starting points

print ('x ...')
print (x)

y = x / (x**2).sum (dim = 1, keepdim = True).sqrt()      # project x onto unit circle

print ('y = x projected onto unit circle ...')
print (y)

torch.nn.MSELoss (reduction = 'sum') (x, t).backward()

print ('x.grad (no projection) ...')
print (x.grad)
print ('x.grad &lt;dot&gt; x ...')
print ((x * x.grad).sum (dim = 1))                       # x.grad not tangent to unit circle

x.grad = None

torch.nn.MSELoss (reduction = 'sum') (y, t).backward()

print ('x.grad (with projection) ...')
print (x.grad)
print ('x.grad &lt;dot&gt; x ...')
print ((x * x.grad).sum (dim = 1))                       # this version of x.grad is tangent to unit circle
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.5.1
x ...
tensor([[-0.0404,  1.7260],
        [-0.8140,  1.3722],
        [ 0.5060, -0.4823],
        [-0.7853,  0.6681],
        [-0.4439,  0.1888]], requires_grad=True)
y = x projected onto unit circle ...
tensor([[-0.0234,  0.9997],
        [-0.5102,  0.8600],
        [ 0.7238, -0.6900],
        [-0.7616,  0.6480],
        [-0.9202,  0.3914]], grad_fn=&lt;DivBackward0&gt;)
x.grad (no projection) ...
tensor([[-4.0809, -0.5480],
        [-5.6281, -1.2557],
        [-2.9881, -4.9647],
        [-5.5705, -2.6638],
        [-4.8879, -3.6223]])
x.grad &lt;dot&gt; x ...
tensor([-0.7809,  2.8585,  0.8828,  2.5946,  1.4859], grad_fn=&lt;SumBackward1&gt;)
x.grad (with projection) ...
tensor([[ -2.3698,  -0.0555],
        [ -2.9546,  -1.7528],
        [ -5.5822,  -5.8556],
        [ -3.5439,  -4.1655],
        [ -4.2567, -10.0074]])
x.grad &lt;dot&gt; x ...
tensor([-7.4506e-09,  2.3842e-07,  0.0000e+00,  4.7684e-07, -5.9605e-07],
       grad_fn=&lt;SumBackward1&gt;)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

jumdc: <p>Hi Frank,</p>
<p>Thanks a lot for your detailed answer and insights.</p>
<p>Best,<br>
Julie</p>