Gradients being returned as zeros in custom likelihood function

armal: <p>I have created a likelihood function (given below) where I manually input vectors of potential parameters into a NN (referred to as self.model in the code) and return the likelihood.</p>
<pre><code>def log_likelihood(self, x):
    ll = torch.zeros(len(x))

    for i in range(len(x)):
        new_params = x[i, :]
        index = 0
        for p in self.model.parameters():
            size = p.numel()
            p.data = new_params[index:index+size].view(p.shape)
            index += size

        logits = self.model(self.X)
        labels = self.Y 
        crit = nn.CrossEntropyLoss(reduction='sum')
        ll[i] = - crit(logits, labels)

    return ll 
</code></pre>
<p>I am trying to also calculate the gradient of the likelihood function w.r.t the parameters I input using the functional jacobian method.</p>
<pre><code>def log_prob_grad(x)
    grad = torch.autograd.functional.jacobian(log_likelihood, x, vectorize=True)
    return grad 
</code></pre>
<p>This gradient function returns a tensor full of zeros. Any help on this problem would be much appreciated.</p>

KFrank: <p>Hi Armal!</p>
<aside class="quote no-group" data-username="armal" data-post="1" data-topic="195494" 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/a/e9bcb4/48.png" class="avatar"> armal:</div>
<blockquote>
<pre><code>        for p in self.model.parameters():
            size = p.numel()
            p.data = new_params[index:index+size].view(p.shape)
</code></pre>
</blockquote>
</aside>
<p>Two related comments:</p>
<p>You shouldn’t be using <code>.data</code>.  It is deprecated (in the public api) and can<br>
lead to errors.</p>
<p>Pytorch does not support <code>Parameter</code>s being modified while gradients are<br>
being tracked.</p>
<p>You should probably do something like:</p>
<pre data-code-wrap="python"><code class="lang-python">for p in self.model.parameters():
    p.grad = None
    with torch.no_grad():
        p.copy_ (new_params ...)   # whatever the details are

logits = self.model (self.X)
labels = self.Y 
crit = nn.CrossEntropyLoss (reduction = 'sum')
loss = - crit (logits, labels)
loss.backward()   # this will populate the grads of model.parameters()
</code></pre>
<p>The <code>.grad</code> properties of the <code>model.parameters()</code> will be the gradient of<br>
<code>loss</code> with respect to <code>new_params</code> (basically <code>x</code>) – they just won’t be directly<br>
attached to <code>new_params</code> (nor <code>x</code>).</p>
<aside class="quote no-group">
<blockquote>
<p>I am trying to also calculate the gradient of the likelihood function w.r.t the parameters I input using the functional jacobian method.</p>
</blockquote>
</aside>
<p>If you really want to do this using <code>functional.jacobian()</code>, you will likely be<br>
better off implementing <code>model</code> with purely functional calls (e.g., use things like<br>
<code>torch.nn.functional.linear()</code> instead of <code>torch.nn.Linear</code>).  This way<br>
you can pass the parameter values in your <code>x</code> into the functional calls, rather<br>
than work around the fact that pytorch doesn’t want you to load new values<br>
into <code>Parameter</code>s while gradients are being tracked.</p>
<p>Best.</p>
<p>K. Frank</p>

AlphaBetaGamma96: <p>Hi Armal (and <a class="mention" href="/u/kfrank">@KFrank</a>)</p>
<p>Another way (expanding upon <a class="mention" href="/u/kfrank">@KFrank</a>’s original answer) would be to use the <code>torch.func</code> namespace to compute derivatives w.r.t parameters in a functionalized manner.</p>
<p>For example,</p>
<pre><code class="lang-auto">net = Model(*args, **kwargs)

crit = nn.CrossEntropyLoss(reduction='sum')

from torch.func import functional_call, jacrev, vmap

def fcall_model(params, inputs): #functionalized call (params now an input)
  return functional_call(net, params, inputs)
  
def log_likelihood(params, inputs, labels):
  logits = fcall_model(params, inputs)
  ll_loss = -1.0 * crit( logits, labels)
  return ll_loss

grads_ll_wrt_params = vmap(jacrev(log_likelihood, argnums=(0)), in_dims=(None, 0, 0))
#You could also used forward-mode AD with jacfwd instead.
</code></pre>
<p>Also, I believe using <code>torch.autograd.functional.jacobian</code> computes the derivatives between samples in the batch, which should be zero by definition. (Unless you have any norm layers in your network).</p>

KFrank: <p>Hi Alpha!</p>
<aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="3" data-topic="195494" 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/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<pre><code class="lang-auto">def fcall_model(params, inputs): #functionalized call (params now an input)
  return functional_call(net, params, inputs)
</code></pre>
</blockquote>
</aside>
<p>Thanks for that.  It looks like <code>functional_call()</code> is the right tool for<br>
painlessly turning a “<code>Module</code>” model that contains its parameters into<br>
a functional model into which you can pass the parameters.</p>
<p>Best.</p>
<p>K. Frank</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>,</p>
<p>Also, I forgot to mention the parameters of the network should be defined as a dictionary with its corresponding names rather than just a list.</p>
<p>For example,</p>
<pre><code class="lang-auto">net = Model(*args, **kwargs)

params = dict(net.named_parameters()) #changed from net.parameters()
</code></pre>

armal: <p>Thank you so much for this. All the suggestions by <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> and <a class="mention" href="/u/kfrank">@KFrank</a> were very helpful in helping me achieve what I needed. The vmap function was also very helpful in being able to vectorize the approach.</p>
<p>This forum is worth its weight in gold due to the community.</p>