Answers.Body: <p>To back-propagate through non-differentiable operations, without designing a soft approximation for them, is possible to recur to a <a href="https://pytorch.org/docs/stable/distributions.html#score-function" rel="nofollow noreferrer">score gradient estimator</a>:</p>
  <pre><code>class CombinedNetwork(torch.nn.Module):
      def __init__(self, input_size, hidden_size, num_classes):
          super().__init__()
          self.classifier_fc1 = nn.Linear(input_size, hidden_size)
          self.classifier_fc2 = nn.Linear(hidden_size, num_classes)
          self.regression_fc = nn.Linear(input_size, num_classes)
  
  
      def forward(self, x):
          # Classification branch
          classifier_x = F.relu(self.classifier_fc1(x))
          logits = self.classifier_fc2(classifier_x)
  
          regression_value = self.regression_fc(x)
          return regression_value, logits
  
  # forward
  v, logits = model(x)
  
  # define Categorical distribution from predicted probabilities
  cat = torch.distributions.Categorical(logits=logits)
  indices = logits.argmax(-1, keepdims=True)
  
  # obtain values
  values = torch.gather(v, dim=1, index=indices)
  
  # setup score function, e.g. from MSE loss
  reward = (targets - values).square()
  loss = (-cat.log_prob(indices) * reward).mean()
  
  # optimization step
  ...
  </code></pre>
  <p>Notice that I've assumes a squared loss between the <code>targets</code> and the predicted values: you can change it as you need. Once the model is trained, you can avoid defining the Categorical and directly index the value by according to the argmax of the <code>logits</code>.</p>
  <p>NOTE: training may be unstable, if so clip the gradients and rescale the <code>targets</code> (and so <code>values</code>).</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-07-16T16:05:44.013Z
Answers.LastActivityDate: 2023-07-16T16:05:44.013Z
Answers.UserDisplayName: Luca Anzalone
Body: <p>I am attempting to utilize two networks: a classifier and a linear network. Based on the output class of the first network, my goal is to retrieve the corresponding value from the linear network using indexing. For instance, if the first network's output is [0.1, 0.8, 0.1], and the second network's output is [22, 30, 10], it implies that class two is selected, and the final output should be 30. Since this involves backpropagation through argmax, I am using the soft-argmax strategy described in the literature.</p>
  <p>Following is the code for soft-argmax</p>
  <pre><code>class SoftArgmax1D(torch.nn.Module):
      def __init__():
          self.base_index = 0
          self.step_size = 1
          self.softmax = torch.nn.Softmax(dim=1)
  
      def forward(self, x):
          smax = self.softmax(x*100)
          end_index = self.base_index + x.size()[1] * self.step_size
          indices = torch.arange(start=self.base_index,
                                 end=end_index,
                                 step=self.step_size)
          return torch.matmul(smax, indices.float()) 
  </code></pre>
  <p>Following is the final architecture:</p>
  <pre><code>class CombinedNetwork(nn.Module):
      def __init__(self, input_size, hidden_size, num_classes):
          super(CombinedNetwork, self).__init__()
          self.classifier_fc1 = nn.Linear(input_size, hidden_size)
          self.classifier_fc2 = nn.Linear(hidden_size, num_classes)
          self.regression_fc = nn.Linear(input_size, num_classes)
          self.smax = SoftArgmax1D()
  
  
      def forward(self, x):
          # Classification branch
          classifier_x = F.relu(self.classifier_fc1(x))
          classifier_x = self.classifier_fc2(classifier_x)
          classification_value = self.smax(classifier_x)
          
          regression_value = self.regression_fc(x)
          
          reg_op = regression_value.gather(1, classification_value.round().unsqueeze(1)).squeeze(1)
          
          return reg_op
  </code></pre>
  <p>The issue arises when indexing the regression_value requires an integer input. However, converting classification_value, which essentially represents the index, to an integer causes the backpropagation to be truncated.</p>
  <p>Do you have any suggestions on how to achieve backpropagation for this code?</p>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2023-07-13T17:35:28.340Z
Comments.Text: Can you cast the value to `int` instead of rounding? Because `.round()` seems to be not differentiable
Comments.UserDisplayName: Luca Anzalone
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2023-07-14T09:37:30.403Z
Comments.Text: Upon rechecking, it appears that using `.round()` is fine. However, converting it to an integer (using `.long()`) disrupts the flow in backpropagation.
Comments.UserDisplayName: Subrat Prasad
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2023-07-16T16:07:48.237Z
Comments.Text: I've added a new answer. Please review it: if addresses your issue I'll delete the old one.
Comments.UserDisplayName: Luca Anzalone
Tags: <backpropagation><pytorch><regression><softmax>
Title: How to do backpropagation with argmax?
UserDisplayName: Subrat Prasad
ThreadText: # How to do backpropagation with argmax?
  
  (Asked by: username_0 on 2023-07-13)
  
  I am attempting to utilize two networks: a classifier and a linear network. Based on the output class of the first network, my goal is to retrieve the corresponding value from the linear network using indexing. For instance, if the first network's output is [0.1, 0.8, 0.1], and the second network's output is [22, 30, 10], it implies that class two is selected, and the final output should be 30. Since this involves backpropagation through argmax, I am using the soft-argmax strategy described in the literature.
  Following is the code for soft-argmax
  class SoftArgmax1D(torch.nn.Module):
      def __init__():
          self.base_index = 0
          self.step_size = 1
          self.softmax = torch.nn.Softmax(dim=1)
  
      def forward(self, x):
          smax = self.softmax(x*100)
          end_index = self.base_index + x.size()[1] * self.step_size
          indices = torch.arange(start=self.base_index,
                                 end=end_index,
                                 step=self.step_size)
          return torch.matmul(smax, indices.float()) 
  
  Following is the final architecture:
  class CombinedNetwork(nn.Module):
      def __init__(self, input_size, hidden_size, num_classes):
          super(CombinedNetwork, self).__init__()
          self.classifier_fc1 = nn.Linear(input_size, hidden_size)
          self.classifier_fc2 = nn.Linear(hidden_size, num_classes)
          self.regression_fc = nn.Linear(input_size, num_classes)
          self.smax = SoftArgmax1D()
  
  
      def forward(self, x):
          # Classification branch
          classifier_x = F.relu(self.classifier_fc1(x))
          classifier_x = self.classifier_fc2(classifier_x)
          classification_value = self.smax(classifier_x)
          
          regression_value = self.regression_fc(x)
          
          reg_op = regression_value.gather(1, classification_value.round().unsqueeze(1)).squeeze(1)
          
          return reg_op
  
  The issue arises when indexing the regression_value requires an integer input. However, converting classification_value, which essentially represents the index, to an integer causes the backpropagation to be truncated.
  Do you have any suggestions on how to achieve backpropagation for this code?
  
  **Tags:** backpropagation, pytorch, regression, softmax
  > **Comments:**
  > **username_1**: Can you cast the value to `int` instead of rounding? Because `.round()` seems to be not differentiable  
  > **username_0**: Upon rechecking, it appears that using `.round()` is fine. However, converting it to an integer (using `.long()`) disrupts the flow in backpropagation.  
  > **username_1**: I've added a new answer. Please review it: if addresses your issue I'll delete the old one.  
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2023-07-16. Score: 0 | Accepted answer
  
  To back-propagate through non-differentiable operations, without designing a soft approximation for them, is possible to recur to a score gradient estimator:
  class CombinedNetwork(torch.nn.Module):
      def __init__(self, input_size, hidden_size, num_classes):
          super().__init__()
          self.classifier_fc1 = nn.Linear(input_size, hidden_size)
          self.classifier_fc2 = nn.Linear(hidden_size, num_classes)
          self.regression_fc = nn.Linear(input_size, num_classes)
  
  
      def forward(self, x):
          # Classification branch
          classifier_x = F.relu(self.classifier_fc1(x))
          logits = self.classifier_fc2(classifier_x)
  
          regression_value = self.regression_fc(x)
          return regression_value, logits
  
  # forward
  v, logits = model(x)
  
  # define Categorical distribution from predicted probabilities
  cat = torch.distributions.Categorical(logits=logits)
  indices = logits.argmax(-1, keepdims=True)
  
  # obtain values
  values = torch.gather(v, dim=1, index=indices)
  
  # setup score function, e.g. from MSE loss
  reward = (targets - values).square()
  loss = (-cat.log_prob(indices) * reward).mean()
  
  # optimization step
  ...
  
  Notice that I've assumes a squared loss between the targets and the predicted values: you can change it as you need. Once the model is trained, you can avoid defining the Categorical and directly index the value by according to the argmax of the logits.
  NOTE: training may be unstable, if so clip the gradients and rescale the targets (and so values).
  
  
  ---
  
  