Pytorch: no speedup when doing backward pass only for a part of the samples in minibatch - why?

Question

I have a function (see below) which modifies the loss function so that it returns only the loss for the K samples in the minibatch with the lowest loss. The idea is to focus in each optimization step on these samples.

So I do first a forward pass to get the loss value for each sample in the mini-batch, then adapt the loss via the fn. “get_adapted_loss_for_minibatch”.

As the adapted loss takes into account only a certain fraction of the samples in the minibatch (I am using currently 60% of the samples), I was expecting that I get also a measurable speedup during training, as the backward step has to be done only for a fraction of the samples in the minibatch.

But unfortunately this is not the case, the training takes practically the same amount of time as when I am using all samples in the minibatch (so when I do not adapt the loss). I am using a ‘densenet121’ network, and training is done on CIFAR-100.

Am I doing something wrong ? Should I disabled autograd for some samples in the minibatch manually ? I though the ‘topk’ function would do that automatically.

def get_adapted_loss_for_minibatch(loss):
    # Returns the loss containing only the samples of the mini-batch with the _lowest_ loss
    # Parameter 'loss' must be a vector containing the per-sample loss for all samples in the (original) minibatch
    minibatch_size = loss.size()[0]
    r = 0.6 * minibatch_size
    # round r to integer, safeguard if r is 0
    r = max(round(r), 1)
    # The 'topk' function returns the loss for the 'r' samples with the _lowest_ loss in the minibtach
    # See documentation at https://pytorch.org/docs/stable/generated/torch.topk.html
    # Note the 'topk' operation is differentiable, see https://stackoverflow.com/questions/67570529/derive-the-gradient-through-torch-topk
    # and https://math.stackexchange.com/questions/4146359/derivative-for-masked-matrix-hadamard-multiplication
    loss_adapted = torch.topk(loss, r, largest = False, sorted = False, dim = 0)[0]
    # return it
    return loss_adapted

Answer 1

The reason why you see no difference in the speed for training is that you are using batch normalization. In turn, this means your gradient still depends on the entirety of the batch, even if you are only using part of the batch's content to compute the final loss term and backpropagate.

Mathematically speaking the running statistics measured in each batch normalization layer will involve all elements in the batch.

If you look at the mean computation (of course batchnorms also involve standard deviation measurements). Intuitively when you normalize a given vector by its average, the resulting vector's elements will depend on all of the elements of the initial vector since all of them were used precisely to compute the average.

If you want to read more about this, you can read more on this post which is about backpropagating through x / x.mean(0).

---

Following our discussion, here is a way to replace your BatchNorm2d layers with GroupNorm. Go through the sub-modules of your network, then replace all instances of BatchNorm2d with new initialized instances of GroupNorm. I will give you an example with Bottleneck:

>>> net = Bottleneck(10, 2)

>>> for name, module in net.named_children():
...   if isinstance(module, nn.BatchNorm2d):
...     setattr(net, name, 
...        nn.GroupNorm(num_channels=module.num_features, num_groups=num_groups))

This way net will look like:

>>> net
Bottleneck(
  (bn1): GroupNorm(2, 10, eps=1e-05, affine=True)
  (conv1): Conv2d(10, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (bn2): GroupNorm(2, 8, eps=1e-05, affine=True)
  (conv2): Conv2d(8, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
)