Calculate gradient with respect to data label

Rahn: <p>Hi everyone,</p>
<p>I’m working on implementing a technique from a <a href="https://arxiv.org/pdf/2303.10837" rel="noopener nofollow ugc">research paper, FedML-HE</a> where I need to calculate gradients with respect to data labels. Here’s the approach I’m following:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/5/857e18e9933724b1d450ae09cb71182dfec0c623.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/857e18e9933724b1d450ae09cb71182dfec0c623" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/5/857e18e9933724b1d450ae09cb71182dfec0c623_2_690x235.jpeg" alt="image" data-base62-sha1="j2VFhplt7oxkeEIa7WMbiDTeDCP" width="690" height="235" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/5/857e18e9933724b1d450ae09cb71182dfec0c623_2_690x235.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/8/5/857e18e9933724b1d450ae09cb71182dfec0c623.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/5/857e18e9933724b1d450ae09cb71182dfec0c623.jpeg 2x" data-dominant-color="E6E6E6"><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">image</span><span class="informations">976×333 105 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>Calculate the gradient of the loss with respect to the model’s parameters (grad1).<br>
Calculate the gradient of grad1 with respect to the data labels.<br>
However, I’m encountering an issue where the gradient with respect to the data labels is always None. It seems that the data labels (y) are not part of the computational graph (y.grad_fn returns None).</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.nn.utils import parameters_to_vector

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Load CIFAR10 dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=1, shuffle=True)

model = LeNet()
criterion = nn.CrossEntropyLoss()

# Get a batch of data
x, y = next(iter(trainloader))

y = y.float().requires_grad_(True)

output = model(x)
loss = criterion(output, y.long())

first_order_grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)

Jm_list = []

for grad in first_order_grads:
    if grad is not None:
        for grad_element in parameters_to_vector(grad): 
            Jm = torch.autograd.grad(grad_element, y, retain_graph=True, allow_unused=True)[0]
            Jm_list.append(Jm)
</code></pre>
<p>I’m looking for advice on two points:</p>
<p>How can I include data labels in the computational graph?<br>
Are there alternative methods to calculate gradients with respect to data labels effectively?<br>
Any insights or references to similar implementations would be greatly appreciated!</p>

ptrblck: <aside class="quote no-group" data-username="Rahn" data-post="1" data-topic="206748">
<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/r/838e76/48.png" class="avatar"> Rahn:</div>
<blockquote>
<p>However, I’m encountering an issue where the gradient with respect to the data labels is always None. It seems that the data labels (y) are not part of the computational graph (y.grad_fn returns None).</p>
</blockquote>
</aside>
<p>This is expected since the target wasn’t used in the forward pass and no trainable parameters of the model were used to create it. Thus changing the parameters in any direction will not influence the target at all, which is represented by the <code>None</code> gradient.</p>
<p>Could you explain your use case in more detail and why you are expecting to see gradients?</p>

Rahn: <p>Thank you for the explanation for why taking derivative to the data label always returns None.</p>
<p>The main idea from the <a href="https://arxiv.org/pdf/2303.10837" rel="noopener nofollow ugc">research paper</a> that discusses the sensitivity of model gradients in response to changes in input data labels. The paper suggests that a larger gradient in response to a particular data label indicates a higher sensitivity of the model to that specific input.</p>
<p>Here’s what I understand: The sensitivity analysis aims to quantify how much the gradient (or direction and magnitude of model learning) will change with the data label for each data point, which to identifies which sensitive the model parameters to the training data.</p>

qq-me: <p>The way I would do it is</p>
<pre data-code-wrap="py"><code class="lang-py">y = y.long()
y.requires_grad = True
yhat = model(x)
loss = criterion(yhat, y)
loss.backward()
print(y.grad)
</code></pre>
<p>Also I am not exactly sure about this, but you might also need to add this:</p>
<pre data-code-wrap="py"><code class="lang-py">yopt = torch.optim.SGD([y])
</code></pre>
<p>just to make it so that pytorch calculates gradients for y. But maybe this isn’t needed.</p>

ptrblck: <aside class="quote no-group" data-username="Rahn" data-post="3" data-topic="206748">
<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/r/838e76/48.png" class="avatar"> Rahn:</div>
<blockquote>
<p>sensitivity of model gradients in response to changes in input data labels.</p>
</blockquote>
</aside>
<p>Would this mean that you would like to compute the gradients of all trainable parameters w.r.t. <code>loss0</code> (let’s call these gradients <code>g0</code>) and compare them to gradients computed w.r.t. to another <code>loss1</code> (let’s call these <code>g1</code>)?<br>
If so, I don’t think the target itself has to be differentiable, but it sounds as if you want to change the labels allowing you to recompute the gradients of model parameters.</p>
<p>Let me know, if I misunderstood your use case.</p>

Rahn: <p>Thanks for replying.</p>
<p>The gradients I’m trying to calculate after getting the loss from the model prediction and ground truth label are:</p>
<ol>
<li>The gradients of the loss w.r.t. all model parameters as <code>g0</code></li>
<li>The gradients of <code>g0</code> w.r.t. the label as <code>g1</code>. Because the ground truth labels are not used in the forward pass, now I modify <code>g1</code> as the gradients of <code>g0</code> w.r.t. the model prediction.</li>
</ol>
<p>Therefore, <code>g0</code> would be regular backward propagation. Next, <code>g1</code> would give how much each gradient responses to the model output.</p>
<pre><code class="lang-auto">class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Load CIFAR10 dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=1, shuffle=True)

model = LeNet()
criterion = nn.CrossEntropyLoss()

# Get a batch of data
x, y = next(iter(trainloader))

output = model(x)
loss = criterion(output, y.long())

for param_layer in model.parameters():
    g0= torch.autograd.grad(outputs=loss, inputs=param_layer, create_graph=True)[0]
    print(f"    g0 size: {g0.shape}")

    for param in parameters_to_vector(g0):
        # g1 = torch.autograd.grad(outputs=param, inputs=y, create_graph=True, allow_unused=True)[0] # replace y as output because y is not used in the model
        g1 = torch.autograd.grad(outputs=param, inputs=output, create_graph=True, allow_unused=True)[0]
</code></pre>

Rahn: <p>Thank you for the reply.</p>
<p>y.grad would give me None values because it’s not used in the model except for calculating the loss.</p>

KFrank: <p>Hi Rahn!</p>
<p>The short story is pass your labels as a <code>FloatTensor</code> to<br>
<code>CrossEntropyLoss</code>.  Specifically, if <code>output</code> is a <code>FloatTensor</code> with<br>
shape <code>[nBatch, nClass]</code>, then <code>y</code> should also have that shape,<br>
that is, have an explicit class dimension.</p>
<aside class="quote no-group quote-modified" data-username="Rahn" data-post="1" data-topic="206748" 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/r/838e76/48.png" class="avatar"> Rahn:</div>
<blockquote>
<p>I’m encountering an issue where the gradient with respect to the data labels is always None.<br>
…</p>
<pre><code class="lang-auto">criterion = nn.CrossEntropyLoss()
...
y = y.float().requires_grad_(True)

output = model(x)
loss = criterion(output, y.long())
</code></pre>
</blockquote>
</aside>
<p>Your use of <code>CrossEntropyLoss</code> and the conversion of <code>y</code> to <code>y.float()</code><br>
suggests that <code>y</code> consists of <em>integer</em> categorical class labels.  Because<br>
these are discrete, they are inherently not differentiable.</p>
<p>Converting to <code>float</code>, setting <code>requires_grad = True</code>, and converting<br>
back to <code>long</code> doesn’t fix this.  (The result of <code>y.long()</code> no longer<br>
carries <code>requires_grad = True</code>.)</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">first_order_grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)

Jm_list = []

for grad in first_order_grads:
    if grad is not None:
        for grad_element in parameters_to_vector(grad): 
            Jm = torch.autograd.grad(grad_element, y, retain_graph=True, allow_unused=True)[0]
            Jm_list.append(Jm)
</code></pre>
</blockquote>
</aside>
<p>But yes, once you address your issue of discrete, non-differentiable<br>
labels and how you use <code>CrossEntropyLoss</code>, your approach to<br>
computing your desired second partial derivative is basically correct.</p>
<aside class="quote no-group">
<blockquote>
<p>How can I include data labels in the computational graph?</p>
</blockquote>
</aside>
<p><code>CrossEntropyLoss</code> has two modes:  In the first, you supply a single<br>
integer class label per sample (so your labels would be a <code>LongTensor</code><br>
with no class dimension).  In the second, you supply a set of class<br>
probabilities (so your labels are a <code>FloatTensor</code> with an explicit class<br>
dimension).  Such class probabilities are sometimes called “soft” or<br>
“probabilistic” labels.</p>
<p>In the simplest approach you could:</p>
<pre data-code-wrap="python"><code class="lang-python">y = torch.nn.functional.one_hot (y, nClass).float().requires_grad_ (True)
</code></pre>
<p><code>y</code> will now have an explicit class dimension and you will be able to<br>
compute gradients with respect to <code>y</code>.</p>
<p>There are some nuances (which may or may not matter for your use<br>
case).  <code>y</code> should be understood as a (discrete) probability distribution<br>
over classes.  Therefore each element should be between <code>0.0</code> and<br>
<code>1.0</code> and <code>y</code> should sum to <code>1.0</code>.  (The result of <code>one_hot()</code> satisfies<br>
this.)</p>
<p>But when you think about varying an element of <code>y</code> so as to compute<br>
a gradient with respect to <code>y</code>, it doesn’t necessarily make logical sense<br>
to have an element of <code>y</code> be exactly equal to <code>0.0</code> or <code>1.0</code>, because<br>
if it were to vary below <code>0.0</code> or above <code>1.0</code>, the resulting <code>y</code> would fail<br>
(by an infinitesimal amount) to be a valid probability distribution.</p>
<p>In a similar vein, if you vary a single element of <code>y</code> in isolation, <code>y</code> will<br>
no longer sum to <code>1.0</code> and again not be a valid probability distribution.</p>
<p>Last, the “sensitivities” of your parameter gradients with respect to your<br>
labels are just the second partial derivatives of your loss with respect<br>
to your model parameters and your labels.  (These form a subset of the<br>
full “hessian” of your loss with respect to the parameters and labels.)</p>
<p>You’re allowed, if you so choose, to first compute the gradient of the<br>
loss with respect to the labels, and then compute the jacobian of that<br>
gradient with respect to the model parameters.  (That is to say, the<br>
hessian is symmetric.)  It turns out to be a little more convenient to<br>
compute the gradients is that order (but you can do it either way).</p>
<p>Here is a script that shows how to compute your desired sensitivities<br>
(computing the label gradient first and and accounting for the nuances<br>
due to the “soft” labels being a probability distribution):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

nClass = 5

model = torch.nn.Linear (3, nClass)
inp = torch.randn (3)

prd = model (inp)

targetIndex = torch.tensor (2)                                             # target as integer categorical index
targetOneHot = torch.nn.functional.one_hot (targetIndex, nClass).float()   # target as "soft" probabilities

lossIndex = torch.nn.CrossEntropyLoss() (prd, targetIndex)
lossOneHot = torch.nn.CrossEntropyLoss() (prd, targetOneHot)

print ('torch.equal (lossIndex, lossOneHot):', torch.equal (lossIndex, lossOneHot))

def softHot (prob, ind, nc):                                               # soft target not saturated at 0 and 1
    pLo = (1.0 - prob) / (nc - 1)
    soft = torch.full ((nc,), pLo)
    soft[ind] = prob
    return  soft

targetSoftHot = softHot (0.95, targetIndex, nClass)

lossSoftHot = torch.nn.CrossEntropyLoss() (prd, targetSoftHot)

print ('lossOneHot: ', lossOneHot)                                         # the two loss values are similar
print ('lossSoftHot:', lossSoftHot)

targetSoftHot.requires_grad = True
targetNorm = targetSoftHot / targetSoftHot.sum()                           # affects grad, so not a no-op

loss = torch.nn.CrossEntropyLoss() (prd, targetNorm)

grad = torch.autograd.grad (loss, targetSoftHot, create_graph = True)[0]
print ('grad.shape:', grad.shape)

hessW = []
for  g in grad:
    hessW.append (torch.autograd.grad (g, model.weight, retain_graph = True)[0])

hessW = torch.stack (hessW)
print ('hessW.shape:', hessW.shape)

hessB = []
for  g in grad:
    hessB.append (torch.autograd.grad (g, model.bias, retain_graph = True)[0])

hessB = torch.stack (hessB)
print ('hessB.shape:', hessB.shape)

print ('hessB = ...')                                                      # second partials of loss with respect to
print (hessB)                                                              # target and bias, 5 elements each
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.3.1
torch.equal (lossIndex, lossOneHot): True
lossOneHot:  tensor(1.4802, grad_fn=&lt;DivBackward1&gt;)
lossSoftHot: tensor(1.4942, grad_fn=&lt;DivBackward1&gt;)
grad.shape: torch.Size([5])
hessW.shape: torch.Size([5, 5, 3])
hessB.shape: torch.Size([5, 5])
hessB = ...
tensor([[-0.9875,  0.0125,  0.9500,  0.0125,  0.0125],
        [ 0.0125, -0.9875,  0.9500,  0.0125,  0.0125],
        [ 0.0125,  0.0125, -0.0500,  0.0125,  0.0125],
        [ 0.0125,  0.0125,  0.9500, -0.9875,  0.0125],
        [ 0.0125,  0.0125,  0.9500,  0.0125, -0.9875]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>