Is max operation differentiable in Pytorch?

Viewed 2123

I am using Pytorch to training some neural networks. The part I am confused about is:

prediction = myNetwork(img_batch)
max_act = prediction.max(1)[0].sum()
loss = softcrossentropy_loss - alpha * max_act

In the above codes, "prediction" is the output tensor of "myNetwork". I hope to maximize the larget output of "prediction" over a batch.

For example: [[-1.2, 2.0, 5.0, 0.1, -1.5] [9.6, -1.1, 0.7, 4,3, 3.3]] For the first prediction vector, the 3rd element is the larget, while for the second vector, the 1st element is the largets. And I want to maximize "5.0+9.6", although we cannot know what index is the larget output for a new input data.

In fact, my training seems to be successful, because the "max_act" part was really increased, which is the desired behavior to me. However, I heard some discussion about whether max() operation is differentiable or not:

Some says, mathmatically, max() is not differentiable.
Some says, max() is just an identity function to select the largest element, and this largest element is differentiable.

So I got confused now, and I am worried if my idea of maximizing "max_act" is wrong from the beginning. Could someone provide some guidance if max() operation is differentiable in Pytorch?

1 Answers

max is differentiable with respect to the values, not the indices. It is perfectly valid in your application.

From the gradient point of view, d(max_value)/d(v) is 1 if max_value==v and 0 otherwise. You can consider it as a selector.

d(max_index)/d(v) is not really meaningful as it is a discontinuous function, with only 0 and undefined as possible gradients.

Related