Note on the quadratic penalties in elastic weight consolidation

Catastrophic forgetting is an undesired phenomenon which occurs when neural networks are trained on different tasks sequentially. Elastic weight consolidation (EWC; ref. 1), published in PNAS, is a novel algorithm designed to safeguard against this. Despite its satisfying simplicity, EWC is remarkably effective.

Motivated by Bayesian inference, EWC adds quadratic penalties to the loss function when learning a new task. The purpose of penalties is to approximate the loss surface from previous tasks. The authors derive the penalty for the two-task case and then extrapolate to handling multiple tasks. I believe, however, that the penalties for multiple tasks are applied inconsistently.

In ref. 1 a separate penalty is maintained for each task $T$, centered at ${\theta }_T^{\ast }$, the value of $\theta$ obtained after training on task $T$. When these penalties are combined (assuming ${\lambda }_T=\mathrm{1}$), the aggregate penalty is anchored at

$${\mu }_T={\left(F_A+F_B\dots +F_T\right)}^{-1}\left(F_A{\theta }_A^{\ast }+F_B{\theta }_B^{\ast }\dots +F_T{\theta }_T^{\ast }\right).$$

From the third task onward this is inconsistent with Bayesian inference. In the Bayesian paradigm the posterior $p\left(\theta \hspace{0.17em}|D_A,D_B\right)$ encapsulates the agent’s experience in both tasks $A$ and $B$, thus rendering the previous posterior $p\left(\theta \hspace{0.17em}|D_A\right)$ irrelevant. Analogously, as ${\theta }_B^{\ast }$ was obtained while incorporating the penalty around ${\theta }_A^{\ast }$, once we have ${\theta }_B^{\ast }$, ${\theta }_A^{\ast }$ is not needed anymore.

The correct form of penalties can be obtained by recursive application of the two-task derivation (2). It turns out that a single penalty is sufficient; its center is the latest optimum ${\theta }_T^{\ast }$ and its weights are given by the sum of diagonal Fisher information matrices from previous tasks $F_A+F_B+\mathrm{\dots }+F_T$. This single penalty version is akin to Bayesian online learning (3).

If the agent has to revisit data from past training episodes, multiple penalties should be maintained similarly to expectation propagation (4, 5). The aggregate penalty should be centered at ${\theta }_T^{\ast }$ rather than ${\mu }_T$. The anchor point for task $T$’s penalty should therefore be

$${\stackrel{\sim }{\theta }}_T=F_T^{-1}\left(\left(F_A+F_B+\mathrm{\dots }{F}_T\right){\theta }_T^{\ast }-F_A{\stackrel{\sim }{\theta }}_A-\mathrm{\dots }-F_S{\stackrel{\sim }{\theta }}_S\right)$$

rather than ${\theta }_T^{\ast }$, where tasks $A\mathrm{\dots }S$ precede task $T$, and ${\stackrel{\sim }{\theta }}_A\mathrm{\dots }{\stackrel{\sim }{\theta }}_S$ are the respective penalty centers for these tasks.

To illustrate the behavior and effect of different penalties I applied EWC to a sequence of linear regression tasks. The tasks define quadratic losses each with a diagonal Hessian (Fig. 1, Left). As all simplifying assumptions of EWC hold exactly, one should expect it to match exact Bayesian inference: Reach the global minimum of the combined loss and do so irrespective of the order in which tasks were presented. Fig. 1, Center shows that although a quadratic penalty could model task B perfectly the penalty is placed around the wrong anchor point. Table 1 shows that this leads to suboptimal performance and an unwanted sensitivity to task ordering. By contrast, EWC using the corrected penalties around ${\stackrel{\sim }{\theta }}_T$ models the losses perfectly (Fig. 1, Right) and achieves optimal performance in an order-agnostic fashion (Table 1).

Fig. 1.

(Left) Elliptical level sets of quadratic loss functions for tasks A, B, and C also used in Table 1. (Center) When learning task C via EWC, losses for tasks A and B are replaced by quadratic penalties around ${\theta }_A^{*}$ and ${\theta }_B^{*}$. (Right) Losses are approximated perfectly by the correct quadratic penalties around ${\stackrel{\sim }{\theta }}_A={\theta }_A^{\ast }$ and ${\stackrel{\sim }{\theta }}_B$.

Table 1.

Final performance of a model trained via different versions of EWC on a sequence of three linear regression tasks

Algorithm	Task order	Task $A$	Task $B$	Task $C$	Combined
EWC, original penalties (1)
	$\mathit{A}\to B\to C$	0.06880	0.09033	0.00042	0.15955
	$A\to C\to B$	0.07410	0.08036	0.00036	0.15481
	$B\to A\to C$	0.09033	0.06880	0.00042	0.15955
	$B\to C\to A$	0.08036	0.07410	0.00036	0.15481
	$C\to A\to B$	0.07981	0.07487	0.00014	0.15481
	$C\to B\to A$	0.07487	0.07981	0.00014	0.15481
EWC, corrected penalties (2)	Agnostic	0.07693	0.07693	0.00022	0.15407
Simultaneous training	Agnostic	0.07693	0.07693	0.00022	0.15407

${\lambda }_T$ is set to $1$ for all tasks. Using the penalties of ref. 1, the final performance is sensitive to the order in which the tasks were presented: The model incurs lower loss on tasks presented first, while performance on later tasks may degrade. Irrespective of task ordering the combined loss is suboptimal. Using the corrected penalties (2) EWC becomes order-agnostic and finds the same global minimum as simultaneous training, shown in the last row. Lowest values in each column are highlighted in boldface type.

I expect the negative impact of using incorrect penalties to be negligible until the network’s capacity begins to saturate. Furthermore, optimizing $\lambda$ can compensate for the effects to some degree. As a result, it is possible that using incorrect penalties would result in no observable performance degradation in practice. However, since the correct penalties are just as easy to compute, and are clearly superior in some cases, I see no reason against adopting them.