Unexpected behavior when using torch.autograd.functional.jacobian with multiple inputs/outputs neural network

username: ziv_chen
cooked: <p>I’m implementing a neural network to solve ODEs. I want to encode several points at once and get several results of the ODE (one result per corresponding input) at the same run.<br>
Currently I’m running a toy model where I try to get <code>u(x)=sin(x)</code> according to the loss function: <code>du_dx - cos(x)</code>.</p>
<p>The network structure is:</p>
<details>
<summary>
network structure</summary>
<pre><code class="lang-auto">class NetworkClass(nn.Module):
    def __init__(self, n_nodes=8):
        super().__init__()
        self.fc_in = nn.Linear(n_nodes, n_nodes)
        self.activation = nn.Tanh()
        self.q_layer = nn.Linear(n_nodes, n_nodes)
        self.activation_2 = nn.Tanh()
        self.fc_out = nn.Linear(n_nodes, n_nodes)

    def forward(self, inputs):
        out_fc_in = self.fc_in(inputs.unsqueeze(0))
        out_activation = self.activation(out_fc_in)
        out_q_layer = self.q_layer(out_activation.squeeze(0).squeeze(0))
        out_activation_2 = self.activation_2(out_q_layer)
        out_fc_out = self.fc_out(out_activation_2).squeeze(0).squeeze(0)
        return out_fc_out

</code></pre>
</details>
<p>The gradients:</p>
<details>
<summary>
gradients</summary>
<pre><code class="lang-auto">for epoch in range(self.epochs):
    u_out_arr = []
    d_u_dx_arr = []
    loss_single_batch_arr = []
    loss_single_batch_no_beta_arr = []
    self.opt.zero_grad(set_to_none=False)
    collocation_points = torch.reshape(self.collocation_points, [self.n_points//self.n_nodes, self.n_nodes])
    u_out = self.model(collocation_points)
    u_out_arr.append(u_out)  # for plotting and derivative calculation check
    d_u_dx = torch.zeros(self.n_points)
    for i in range(self.n_batches):
        jac_func = jacobian(self.model, collocation_points[i], create_graph=True, strict=True)
        d_u_dx[i*self.batch_size:(i+1)*self.batch_size] = torch.diag(jac_func, 0)
    d_u_dx_arr.append(d_u_dx)
    loss= self.loss_fn(collocation_points, u=torch.flatten(u_out), d_u_dx=d_u_dx,
                                     target_fn=self.target_fn_dict)
    loss.backward()
    loss_single_batch_arr.append(loss)

    self.opt.step()
</code></pre>
</details>
<p>When computing derivatives using jacobian (only the diagonal because I need each output’s derivative with respect to the corresponding input):</p>
<ul>
<li>With <code>in_out_size=1</code>: derivatives correctly correspond to spatial derivatives</li>
<li>With <code>in_out_size=n_nodes</code>: derivatives don’t match expected spatial derivatives</li>
</ul>
<p>Question: Why does increasing input/output dimensions affect the derivative computation?</p>
<p>Also, I’ve found that when adding the term of the original_function_loss to the loss_fn in that manner, when beta weights the loss to the original only (beta =1) or the derivative only (beta = 0) or somewhere in between, enables the network to get the accurate results (which is not surprising) by beta as small as 1e-7 (which is very surprising!):</p>
<pre><code class="lang-auto">pde_loss_1 = d_u_dx - target_fn["d_u_dx"](torch.flatten(collocation_point))
original_function_loss  = u - target_fn["u"](torch.flatten(collocation_point))
mse_loss = torch.mean(self.beta * original_function_loss ** 2 + (1 - self.beta) * pde_loss_1 ** 2)
</code></pre>
<p>results for the two cases (without the beta):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff" title="results_pytorch"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff_2_690x388.jpeg" alt="results_pytorch" data-base62-sha1="fH69LmCggKEVzjRUinT5AiZhFEz" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff_2_690x388.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff_2_1035x582.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff.jpeg 2x" data-dominant-color="FBFBFB"><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">results_pytorch</span><span class="informations">1280×720 65.5 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>Thanks in advance!</p>

username: KFrank
cooked: <p>Hi Ziv!</p>
<aside class="quote no-group quote-modified" data-username="ziv_chen" data-post="1" data-topic="215580" 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/z/49beb7/48.png" class="avatar"> ziv_chen:</div>
<blockquote>
<pre><code class="lang-auto">class NetworkClass(nn.Module):
    def __init__(self, n_nodes=8):
        super().__init__()
        self.fc_in = nn.Linear(n_nodes, n_nodes)
        self.activation = nn.Tanh()
        self.q_layer = nn.Linear(n_nodes, n_nodes)
        self.activation_2 = nn.Tanh()
        self.fc_out = nn.Linear(n_nodes, n_nodes)
</code></pre>
<p>…<br>
When computing derivatives using jacobian (only the diagonal because I need each output’s derivative with respect to the corresponding input):<br>
…<br>
Question: Why does increasing input/output dimensions affect the derivative computation?</p>
</blockquote>
</aside>
<p>If I understand correctly what you are asking, when <code>in_out_size</code> is<br>
greater than one, the various values along the <code>n_nodes</code> dimension<br>
are mixed together – multiple time with non-linearities – so the<br>
dependence of the, say, first output value on the the first input value<br>
is not the same as it would have been if <code>n_nodes</code> had been one.<br>
So even though you only look at the diagonal values of the jacobian,<br>
even the diagonal values of the jacobian depend on <code>n_nodes</code>.</p>
<p>Furthermore – more trivially – the weights in the <code>Linears</code> are generated<br>
randomly and even the “same” weight value will differ when using<br>
different values of <code>n_nodes</code>, so the different weight values will lead<br>
to different values of the jacobian elements.</p>
<p>If this is not really what you are asking about, please you post a simplified<br>
(e.g., no need for any loop over epochs), fully-self-contained, runnable<br>
example script that shows the differing derivative values, together with<br>
the output you get when you run it.</p>
<p>Best.</p>
<p>K. Frank</p>

username: ziv_chen
cooked: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>,</p>
<p>Thanks for the answer, I really appreciate it!</p>
<p>I understand what you say about the weights but if I look at the network as a “black-box” that is just some function that depends on <code>theta</code> (weights and biases) and have <code>n_nodes</code> inputs and <code>n_nodes</code> outputs, then I still have trouble seeing why does changing the nodes amount will make a difference for the jacobian…</p>
<p>Can you please elaborate about it a little bit? both about the non-linearities and about the weights.</p>
<p>Thanks a lot in advance,<br>
Ziv</p>

username: KFrank
cooked: <p>Hi Ziv!</p>
<aside class="quote no-group" data-username="ziv_chen" data-post="4" data-topic="215580" 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/z/49beb7/48.png" class="avatar"> ziv_chen:</div>
<blockquote>
<p>I still have trouble seeing why does changing the nodes amount will make a difference for the jacobian…</p>
</blockquote>
</aside>
<p>The real question is why would you expect changing <code>n_nodes</code> <em>not</em> to<br>
change the jacobian?</p>
<p>Here’s a concrete example:  Let’s say we have a super-trivial network<br>
that consists just of two <code>Linears</code> (without <code>bias</code> and no intervening<br>
non-linear “activation” layer) and let’s give each <code>Linear</code> the same<br>
<code>weight</code> matrix.</p>
<p>For <code>n_nodes = 1</code>, let the weight matrix be <code>[[1.]]</code> and for <code>n_nodes = 2</code>,<br>
be <code>[[1., 1.], [1., 1]]</code>.  When combined into a  two-layer network,<br>
the two weight matrices collapse together into a single weight matrix,<br>
basically <code>weight_combined = weight_1 @ weight_2</code>, which becomes<br>
<code>[[1.]]</code> for <code>n_nodes = 1</code> and <code>[[2., 2.], [2., 2.]]</code> for <code>n_nodes = 2</code>.</p>
<p>You can read right off the <code>weight_combined</code> matrix that the (single)<br>
diagonal element of the (one-element) jacobian matrix for <code>n_nodes = 1</code><br>
is <code>1.</code>, while the <code>[0, 0]</code> diagonal element of the jacobian matrix for<br>
<code>n_nodes = 2</code> is <code>2.</code>.  So the value of the jacobian depends on the<br>
value of <code>n_nodes</code>.  But, again, why would it not?</p>
<p>Best.</p>
<p>K. Frank</p>