Why mini-batches larger than 1 doesn't work, but larger accumulating gradients work?

Viewed 61

I am trying to implement a neural network approximating the logical XOR function, however, the network only converges when using a batch size of 1.

I don't understand why: when I use gradient accumulation with multiple mini-batches of size 1, the convergence is very smooth, but mini-batches of size 2 or more don't work at all.

This issue arises, whatever the learning rate, and I have the same issue with another problem(more complex) than XOR.

I join my code for reference:

import numpy as np
import torch.nn as nn
import torch
import torch.optim as optim
import copy

#very simple network
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(2,3,True)
        self.fc1 = nn.Linear(3,1, True)

    def forward(self, x):
        x = torch.sigmoid(self.fc(x))
        x = self.fc1(x)
        return x

def data(n): # return n sets of random XOR inputs and output
    inputs = np.random.randint(0,2,2*n)
    inputs = np.reshape(inputs,(-1,2))
    outputs = np.logical_xor(inputs[:,0], inputs[:,1])
    return torch.tensor(inputs, dtype = torch.float32),torch.tensor(outputs, dtype = torch.float32)


N = 4
net = Net() # first network, is updated with minibatches of size N
net1 = copy.deepcopy(net) # second network, updated with N minibatches of size 1
inputs = torch.tensor([[0,0],[0,1],[1,0],[1,1]], dtype = torch.float32)
labels = torch.tensor([0,1,1,0], dtype = torch.float32)
optimizer = optim.SGD(net.parameters(), lr=0.01)
optimizer1 = optim.SGD(net1.parameters(), lr=0.01)
running_loss = 0
running_loss1 = 0
for epoch in range(25000):  # loop over the dataset multiple times
    # get the inputs; data is a list of [inputs, labels]
    input, labels = data(N)

    # zero the parameter gradients
    optimizer.zero_grad()
    optimizer1.zero_grad()
    # forward + backward + optimize
    loss1_total = 0
    for i in range(N):
        outputs1 = net1(input[i])
        loss1 = (outputs1-labels[i]).pow(2)/N # I divide by N to get the effective mean
        loss1.backward()
        loss1_total += loss1.item()


    outputs = net(input)
    loss = (outputs-labels).pow(2).mean()
    loss.backward()
    
    # optimization
    optimizer.step()
    optimizer1.step()

    # print statistics
    running_loss += loss.item()
    running_loss1 += loss1_total
    if epoch % 1000 == 999:    # print every 1000 mini-batches
        print(f'[{epoch + 1},  loss: {running_loss/1000 :.3f}, loss1: {running_loss1/1000 :.3f}')
        running_loss1 = 0.0
        running_loss = 0.0
        
print('Finished Training')
 # exemples of data and outputs for reference ; network 2 always converge to the sub-optimal point(0.5,0.5)
datatest = data(4)
outputs = net(datatest[0])
outputs1 = net1(datatest[0])
inputs = datatest[0]
labels = datatest[1]
print("input",inputs)
print("target",labels)
print("net output",outputs)
print("net output",outputs1)

[EDIT] Improved readability and updated the code

result :

[1000,  loss: 0.259, loss1: 0.258
[2000,  loss: 0.252, loss1: 0.251
[3000,  loss: 0.251, loss1: 0.250
[4000,  loss: 0.252, loss1: 0.250
[5000,  loss: 0.251, loss1: 0.249
[6000,  loss: 0.251, loss1: 0.247
[7000,  loss: 0.252, loss1: 0.246
[8000,  loss: 0.251, loss1: 0.244
[9000,  loss: 0.252, loss1: 0.241
[10000,  loss: 0.251, loss1: 0.236
[11000,  loss: 0.252, loss1: 0.230
[12000,  loss: 0.252, loss1: 0.221
[13000,  loss: 0.250, loss1: 0.208
[14000,  loss: 0.251, loss1: 0.193
[15000,  loss: 0.251, loss1: 0.175
[16000,  loss: 0.251, loss1: 0.152
[17000,  loss: 0.252, loss1: 0.127
[18000,  loss: 0.251, loss1: 0.099
[19000,  loss: 0.251, loss1: 0.071
[20000,  loss: 0.251, loss1: 0.048
[21000,  loss: 0.251, loss1: 0.029
[22000,  loss: 0.251, loss1: 0.016
[23000,  loss: 0.250, loss1: 0.008
[24000,  loss: 0.251, loss1: 0.004
[25000,  loss: 0.251, loss1: 0.002

Finished Training

input tensor([[1., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.]])
target tensor([1., 0., 0., 0.])
net output tensor([[0.4686],
        [0.4472],
        [0.4472],
        [0.4472]], grad_fn=<AddmmBackward0>)
net1 output tensor([[0.9665],
        [0.0193],
        [0.0193],
        [0.0193]], grad_fn=<AddmmBackward0>)

Please, could you explain me why this strange phenomena is appearing ? I searched for a long time on the net, without success...

Excuse me if my question is not well formatted, it is the first time I ask a question on stack overflow.

EDIT : I found, comparing accumulated gradients of size 1 minibatches and gradients from minibatches of size N, that the computed gradients are mostly the same, only small(but noticeable) differences appear probably due to approximation errors, so my implementation looks fine at first sight. I still don't get where does this strong convergence property of minibatches of size 1 come from.

2 Answers

The problem lies in the way you define labels / compute the loss in

 loss = (outputs-labels).pow(2).mean()

We have labels.shape = [4] but outputs.shape =[4, 1]. This due to the broadcasting, the difference

(outputs - labels).shape = [4, 4]

which means we compute all pairwise differences between outputs and labels (and then take their 2nd power and average them), which basically means that no meaningful supervision takes place.

The quick way to fix that here would be adding a dummy dimension here:

loss = (outputs-labels[:, None]).pow(2).mean()

but the clean way would be doing it the correct way right from the start, that is defining your labels in a ways such that labels.shape = [_, 1]:

labels = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32)

(and similar in your data() function).

It seems there is a minor issue with the dimensions of labels and outputs.

This:

labels = torch.tensor([0,1,1,0], dtype = torch.float32)

Needs to become this:

labels = torch.tensor([[0],[1],[1],[0]], dtype = torch.float32)

Otherwise, the mismatch between the model output and the labels messes up the loss in the minibatch example.

This can be fixed in data(n), if you add an extra dimension to outputs:

outputs = np.logical_xor(inputs[:,0], inputs[:,1]).reshape((n, 1))

After fixing that, there will be a floating-point underflow issue as well. The gradient accumulation method divides then sums gradients, but in the minibatch method first sums then divides the values. Mathematically they are the same, but in practice, there will be drift between them in long run.

Check this example:

x = np.array([0.00649802, 0.24420964, 0.05081264,])
(x/3).sum() - x.mean()
# -1.3877787807814457e-17
Related