# How to do backpropagation with argmax?

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)


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?

• Can you cast the value to int instead of rounding? Because .round() seems to be not differentiable Jul 13, 2023 at 17:35
• Upon rechecking, it appears that using .round() is fine. However, converting it to an integer (using .long()) disrupts the flow in backpropagation. Jul 14, 2023 at 9:37

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).

• This works. Thanks! Jul 17, 2023 at 5:40