How does the autograd function works in nonlinear equations?

username: vmtenorio
cooked: <p>Hi all</p>
<p>I was trying to understand how does the <code>autograd</code> module and <code>backward</code> function work. In some projects, I’ve seen non-linear and non-convex operations like the following:</p>
<pre><code class="lang-auto">z1, z2 = model(x1, x2)

z1 = (z1 - z1.mean(0)) / z1.std(0)
z2 = (z2 - z2.mean(0)) / z2.std(0)
N = z1.shape[0]

c = (z1.T @ z2) / N

loss = -th.diagonal(c).sum()
loss.backward()
</code></pre>
<p>Code adapted from <a href="https://github.com/hengruizhang98/CCA-SSG/blob/main/main.py" rel="noopener nofollow ugc">here</a>.</p>
<p>This operations although they are not complex, they are highly non-linear (std involves squaring the term, we have a division and the derivative of the division is complex, etc.), and it would take a while to do it by hand. However, Pytorch does it automatically, and I’m trying to understand how this is possible.</p>
<p>I’ve tried to check what are the <code>grad_fn</code> of <code>loss</code> and <code>c</code> and <code>z1</code> variables and it involve some <code>MmBackward</code> and <code>DivBackward</code>, but I don’t really see how these can be done using Jacobian and gradient products. I was trying to find the code for these backward functions in Pytorch’s Github repository to see if it helps to understand it, but I’m not able to find it.</p>
<p>I would be really grateful if someone could tell me how this work, high level.</p>
<p>Thanks in advance</p>

username: KFrank
cooked: <p>Hi Victor!</p>
<aside class="quote no-group" data-username="vmtenorio" data-post="1" data-topic="147626" 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/v/e0b2c6/48.png" class="avatar"> vmtenorio:</div>
<blockquote>
<p>I was trying to understand how does the <code>autograd</code> module and <code>backward</code> function work. In some projects, I’ve seen non-linear and non-convex operations</p>
</blockquote>
</aside>
<p>Autograd works in two pieces:  First, “building-block” functions know<br>
how to compute their own gradients.  Second, the autograd machinery<br>
uses the chain rule to compute (numerically) the gradient of a “composite”<br>
function that has been constructed by stringing together a sequence of<br>
these building-block functions.</p>
<p>So, on the one hand, you can build up complicated nonlinearities by<br>
stringing together simple building-block functions, and autograd will<br>
do the rest by using the chain rule.</p>
<p>On the other hand, a single building-block function can be very complicated<br>
and highly nonlinear, provided that whoever wrote it also implemented its<br>
corresponding gradient function (typically packaged as a <code>backward()</code><br>
method).</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt; t = torch.arange (5).float()
&gt;&gt;&gt; t.requires_grad = True
&gt;&gt;&gt; t
tensor([0., 1., 2., 3., 4.], requires_grad=True)
&gt;&gt;&gt; t.pow (3).sum().backward()
&gt;&gt;&gt; t.grad
tensor([ 0.,  3., 12., 27., 48.])
</code></pre>
<p>Ignoring the <code>sum()</code> (which is just a technical device to illustrate the<br>
derivative of <code>pow()</code> for a number of different arguments), this example<br>
has only a single building-block function, namely <code>pow()</code>.</p>
<p>The nonlinear function <code>torch.pow()</code> knows how to compute its gradient<br>
(derivative), and when you call <code>.backward()</code> on a chain of functions that<br>
includes <code>torch.pow()</code>, autograd asks <code>torch,pow()</code> to compute its gradient<br>
when it hits that step in applying the chain rule to the overall composite<br>
function.</p>
<p>There is no magic – just a lot of work in making autograd’s chain-rule<br>
processing work correctly and in making sure that each and every<br>
building-block function computes its gradient correctly.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Marcraven
cooked: <p>This makes me think…</p>
<p>I have a non-linear function I would like to minimize based on some loss. This function involves finding roots of polynomials, rotations and translations in space, intersections with lines and planes, bucketize<br>
… And finally a loss is computed between this and some experimental data.</p>
<p>Could autograd deal with all these steps?</p>