How to calculate FIM for EWC ?

[DragonRed18]

Hi everyone !
First of all, I would like to thank all contributors of avalanche for their work.
I was looking for the code of EWC strategy and I found two possible versions.

The first one was proposed here https://github.com/ContinualAI/colab/blob/master/notebooks/intro_to_continual_learning.ipynb.
The second is defined in avalanche code https://github.com/ContinualAI/avalanche/blob/master/avalanche/training/plugins/ewc.py.

In the first one all data is passed to the model.forward and loss.backward is applied (and optimizer.zero_grad is never called).
Since loss.backward is used the gradients are accumulated into the parameters of the model.
When there is no more data, gradients accumulated into parameters are used for fisher.

Instead in the second one, while data are forwarded the gradients of the current batch are saved in another variable, which accumulates gradients over time.

My question is if in your opinion these two versions can be considered equivalent and therefore they can be both used.

First version

def on_task_update(task_id, x_mem, t_mem):
      model.train()
      optimizer.zero_grad()
      for start in range(0, len(t_mem)-1, 256):
          end = start + 256
          x, y = torch.from_numpy(x_mem[start:end]), torch.from_numpy(t_mem[start:end]).long()
          x, y = x.to(device), y.to(device)
          output = model(x)
          loss = F.cross_entropy(output, y)
          loss.backward()
    
      fisher_dict[task_id] = {}
      optpar_dict[task_id] = {}
    
      for name, param in model.named_parameters():
        
        optpar_dict[task_id][name] = param.data.clone()
        fisher_dict[task_id][name] = param.grad.data.clone().pow(2)

Second version

def compute_importances(self, model, criterion, optimizer,
                        dataset, device, batch_size):


    model.eval()

    importances = zerolike_params_dict(model)
    dataloader = DataLoader(dataset, batch_size=batch_size)
    for i, (x, y, task_labels) in enumerate(dataloader):
        x, y = x.to(device), y.to(device)

        optimizer.zero_grad()
        out = avalanche_forward(model, x, task_labels)
        loss = criterion(out, y)
        loss.backward()

        for (k1, p), (k2, imp) in zip(model.named_parameters(),
                                      importances):
            assert (k1 == k2)
            if p.grad is not None:
                imp += p.grad.data.clone().pow(2)

    for _, imp in importances:
        imp /= float(len(dataloader))

    return importances

[AntonioCarta]

Avalanche does an average at the end:

for _, imp in importances:
    imp /= float(len(dataloader))

while the colab version does not. Apart from that, they are basically equivalent.

Notice that, to really compute the FIM, you should have batch_size=1. In practice, batch_size>1 works as well and it is much more efficient.

[DragonRed18]

Thank you for the answer.

I agree that there is not a practical difference between the two approaches.