0
$\begingroup$

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?

$\endgroup$
3
  • $\begingroup$ Can you cast the value to int instead of rounding? Because .round() seems to be not differentiable $\endgroup$ Jul 13, 2023 at 17:35
  • $\begingroup$ Upon rechecking, it appears that using .round() is fine. However, converting it to an integer (using .long()) disrupts the flow in backpropagation. $\endgroup$ Jul 14, 2023 at 9:37
  • $\begingroup$ I've added a new answer. Please review it: if addresses your issue I'll delete the old one. $\endgroup$ Jul 16, 2023 at 16:07

1 Answer 1

0
$\begingroup$

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

$\endgroup$
1
  • 1
    $\begingroup$ This works. Thanks! $\endgroup$ Jul 17, 2023 at 5:40

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .