I first did this rederivation when trying to understand EWC in 2019, when experimenting with continual learning at SDSUG AIML lab with Davit Soselia and others. No guarantees that everything is correct, but it mostly seems to check out.

• Data-generating distribution: The mixture of the underlying distributions of the dataset from the first task and the dataset from the second task.

• $W = w$: The event that our current parameters $w$ are optimal for our model's ability to model the data-generating distribution.

• $x$: The event that some data point or collection of data points $x$ has been generated from the data-generating distribution.

• $D_a$: The event that the data in some array $D_a$ has been generated from the data-generating distribution.

• $D_b$: The event that the data in some array $D_b$ has been generated from the data-generating distribution.

• $D$: The event that the data in some array $D$ (which is the union of the arrays $D_a$ and $D_b$) has been generated from the data-generating distribution. This can be expressed as:

1. We will use a flat prior for $P(W=w)$. This means that when we have seen no data, we assume that any value of $w$ is just as likely to be the optimal weights for our data-generating distribution as any other.

2. We will assume that the non-diagonal elements of the Hessian of the model's negative log-likelihood on the data are negligible. In other words, we assume that the variation in the loss with respect to any given parameter $w_i$ does not significantly depend on any other parameter $w_j$. Mathematically, this means that:

Intuitively, this assumption means that how the loss varies with any given parameter $w_i$ does not significantly change due to the influence of other parameters.

3. $D_a$ and $D_b$ are conditionally independent given $W=w$. This can be written as:

Intuitively, this assumption makes sense because the contents of $D_a$ and $D_b$ are dissimilar. Thus, knowing how likely it is for our model to have assigned a high probability to an item in $D_a$ gives us no information about how likely it is for our model to assign a high probability to an item in $D_b$.

• In supervised learning terms, if we imagine the datasets as key-value pairs (with samples as keys and labels as values), $D_a$ and $D_b$ will not share keys or have similar keys to each other.

4. All gradients and Hessians are with respect to $w$.

5. $N$ is the number of samples $x$ in $D_a$.

6. $M$ is the dimension of $w$.

Let us consider the probability of $W = w$ given $D$ , which represents the probability that our current weights are optimal, assuming that the data in $D$ came from the distribution we’re trying to model.

According to Bayes' Theorem:

Expanding the definition of $P(D)$, we get:

From our presupposition, we know that $D_a$ and $D_b$ are conditionally independent given $W = w$. Therefore,

Substituting this into the previous equation:

Rearranging the terms:

We know that:

Thus, we can rewrite the equation as:

We also know that:

Therefore, we get:

Next, we know that:

Thus, we can simplify it to:

For visibility, let’s take the logarithm of both sides:

Rearranging the terms:

• $\log P(D_b \mid W = w)$ is the log-likelihood of our model on dataset $D_b$. Thus, the probability of $W$ being optimal on $D$ depends on the log-likelihood of our model on dataset $D_b$.
• $\log P(W = w \mid D_a)$ is an intractable posterior term that represents the probability of $W$ being optimal given that $D_a$ has been generated from the data-generating distribution.
• Finally, we don't care about the $\log P(D_b \mid D_a)$ term because it is independent of and not conditioned on the probability of our weights being optimal, so it's a constant regardless of what parameters we have.

In other words, to maximize $P(W = w \mid D)$, we need to:

1. Maximize $P(D_b \mid W = w)$
2. Maximize $P(W = w \mid D_a)$

To achieve the first, we need to train the model with a negative log-likelihood loss on dataset $D_b$.

The second term, $P(W = w \mid D_a)$, is intractable. However, we can approximate it and instead maximize the approximation.

Let’s suppose that we have randomly initialized our model and then trained it with a negative log-likelihood loss on $D_a$. As a result, we have obtained weights $w^*$ such that:

Thus, we have:

Furthermore, we know that the array $D_a$ consists of some independent and identically distributed (iid) points $x$. So, we can express $P(D_a \mid W = w^*)$ as:

Taking the logarithm of both sides, we get:

The gradient of the sum can be written as:

Dividing by $N$, the number of samples, we get:

Thus, we have:

Next, we know that:

Thus, the negative log of $P(W = w \mid D_a)$ is:

Taking the gradient of both sides:

Since $\nabla \log P(D_a) = 0$, we get:

We also know that:

Thus, the negative log of $P(D_a \mid W = w)$ is:

Taking the gradient of both sides:

Since we are using a flat prior for $P(W = w)$, $\nabla \log P(W = w) = 0$. Thus, we get:

Therefore, we conclude that:

As such, we have:

Similarly, for the Hessian:

From Lemma 3 and Lemma 1, we get:

From Lemma 4, we get:

The Laplace approximation procedure here follows the one outlined in chapter 4.4 of Pattern Recognition and Machine Learning (Bishop).

Let’s approximate $\log P(W = w \mid D_a)$ using a second-order Taylor series expansion.

From Lemma 5, we know that $\nabla \log P(W = w^* \mid D_a) = 0$. Thus, the above expression simplifies to:

Let $A = -H(\log P(W = w^* \mid D_a))$. Then we can write:

Taking the exponential of both sides, we get:

Next, we approximate $P(W = w \mid D_a)$ as a Gaussian distribution:

Substituting back $A$, we get:

From Lemma 6, we know that:

Thus:

We know that $H(-\log P(x \mid W = w^*))$ is the observed Fisher information on $x$. Its expected value is the covariance of $\nabla (-\log P(x \mid W = w^*))$. Therefore:

Assuming that the non-diagonal elements of this covariance matrix are zero, we get:

The variance of $\nabla (-\log P(x \mid W = w^*))$ is given by:

From Lemma 2, we know that:

Thus, the variance simplifies to:

Therefore, we get:

Thus, the Laplace approximation for $P(W = w \mid D_a)$ becomes:

Let $H^* = \text{diag} \left( \sum_{x} \nabla (-\log P(x \mid W = w^*))^2 \right)$. Then, we have:

Let $C = \frac{\sqrt{\det(H^*)}}{(2\pi)^{M/2}}$. We can express the approximation as:

Taking the logarithm of both sides, we get:

Since $\log C$ is a constant, we can ignore it in the optimization process. Thus, to maximize $\log P(W = w \mid D_a)$, we need to minimize:

Returning to our expression for $\log P(W = w \mid D)$, the full expression is:

This expression is maximized by minimizing the following loss function:

The second term can be viewed as a quadratic regularization term. To control its influence on the training process, we introduce a regularization parameter $\lambda$:

(Note that in the original paper, $\lambda$ is introduced slightly differently. The diagonal precision matrix is divided by $N$, and our derivation with $\lambda = 1$ can be reobtained by setting $\lambda$ in the original expression to $N$).

To add L2 weight decay, we can include the decay term:

Let us introduce a new task $D_c$, and redefine $D$ as the union of arrays $D_a$, $D_b$, and $D_c$:

We have already derived an approximation for $\log P(W = w \mid D_a \cap D_b)$. Let $w_a^*$ be the optimal point found after training on $D_a$ (the same value as $w^*$ in the previous section).

This simplifies to:

After training on $D_a$ and $D_b$, the gradient of the above expression is zero. Let $w_b^*$ denote the weights at this stage. Now, we compute the second-order Taylor approximation of the above expression, term by term.

• 0th degree term:

• 1st degree term:

This term is zero because the gradient of the expression is zero at $w_b^*$.

• 2nd degree term:

  • The Hessian of the constant part with 0 gradient is zero.
  • The Hessian of $\frac{1}{2} (w - w_a^*)^T \text{diag} \left( \sum_{x \in D_a} \nabla (-\log P(x \mid W = w_a^*))^2 \right) (w - w_a^*)$ is $\text{diag} \left( \sum_{x \in D_a} \nabla (-\log P(x \mid W = w_a^*))^2 \right)$.
  • The Hessian of $\log P(D_b \mid W = w)$ is the same as the Hessian of $\log P(W = w \mid D_b)$ (Lemma 4). Thus, it can be approximated as $\text{diag} \left( \sum_{x \in D_b} \nabla (-\log P(x \mid W = w_b^*))^2 \right)$.

The complete second-degree term is:

Thus, the full second-degree Taylor approximation becomes:

Therefore:

Generalizing, let $D$ be the union of arrays $D_1, D_2, \dots, D_k$, and let $w_t^*$ be the optimum found after training on task $t$.

The loss becomes:

The EWC loss becomes:

Adding per-task $\lambda$'s and $\ell_2$ decay, the loss becomes:

As outlined by Ferenc Huszár, this is different from the loss recommended by the original paper for multiple tasks. In the original paper, a new quadratic penalty is added for each task:

$J(w)$ is more theoretically sound, but $J_{\text{original}}(w)$ is empirically informed. At the same time, $J(w)$ only requires storing the running sum of the Hessian approximations and the previous-task optimum, whereas $J_{\text{original}}(w)$ requires storing the Hessian and the optimum for each task.