Differentiable Loss Function for Indices

Utkarsh_Upadhyay: <p>Hi, I am trying to build a loss function based on Gaussian distribution. In the calculation, this function uses indices of the matrix which (I think) are breaking the graph. Please can somebody help with this? I am attaching the code for the loss function and the output when I print out the <code>grad_fn</code> of each operation.</p>
<p><strong>Code</strong></p>
<pre><code class="lang-auto">class Gaussian(nn.Module):
    def __init__(self, var: float) -&gt; None:
        """
        Initializes Gaussian metric.

        Args:
            var (float): Variance parameter.
        """

        super(Gaussian, self).__init__()
        self.var = var

    def forward(self, preds: torch.Tensor, target: torch.Tensor) -&gt; torch.Tensor:
        """
        Computes the Gaussian metric.

        Args:
            preds (torch.Tensor): Predictions.
            target (torch.Tensor): Target data.

        Returns:
            torch.Tensor: Gaussian metric value.
        """
        preds = torch.softmax(preds, dim=1)
        print('Preds:', preds.grad_fn)
        contact_matrix = preds[:, 1, :, :]
        # print('Contact matrix:', contact_matrix.grad_fn)
        contact_matrix = torch.ge(contact_matrix, 0.5).float()
        # print(contact_matrix.int())



        gaussian_metric_value = self.calculate_distances(contact_matrix, self.var)
        # print('Gaussian metric:', gaussian_metric_value.requires_grad)
        return gaussian_metric_value

    def calculate_distances(self, contact_matrix, var):
        """
        Calculate the distance matrix d_{i,j} for a given contact matrix.
        
        Parameters:
        contact_matrix (torch.Tensor): A 2D torch tensor representing contacts (zeros and ones).
        var (float): The variance parameter for the distance calculation.
        
        Returns:
        torch.Tensor: The omega_sum value.
        """

        device = contact_matrix.device
        contact_matrix = contact_matrix.to(device='cpu')
        indices = torch.nonzero(contact_matrix == 1, as_tuple=False).float()
        # print('Indices:', indices.grad)
        indices.requires_grad = True
        indices.retain_grad()
        indices_diff = indices.unsqueeze(1) - indices.unsqueeze(0)
        # print('Indices:', indices.grad)
        
        indices_diff = indices_diff.pow(2)
        # print('Indices diff:', indices_diff.grad_fn)
        sq_diff = indices_diff.sum(dim=2)  # Shape: (num_contacts, num_contacts)
        # print('Sq diff:', sq_diff.grad_fn)
        d_matrix = torch.exp(-sq_diff/(2 * var))  # Shape: (num_contacts, num_contacts)
        # print('D matrix:', d_matrix.grad_fn)
        sum_exp = d_matrix.sum(dim=1)
        # print('Sum exp:', sum_exp.grad_fn)
        omega = 1 / sum_exp
        # print('Omega:', omega.grad_fn)
        omega_sum = omega.mean()
        omega_sum = -omega_sum.to(device=device) # negative sign for minimization
        return omega_sum
</code></pre>
<p><strong>Output</strong></p>
<pre><code class="lang-auto">Preds: &lt;SoftmaxBackward0 object at 0x14e07ca72230&gt;
Contact matrix: &lt;SliceBackward0 object at 0x14e07ca72230&gt;
Indices: None
Indices: None
Indices diff: &lt;PowBackward0 object at 0x14e07ca72230&gt;
Sq diff: &lt;SumBackward1 object at 0x14e07ca72230&gt;
D matrix: &lt;ExpBackward0 object at 0x14e07ca72230&gt;
Sum exp: &lt;SumBackward1 object at 0x14e07ca72230&gt;
Omega: &lt;MulBackward0 object at 0x14e07ca72230&gt;
</code></pre>
<p>Any insights on this would be beneficial.</p>
<p>Thanks</p>

soulitzer: <p>It is expected that indices are not differentiable, do you observe that the gradients produced are not what you expect.</p>
<p>(You can use <a href="https://pytorch.org/docs/stable/generated/torch.autograd.gradcheck.gradcheck.html" class="inline-onebox" rel="noopener nofollow ugc">torch.autograd.gradcheck.gradcheck — PyTorch 2.3 documentation</a> to validate your gradients against numerically computed ones using this function)</p>

Utkarsh_Upadhyay: <p>Thanks for the reply, but the problem is I don’t know what to expect. However, I can still try to check your suggested method.</p>
<p>I want the whole process to be differentiable so that backpropagation happens properly, which won’t happen as long as there are some non-differentiable parts.</p>

Utkarsh_Upadhyay: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a>, can you please suggest something in this? I couldn’t find anything regarding it. Any insights would be helpful.</p>

ptrblck: <p>Indices are not usefully differentiable, since their gradients would be zero everywhere and undefined or Inf at the integer changes.<br>
I don’t know what your use case is exactly, but keep in mind that the actual values of <code>contact_matrix</code> (assuming these are floating point numbers) should be differentiable in case you can use the <code>indices</code> tensor to index <code>contact_matrix</code> and perform further computations.</p>