Differences between gradient calculated by different reduction methods in PyTorch

Question

I'm playing with different reduction methods provided in built-in loss functions. In particular, I would like to compare the following.

1. The averaged gradient by performing backward pass for each loss value calculated with reduction="none"

2. The gradient averaged by dividing the batch size with reduction="sum"

3. The average gradient yielded by reduction="mean"

4. The average gradient calculated by reduction="mean", with the data points fed into the model one at a time.

My code for producing the experiment is as follows:

def estimate_gradient(model, optimizer, batch):
    criterion_no_reduction = nn.CrossEntropyLoss(reduction="none").cuda()
    criterion_sum = nn.CrossEntropyLoss(reduction="sum").cuda()
    criterion_avg = nn.CrossEntropyLoss().cuda()

    input, target = batch
    input, target = input.cuda(), target.cuda()
    output = model(input)
    n = len(output)

    loss_no_reudction = criterion_no_reduction(output, target)
    grad_list_no_reduction = []
    for i in range(n):
        optimizer.zero_grad()
        loss_no_reudction[i].backward(retain_graph=True)
        for j, param in enumerate(model.parameters()):
            if param.requires_grad:
                grad = param.grad.view(-1, 1)
                if i == 0:
                    grad_list_no_reduction.append(grad)
                else:
                    grad_list_no_reduction[j] = torch.cat((grad_list_no_reduction[j], grad), dim=1)
    grad_out_no_reduction = torch.cat(grad_list_no_reduction, dim=0)
    grad_out_no_reduction = (torch.sum(grad_out_no_reduction, dim=1) / n).cpu().detach().numpy().flatten()

    loss_sum = criterion_sum(output, target)
    optimizer.zero_grad()
    loss_sum.backward(retain_graph=True)
    for j, param in enumerate(model.parameters()):
        if param.requires_grad:
            if j == 0:
                grad_list_sum = param.grad.view(-1)
            else:
                grad_list_sum = torch.cat((grad_list_sum, param.grad.view(-1)))
    grad_out_sum = (grad_list_sum / n).cpu().detach().numpy().flatten()

    loss_avg = criterion_avg(output, target)
    optimizer.zero_grad()
    loss_avg.backward(retain_graph=True)
    for j, param in enumerate(model.parameters()):
        if param.requires_grad:
            if j == 0:
                grad_list_avg = param.grad.view(-1)
            else:
                grad_list_avg = torch.cat((grad_list_avg, param.grad.view(-1)))
    grad_out_avg = grad_list_avg.cpu().detach().numpy().flatten()

    target = target.view(-1, 1)
    grad_list_one_by_one = []
    for i in range(n):
        optimizer.zero_grad()
        curr_output = output[i].view(1, -1)
        loss = criterion_avg(curr_output, target[i])
        loss.backward(retain_graph=True)
        for j, param in enumerate(model.parameters()):
            if param.requires_grad:
                grad = param.grad.view(-1, 1)
                if i == 0:
                    grad_list_one_by_one.append(grad)
                else:
                    grad_list_one_by_one[j] = torch.cat((grad_list_one_by_one[j], grad), dim=1)
    grad_out_one_by_one = torch.cat(grad_list_one_by_one, dim=0)
    grad_out_one_by_one = (torch.sum(grad_out_one_by_one, dim=1) / n).cpu().detach().numpy().flatten()

    assert grad_out_no_reduction.shape == grad_out_sum.shape == grad_out_avg.shape == grad_out_one_by_one.shape
    print("Maximum discrepancy between reduction = none and sum: {}".format(np.max(np.abs(grad_out_no_reduction - grad_out_sum))))
    print("Maximum discrepancy between reduction = none and avg: {}".format(np.max(np.abs(grad_out_no_reduction - grad_out_avg))))
    print("Maximum discrepancy between reduction = none and one-by-one: {}".format(np.max(np.abs(grad_out_no_reduction - grad_out_one_by_one))))
    print("Maximum discrepancy between reduction = sum and avg: {}".format(np.max(np.abs(grad_out_sum - grad_out_avg))))
    print("Maximum discrepancy between reduction = sum and one-by-one: {}".format(np.max(np.abs(grad_out_sum - grad_out_one_by_one))))
    print("Maximum discrepancy between reduction = avg and one-by-one: {}".format(np.max(np.abs(grad_out_avg- grad_out_one_by_one))))

The results are as follows:

Maximum discrepancy between reduction = none and sum: 0.0316
Maximum discrepancy between reduction = none and avg: 0.0316
Maximum discrepancy between reduction = none and one-by-one: 0.0
Maximum discrepancy between reduction = sum and avg: 0.0
Maximum discrepancy between reduction = sum and one-by-one: 0.0316
Maximum discrepancy between reduction = avg and one-by-one: 0.0316

That is, the result produced by reduction=none and one-by-one backward pass appear to be identical, while reduciton=sum and reduction=mean yields different results from the previous pair. It would be really helpful to explain the discrepancy (maybe due to retain_graph=True?) and thanks in advance for any help!

Answer 1

Let's start by just recalling what each of these means. Reduction 'none' means compute batch_size gradient updates independently for the loss with respect to each input in the batch and then apply (the composition of) them. Reduction 'mean' and 'sum' mean apply the respective operations and the take the gradient with respect to this one value.

Now, let's look at the different comparisons.

None vs. one-by-one

These are the same because by doing one-by-one, you're taking the mean of a 1-item list which is just the value in the list. So the loss will always be identical. Put another way, you're basically just re-creating the 'none' reduction in your code by applying the average to one loss output at a time.

Sum vs. average

Let's next compare sum and average, since this will make it easier to explain none vs. the two of them. To start, let's think about how we expect the gradient updates to differ between 'sum' and 'avg' reduced loss outputs. To be explicit, we know that average is

$$ \frac{1}{n} \sum_{i=1}^n \mathrm{loss} _i $$

whereas sum is $$ \sum_{i=1}^n \mathrm{loss}_i $$ So regardless of the actual formula for our gradients, we expect that the gradient of the 'sum' reduced loss will always equal $ n $ times the gradient of the 'mean' reduced loss. This means that all weights will be updated in the same proportion just with different overall magnitudes.

However, you're seeing that the gradients are the same between these two reductions. Why is that? Well, it turns out, looking at your code, you're dividing the 'sum' reduced gradient by $ n $, so actually it's exactly what we'd expect to see.

As an aside, there's a related question of why both tend to mostly produce the same result. I believe the answer is that smart optimizers like Adam's learning rate tuning will account for the constant difference between the two and factor it into their learning rate choices. I will also note that one nice thing about the 'mean' reduction is that it keeps your gradients constant regardless of your batch size (for the same data).

Sum/avg vs. one-by-one/none

Now that we've explained why the two pairs give the same results, let's understand why the results differ between them. The best guess (hat tip: Ken Arnold in the comments) is that it's the result of your final batch having a different (smaller) size than the rest of your batches and this producing different results between mean and and sum.