1
$\begingroup$

I'm trying to re-implement Elastic Weight Consolidation (EWC) as outlined in this paper. As a reference, I am also using this Github repository (another implementation).

My model/idea is pretty straightforward. Train the network to do the bit operation AND (e.g 1 && 0 = 0), then using EWC, train it to use OR (e.g 1 || 0 = 1). I've got three inputs: bit1, bit2 and operation (0 stands for AND and 1 for OR) and one output neuron - the output of the operation. For example, if I have 0 1 0 the ground truth should be 0.

The problem, however, comes when calculating the EWC loss.

def penalty(self, model: nn.Module):
    loss = 0
    for n, p in model.named_parameters():
        _loss = self._precision_matrices[n] * (p - self._means[n]) ** 2
        loss += _loss.sum()
    return loss

I've got two problems:

  • The current means (p) and the old ones (self._means[n]) are always the same, resulting in multiplication by 0, which completely negates EWC.
  • As I have just one output neuron the calculation of the fisher's matrix is a bit different than the repo. The one I have written seems to be wrong. Any ideas?

I initialise the self._means[n] and self._precision_matrices (fisher's matrix) in the init method of the EWC model:

class EWC(object):
def __init__(self, model: nn.Module, dataset: list, device='cpu'):

    self.model = model
    self.dataset = dataset
    self.device = device

    self._means = {}
    self._precision_matrices = self._diag_fisher()

    for n, p in self.model.named_parameters():
        self._means[n] = p.data.clone()

def _diag_fisher(self):
    precision_matrices = {}

    # Set it to zero
    for n, p in self.model.named_parameters():
        params = p.clone().data.zero_()
        precision_matrices[n] = params

    self.model.eval()

    for input in self.dataset:
        input = input.to(self.device)

        self.model.zero_grad()

        output = self.model(input)
        label = torch.sigmoid(output).round()
        loss = F.binary_cross_entropy_with_logits(output, label)
        # loss = F.nll_loss(F.log_softmax(output, dim=1), label)
        loss.backward()

        for n, p in self.model.named_parameters():
            precision_matrices[n].data += p.grad.data ** 2 / len(self.dataset)

    precision_matrices = {n: p for n, p in precision_matrices.items()}
    return precision_matrices

And this is the actual training:

# Train the model EWC
for epoch in tqdm(range(EPOCS)):

    # Get the loss
    ls = ewc_train(model, opt, loss_func, dataloader[task], EWC(model, old_tasks), importance, device)

def ewc_train(model: nn.Module, opt: torch.optim, loss_func:torch.nn, data_loader: torch.utils.data.DataLoader, ewc: EWC, importance: float, device):
    epoch_loss = 0

    for i, (inputs, labels) in enumerate(data_loader):
        inputs = inputs.to(device).long()
        labels = labels.to(device).float()

        opt.zero_grad()

        output = model(inputs)
        loss = loss_func(output.view(-1), labels) + importance * ewc.penalty(model)
        loss.backward()
        opt.step()

        epoch_loss += loss.item()

    return loss

Note: the loss function that I am using is nn.BCEWithLogitsLoss() and optimisation is: SGD(params=model.parameters(), lr=0.001).

$\endgroup$
1
  • $\begingroup$ Hi Mitch and welcome to this community! Maybe it would be better if you provided a minimal and reproducible example. Bear in mind that not all implementation-related questions are on-topic here. $\endgroup$ – nbro Dec 22 '19 at 14:27

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.