How much does the importance of the input pixel change wrt a certain weight

SchulzKilian: <p>I have kind of a hard derivative to calculate with autograd: I want the importance of an input pixel (so derivative of the loss wrt that pixel) derived wrt a specific weight/layer.</p>
<p>So basically d(dL/dI)/dw. How do I make sure autograd understands this? When I just try to derive it like so:</p>
<aside class="quote no-group">
<blockquote>
<p> second_order_grad = autograd.grad(outputs=self.gradients, inputs=layer.weight, grad_outputs=torch.ones_like(self.gradients), retain_graph=True, create_graph=True)</p>
</blockquote>
</aside>
<p>With self.gradients being the gradient of the loss wrt the input and layer.weight the weights of a layer, I get the worning that self.gradients has an empty computation graph, even though I compute it like this:</p>
<pre><code>    self.model.eval()  # Set the model to evaluation mode
    input_data.requires_grad = True  # Set requires_grad to True to compute gradients
    self.label = label

    if input_data.grad is not None:
        input_data.grad.zero_()
    # Forward pass
    # print("for the real image the size is "+str(input_data.size()))
    outputs = self.model(input_data) 
    if self.last_layer_linear:
        self.activations["output"]=outputs
    self.input = input_data 

    target = torch.zeros(outputs.size(), dtype=torch.float)
    target[0][label] = 1.0
    self.target = target
    self.loss = self.default_loss(outputs, target)

    # Backpropagate to compute gradients with respect to the output
    self.loss.backward(retain_graph=True)
    
    # Get the gradients with respect to the input

    self.gradients = input_data.grad.clone().detach()

    self.gradients.requires_grad = True
</code></pre>
<p>If anyone has any ideas I would be very happy!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/schulzkilian">@SchulzKilian</a>,</p>
<p>If you’re trying to do higher-order derivatives, I’d advise you look at the <code>torch.func</code> library and use their higher-order derivative API. It will allow you to compute exactly what you want, without having to do multiple calls to autograd and track first-order gradients.</p>
<p>The docs are here: <a href="https://pytorch.org/docs/stable/func.html" class="inline-onebox" rel="noopener nofollow ugc">torch.func — PyTorch 2.3 documentation</a></p>

SchulzKilian: <p>thanks a lot, im looking into it right now.</p>
<p>Im always surprised by how much the autograd is doing for you. Like to think that it understands the weight values going into the input pixel importance is cool to me, i hope that works</p>

SchulzKilian: <p>to be honest i dont find enough documentation yet to understand torch.func at my level of pytorch understanding, couldnt I just do a double backward pass maybe to first get the gradients of the input and then another backwards pass witht those in order to get the gradient of the weight with respect to the gradient of the input data with respect to the loss?</p>
<p>So something like this as a minimum example:</p>
<pre><code>def backward(self):

    weights = torch.ones_like(self.gradients)
    self.model.zero_grad()

    self.gradients.backward(weights,retain_graph=True)

    
    for name, param in self.model.named_parameters():
        print(name, param.grad)
</code></pre>

AlphaBetaGamma96: <p>Have a look at this thread to get a general idea of what to do with <code>torch.func</code></p><aside class="quote quote-modified" data-post="6" data-topic="99644">
  <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">
    <a href="https://discuss.pytorch.org/t/implementing-calculation-of-the-laplacian/99644/6">Implementing calculation of the Laplacian</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi <a class="mention" href="/u/rfmiotto">@rfmiotto</a>, 
You can do this quite efficiently within the torch.func namespace, which requires version torch2.1+. Here’s an example below for the laplacian with respect to an input x below, 
import torch
from torch.func import hessian, vmap, functional_call

net = Model(*args, **kwargs) #our model

x = torch.randn(batch_size, num_inputs) #inputs
params = dict(net.named_parameters()) #params (need to be an input for torch.func to work)

#functionalize version (params are now an input to the mod…
  </blockquote>
</aside>


SchulzKilian: <p>I’ve played around with this approach now but cant seem to come to a proper result. Because transferred to my approach, the function i am deriving would be the backpropagation right? cause that gives me the gradients of the input wrt the loss, and then i should do the hessian of that with respect to the weights. But how do I put the weights as an argument into the backward propagation function, they are just implicitly there, or am I missing something?</p>

AlphaBetaGamma96: <p>You could have a function that returns <code>d_Loss/d_input</code>, whose arguments are <code>dict(model.named_parameters())</code> and <code>inputs</code>, then take the hessian of that w.r.t your parameters. You can just compose the gradient functions and it should work.</p>
<p>If you can share a minimal reproducible script for your model (or at least the input and output size), it should be possible to do.</p>

SchulzKilian: <p>ohh i understand what you mean. my problem though is that Im not actually working on a model but on an optimizer that uses the hessian (but the optimizer does take the model as an initial argument)</p>
<p>This makes it tough to create a function that takes the parameters and the inputs and gives out the gradients because it would have to compute that on the fly, correct?</p>

AlphaBetaGamma96: <p>It shouldn’t do because you can create the function that defines the gradients at the initial optimization step, then just pass in the parameters at epoch, t, with the inputs as well and it will return the derivatives for the parameters at that epoch.</p>
<p>When creating <code>jacobian = jacrev(fcall, argnums=(0))</code> you create a function whose argument is <code>fcall</code>, which I’ve just called to be a functionalize call of the model (from <code>torch.func.functionalize</code>), then you can pass <code>params</code> and <code>inputs</code> to the <code>jacobian</code> function at every epoch for each new set of parameters.</p>

SchulzKilian: <p>i tried your way and i think its the right thing for me, now I have space issues and will try to run on google collab as soon as possible. Just the quick snippets, to make sure i understood you well:</p>
<ol>
<li></li>
</ol>
<aside class="quote no-group">
<blockquote>
<p>def modelfunction(self,x,target, layer):<br>
if layer is not None:<br>
layer_obj = setattr(self.model, layer[0], layer[1])</p>
<pre><code>    output = self.model.forward(x)
    loss = self.default_loss(output, target)
    return loss
</code></pre>
</blockquote>
</aside>
<p>This is how i did the model function in order to have the layer i want to propagate to as a meaningful argument (i dont assume theres a way for the backpropagation to just understand from the model parameter that the layer is an argument in there?</p>
<ol start="2">
<li></li>
</ol>
<aside class="quote no-group">
<blockquote>
<p>self.jacobian_func = func.jacrev(self.modelfunction)</p>
</blockquote>
</aside>
<p>This is my first jacrev function that should get the gradients</p>
<ol start="3">
<li></li>
</ol>
<aside class="quote no-group">
<blockquote>
<p>for name, param in self.model.named_parameters():<br>
hessian = func.jacrev(self.jacobian_func,  argnums = 2 )(self.input, self.target, (name,param))</p>
</blockquote>
</aside>
<p>And this one how i then get the hessian from the first jacoban_func with respect to eac</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="SchulzKilian" data-post="10" data-topic="204030">
<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/schulzkilian/48/70261_2.png" class="avatar"> SchulzKilian:</div>
<blockquote>
<p><code>self.model.forward(x)</code></p>
</blockquote>
</aside>
<p>You shouldn’t call <code>model.forward</code> as you should call the method directly via <code>model(x)</code>, and when using <code>torch.func</code> you need to create a functionalize version of your model via,</p>
<pre><code class="lang-auto">fcall = lambda params, x: torch.func.functionalize(model, params, x)
</code></pre>
<aside class="quote no-group quote-modified" data-username="SchulzKilian" data-post="10" data-topic="204030">
<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/schulzkilian/48/70261_2.png" class="avatar"> SchulzKilian:</div>
<blockquote>
<aside class="quote no-group">
<blockquote>
<p>self.jacobian_func = func.jacrev(self.modelfunction)</p>
</blockquote>
</aside>
</blockquote>
</aside>
<p>This <code>jacrev</code> call is fine, but you need to specify the argument to which you differentiate, which defaults to <code>argnum=0</code></p>
<aside class="quote no-group" data-username="SchulzKilian" data-post="10" data-topic="204030">
<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/schulzkilian/48/70261_2.png" class="avatar"> SchulzKilian:</div>
<blockquote>
<p>for name, param in self.model.named_parameters():<br>
hessian = func.jacrev(self.jacobian_func, argnums = 2 )(self.input, self.target, (name,param))</p>
</blockquote>
</aside>
<p>For the hessian calculation there a few mistakes, when computing the hessian you need to pass the <code>params</code> as a <code>dict</code> (no need to include the name), and you need to vectorize the hessian call otherwise you’ll compute the hessian w.r.t all samples in your batch and that will OOM.</p>
<p>You want something like this,</p>
<pre><code class="lang-auto">params = dict(net.named_parameters()) 

fcall = lambda params, x: torch.func.functionalize(model, params, x)

def calc_loss(params, x, target):
  output = fcall(params, x)
  loss = default_loss(output, target)
  return loss

calc_jacobian = jacrev(calc_loss, argnums=(0)) #check the args here
calc_hessian = jacrev(calc_jacobian, argnums=0)

hessian_per_sample = vmap(calc_hessian, in_dims=(None,0,0))(params, input, target)
</code></pre>
<p>Also, wasn’t your loss <code>d(dL/dI)/dw</code> so derivative w.r.t input w.r.t params? So why are you computing the Hessian of the params?</p>
<p>If you are computing the Hessian w.r.t the params it’s usually intractable due to its size, but if you’re multiplying it by a vector afterwards it might become tractable via a Hessian-vector product.</p>

SchulzKilian: <aside class="quote no-group">
<blockquote>
<p>Also, wasn’t your loss <code>d(dL/dI)/dw</code> so derivative w.r.t input w.r.t params? So why are you computing the Hessian of the params?</p>
</blockquote>
</aside>
<p>I was hoping thats what im doing? First getting the jacobian without argnums specified so input, and then again the jacobian with argnums three so the model parameters? But maybe theres something Im getting wrong…</p>
<p>The functionalization I now got, thanks, I didnt know i could specify the model as an argument.</p>
<p>Appreciate your help a lot man</p>

SchulzKilian: <p>I might be implementing it wrong in the code.</p>
<p>I use the jacobian thanks to your help and it gives me all the gradients correctly:</p>
<aside class="quote no-group">
<blockquote>
<p>self.jacobian_func = func.jacrev(self.modelfunction, argnums = 2)</p>
</blockquote>
</aside>
<p>Now I choose parts of the input and get a binary matrix in the shape of the input to specify if I want to take something into account.</p>
<p>Then for each parameter w I want to adjust it according to its d(dL/dI)dw for the sum of the inputs I selected.</p>
<p>I know this is not straight forward, thats why I kind of get stuck.<br>
Do you wanna be my tutor? Serious question Im a broke uni student but if you could help me with my problem over like a half an hour session Id definitely pay like 20 bucks lol (not actually lol im kinda being serious)</p>

AlphaBetaGamma96: <p>I’d stick with consistent argument signatures, but as long as you’re aware of the change in arguments it shouldn’t matter.</p>
<p>If you compose the function to compute the partial derivative w.r.t to <code>params</code> and <code>input</code>, you should get a dictionary containing the derivative for each part of the network. From there, you can select a particular layer and get what you need.</p>

SchulzKilian: <p>I’ve tried around with that way, also using vmap but it still seems to be too big for my laptop.<br>
hessian_per_sample = vmap(hessian, in_dims=(None,0,None))(dict(self.model.named_parameters()), self.input, self.target)</p>
<p>After calling it like this I get the error cannot allocate memory… Can I maybe go through only parts of the input and do it? Like, how can I pass the entire input as the input to the hessian, but only differentiate wrt certain pixels and the same for the weight parameters?<br>
How can I pass all the parameters for the model to work but only differentiate wrt some layers?</p>

AlphaBetaGamma96: <p>How many parameters is your model?</p>
<p>Also, can you share a snippet for how your calculating the entire derivative term? Not just 1 line?</p>
<p>If you only want this partial derivative hessian term to then multiple it by a vector (of the same shape as params), then you could try a hessian-vector product?</p>

SchulzKilian: <p>I have a couple of models Im testing with because I want the optimizer to work in a generalizable way, the smallest has 12855045 parameters but I would love for it to work with bigger models. I read about the hessian vector product and it seems great, I dont fully understand it though. Am I losing a lot of information?</p>
<p>These are the lines where I calculate the hessian:</p>
<aside class="quote no-group">
<blockquote>
<pre><code>def modelfunction(self,params ,x,target):
    output = self.fcall(params, x)
    print(f"the shape for the fcall output is {output.shape} the real shape is {self.model(x).shape}")

    loss = self.default_loss(output, target)
    return loss
</code></pre>
<p>self.jacobian_func = func.jacrev(self.modelfunction, argnums = 1)<br>
hessian = func.jacrev(self.jacobian_func,  argnums = 0)<br>
hessian_per_sample = vmap(hessian, in_dims=(None,0,None))(dict(self.model.named_parameters()), self.input, self.target)</p>
</blockquote>
</aside>
<p>In the end product I would not use all the input though to differentiate</p>

AlphaBetaGamma96: <p>Yeah, so if you have a model of over 12M parameters, computing a second derivative is most likely intractable (at least on a single GPU).</p>
<p>What do you want to with the hessian term once you’ve calculated it?</p>

SchulzKilian: <p>i basically want to get a group of input pixels and see for each weight how much the weight was participating into taking these input pixels into account, and then punish each weight by a term relative to how much it took those pixels into account.</p>
<p>So I would get the hessian of the loss wrt input wrt parameters and take a parameter, then get the sum of how much it took into account the marked input pixels.</p>

AlphaBetaGamma96: <p>So, you’re computing the second derivative with a shape <code>[batch_size, num_input, num_params]</code> and you want you want to multiply that by the input shape and get a resultant tensor of shape: <code>[batch_size, num_params]</code>?</p>
<p>But couldn’t just try the <code>weight_decay</code> option for <code>Adam</code> but modify it for the gradient instead? So, you penalize the l2-norm of the gradient (rather than its weight)</p>