High order derivatives in loss terms

username: M_S
cooked: <p>Hi,<br>
I am implementing high order derivatives for PyTorch. All existing implementation I saw require that the target tensor will be scalar, while my implementation is for any arbitrary tensor.</p>
<p>I developed a functon that for a model (or function) <code>model</code> and input values <code>x</code>, it will return derivatives up to the <code>order</code>th order, for all <code>x</code> values. When passing a NN, this function must be called at the start of each step\epoch, because it zeros the grads of the optimizer’s params.</p>
<pre><code class="lang-auto">def get_high_order_analytic_derivative(model, x, order, optimizer=None):

    result = []

    x_grad = model(x)
    result.append(x_grad)

    while order &gt; 0:
        x.grad = None
        x_grad.backward(torch.ones_like(x), create_graph=True, retain_graph=True)
        x_grad = x.grad
        result.append(x_grad)
        order -= 1

    x.grad = None
    if optimizer is not None:
        optimizer.zero_grad()  # remove trace of calculations from the parameters
    return result
</code></pre>
<p>I confirmed the correctness of the function with this example:</p>
<pre><code class="lang-auto">a = torch.randn(10, requires_grad=True)
f, df, ddf = get_high_order_analytic_derivate(torch.sin, a, 2)
f
Out[95]: 
tensor([ 0.7344,  0.9864,  0.4804, -0.9276, -0.5895, -0.2465, -0.9672,  0.9168,
         0.3063, -0.6023], grad_fn=&lt;SinBackward&gt;)
df
Out[96]: 
tensor([ 0.6788, -0.1645,  0.8770,  0.3736,  0.8078,  0.9691, -0.2539, -0.3994,
         0.9519,  0.7983], grad_fn=&lt;CopyBackwards&gt;)
ddf
Out[97]: 
tensor([-0.7344, -0.9864, -0.4804,  0.9276,  0.5895,  0.2465,  0.9672, -0.9168,
        -0.3063,  0.6023], grad_fn=&lt;CopyBackwards&gt;)
f + ddf
Out[98]: tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], grad_fn=&lt;AddBackward0&gt;)
</code></pre>
<p>I want a NN to learn the sine function, and one of the loss terms is MSE between (f + ddf) and zeros_like(f + ddf), over some grid of inputs. However, I found that this specific term <strong>harms</strong> the convergence to my desired model. It think it smoothes my model.</p>
<p>What could be the problem? Is there a problem with using differential equations as loss terms when training NNs? Or maybe my way of calculating the derivatives is not desirable? (notice how I call zero grad so all the calculations shouldn’t change the state of the model)</p>
<p>This is an illustration of adding it as a loss term with growing significance:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/5/0570c8f4704e90df57bbf723acbf3552b40be65d.png" data-download-href="https://discuss.pytorch.org/uploads/default/0570c8f4704e90df57bbf723acbf3552b40be65d" title="for_pytorch"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/5/0570c8f4704e90df57bbf723acbf3552b40be65d_2_690x183.png" alt="for_pytorch" data-base62-sha1="M81JQNlEf2JsaHdvsZDvbdw5LT" width="690" height="183" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/5/0570c8f4704e90df57bbf723acbf3552b40be65d_2_690x183.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/5/0570c8f4704e90df57bbf723acbf3552b40be65d_2_1035x274.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/5/0570c8f4704e90df57bbf723acbf3552b40be65d_2_1380x366.png 2x" data-dominant-color="F5F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">for_pytorch</span><span class="informations">4019×1069 219 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The blue line is actual sin and the orange is my model.</p>

username: tom
cooked: <p>An unsolicited bit of advice:</p>
<aside class="quote no-group" data-username="M_S" data-post="1" data-topic="97567">
<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/m_s/48/6108_2.png" class="avatar"> M_S:</div>
<blockquote>
<pre><code class="lang-auto">    x_grad.backward(torch.ones_like(x), create_graph=True, retain_graph=True)
        x_grad = x.grad
</code></pre>
</blockquote>
</aside>
<p>don’t do that! Just use <code>x_grad, = torch.autograd.grad(....)</code>. It is easy to get a memory leak if you use <code>x.grad</code> in autograd because it uses C++ shared pointers that are not garbage collected and so cycles will wreck havoc.</p>
<p><strong>Edit</strong>: The sign here is wrong and so it’s bogus, sorry!</p>
<p>So one thing to say is that for the sine function <code>f = sin(ωx)</code>, you have that <code>ddf = ω**2 f</code>, so indeed <code>f + ddf</code> in the loss term is something like <code>(1+ω**2) f</code> and would seem to push <code>f</code> towards 0. It might be better to put the function and <code>ddf</code> into separate terms.</p>
<p>More generally, I think it’s always good to look at each loss terms’s contribution to the loss.</p>
<p>Best regards</p>
<p>Thomas</p>

username: Nikronic
cooked: <p>Hi,</p>
<aside class="quote no-group" data-username="M_S" data-post="1" data-topic="97567">
<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/m_s/48/6108_2.png" class="avatar"> M_S:</div>
<blockquote>
<p>Is there a problem with using differential equations as loss terms when training NNs? Or maybe my way of calculating the derivatives is not desirable?</p>
</blockquote>
</aside>
<p>There is no issue with using derivatives in loss term. For instance, you can solve Poisson PDE using its equation as loss function (and applying boundary/initial condition).</p>
<p>I think there might be an issue with your model definition or the way you obtain grads.<br>
Siren paper might be able to help you, please see <a href="https://colab.research.google.com/github/vsitzmann/siren/blob/master/explore_siren.ipynb" rel="noopener nofollow ugc">this notebook</a>.</p>
<p>It solves PDEs by including the equation in loss function and/or supervising derivative of ground truth.</p>
<p>As <a class="mention" href="/u/tom">@tom</a> pointed precisely, best way is to use <code>torch.autograd.grad</code>. For instance,</p>
<pre><code class="lang-auto">def gradient(y, x, grad_outputs=None):
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]
    return grad
</code></pre>
<p>Bests</p>

username: tom
cooked: <aside class="quote no-group" data-username="Nikronic" data-post="3" data-topic="97567">
<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/nikronic/48/71873_2.png" class="avatar"> Nikronic:</div>
<blockquote>
<p>There is no issue with using derivatives in loss term. For instance, you can solve Poisson PDE using its equation as loss function (and applying boundary/initial condition).</p>
</blockquote>
</aside>
<p>Note that the loss term here seems to have the wrong sign between the terms if the frequency matches…</p>

username: M_S
cooked: <aside class="quote no-group" data-username="Nikronic" data-post="3" data-topic="97567">
<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/nikronic/48/71873_2.png" class="avatar"> Nikronic:</div>
<blockquote>
<p><code>torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]</code></p>
</blockquote>
</aside>
<p>Thanks! I updated my code accordingly.<br>
Unfortunately results stay the same.<br>
I get better results with numerical gradients (f’(x) = f(x + d) - f(x - d) / 2d)), but still want to use the analytical if possible, so I’ll keep investigating it.</p>
<p>BTW: ω=1 in my code.</p>

username: M_S
cooked: <p>ddf is -ω^2 f and not ω^2 f .</p>
<p>And in my case ω = 1, so for learning sin(x), I think MSE(f + ddf, zeros) should be a good loss term</p>
<p>Thanks</p>

username: tom
cooked: <p>Oh right. I’m so stupid sometimes. Sigh.<br>
PS: Do you have the code somewhere?</p>