Paper Note: Overcoming Catastrophic Forgetting with Hard Attention to the Task

Overcoming Catastrophic Forgetting with Hard Attention to the Task

Joan Serra, Didac Suris, Marius Miron, Alexandros Karatzoglou

ICML 2018

[paper] [code]

Related Work

文章评价了一下 rehearsal（Rebuffi et al., 2017; Lopez-Paz & Ranzato, 2017）和 pseudo-rehearsal（Venkatesan et al., 2017; Shin et al., 2017; Nguyen et al., 2017）的方法：rehearsal 的方法对存储资源有要求；pseudo-rehearsal 的方法一般是使用生成网络（generative network），而为一系列任务的数据训练生成网络其实并不容易。而且这两种方法都意味着某种形式上的 concurrent learning，也就是说需要在学习新的任务时去对“旧”的数据再进行处理。

相比于 EWC (Kirkpatrick et al., 2017) 和 SI (Zenke et al., 2017) 给损失函数加了一个 ‘soft’ structural 正则项，HAT 则是施加了一个 ‘hard’ structural 正则——既在损失函数上施加，也在 gradient magnitudes 上显式添加。EWC是在模型训练之后衡量权重的重要性，而SI和HAT是在网络训练的同时计算权重的重要性。EWC 和 SI 使用 specific formulation，而HAT是学习 attention masks 。IMM (Lee et al., 2017) 是对EWC的改进，在学习完一个新任务之后，有一个独立的 model-merging 的步骤。

PNN (Rusu et al., 2016) 按列（column-wise）分配网络权重，为每个任务预先分配列宽（column width）。他们使用 adapters 来重用以前的列/任务中的知识，分配给将来任务的权重数量逐渐增加。HAT不是盲目地预先分配列宽，而是学习每层的“宽度”以及网络权重，并使它们适应当前任务的难度。PathNet (Fernando et al., 2017) 还为每个任务预先分配了一定数量的网络容量，但与PNN相比，它避免了使用 network columns 和 adapters。它使用一种进化的方法来学习在相互连接的恒定数量的 modules（layer subsets）之间的路径。HAT不会维护大量 solutions，而是完全采用反向传播和SGD进行训练，并且不依赖于一组恒定数量的 modules。

与PNN和PathNet相似，PackNet（Mallya＆Lazebnik，2017）采用了 binary mask 来约束网络。 但是，这种约束不是基于列或层模块，而是基于网络权重。因此，它可以更好地利用网络的容量。PackNet是基于启发式权重剪枝，并具有预先指定的剪枝率。HAT也聚焦于网络权重，但是使用 unit-based 的 masks 来限制它们，从而也得到了一种轻量级的结构。尽管它使用 compressibility parameter $c$ 影响所学习模型的紧凑性，但是它避免了任何绝对或预先指定的裁剪率。HAT与前面三种方法之间的另一个区别是它不使用纯 binary masks。相反，stability parameter s_\max 控制着二值化程度。

DEN (Yoon et al., 2018) 也根据手头的任务分配网络容量。但是，他们是在称为“selective retraining”的阶段单独进行的。启发式方法和超参数的复杂混合被用来识别“漂移”了的单元，这些单元在另一个阶段被复制并重新训练。$L_1$-regularization 和 $L_2$-transfer 与相应的正则化常数和其他阈值一起用于限制学习。HAT追求简单性，将超参数的数量限制为两个具有很直接的概念解释的参数。HAT 放弃了对于网络权重的普通 $L_1$ 正则化，而是采用了对于 attention masks 的 attention-weighted $L_1$ 正则化。注意 masks 是一种轻巧的结构，可以在无需对现有网络进行重大更改的情况下加入。

Main Idea

本文提出了一种基于任务的硬注意力机制（hard attention mechanism），该机制可维护先前任务中的信息，而不会影响新任务的学习。在学习一项任务的同时，我们还通过 gated task embeddings 来学习 almost-binary attention vectors。先前任务的注意力向量用于定义 mask，以此限制在当前任务上更新网络的某些权重。由于 mask 几乎是二值的，因此一部分权重会保持不变，而其余部分则会去适应新的任务。我们称我们的方法为 hard attention to the task (HAT)。

Motivation

提出HAT想法的出发点在于：任务定义（task definition），或更实际地说，任务的标识符（task identifier）对于神经网络的运行至关重要。考虑在训练区分鸟和狗图像的任务时，网络可能会学习到一些中间特征。如果第二项任务是使用相同的数据来区分棕色和黑色动物（假设数据只包含棕色或黑色的鸟和狗），则网络可能会学习到一组新的特征，而这些特征与第一组特征可能几乎没有太多交集。因此，如果两个任务中的训练数据相同，则一个重要的区别就是它们的任务描述或标识符。我们的目的是学习使用任务标识符来对每一层神经网络进行条件限制，并在以后使用这种学习到的条件限制来防止忘记先前的任务。

Architecture

为了对任务 t 施加条件限制，我们使用一种 layer-wise 的注意力机制（如 Fig. 1 所示）。当给出 $l$ 层 units （在本文中，units 既指 linear units（即全连接层的神经元），又指卷积层的 filters）的输出 $\boldsymbol{h}_l$ ，我们对其施以 element-wise 的乘法 $\boldsymbol{h}'_l=\boldsymbol{a}_l^t\odot\boldsymbol{h}_l$ 。$\boldsymbol{a}_l^t$ 是一个单层的 task embedding $\boldsymbol{e}_l^t$ 的 gated version：

$\boldsymbol{a}_l^t=\sigma(s\boldsymbol{e}_l^t)\qquad\qquad(1)$

其中 $\sigma(x)\in[0,1]$ 是一个门函数，$s$ 是一个正的的调幅参数。Eq. 1 对 $l=1,...,L-1$ 的所有层都适用。而对于 L 层，也就是网络的最后一层，$\boldsymbol{a}_L^t$ 是二值硬编码的。

Figure 1. 示意图：上半部分是前向传播，下半部分是后向传播

Eq. 1所表示的门机制是想要形成 hard, possibly binary attention masks，以此激活或者抑活每一层units的输出。用这样的方式，我们可以动态地创建或者销毁层与层之间的 paths，这些 paths 可以在以后学习新任务时保留下来。给定某种网络结构，HAT 会学习并自动确定 paths，这最终会影响每一层的权重。 此外，HAT 使用反向传播和SGD同时学习 paths 和网络的其余部分。

Network Training

为了在学习新任务时保留先前任务中学习到的信息，我们根据由先前所有任务所得的累积注意力（cumulative attention）来限制梯度。为了获得一个累积注意力向量，在学习完任务 $t$ 并获得 $\boldsymbol{a}_l^t$ 之后，我们使用 element-wise maximum 递归地计算：

$\boldsymbol{a}_l^{\le t}=\max\Big(\boldsymbol{a}_l^{t}, \boldsymbol{a}_l^{\le t-1}\Big)$

其中 $\boldsymbol{a}_l^{\le 0}$ 为零向量。这样就保留了对于之前任务而言重要的 units 的注意力值，使之能够对将来任务的训练起到限制作用。

为了限制对任务 t+1 的训练，我们使用下面的式子对梯度 $g_{l,ij}$ 进行修改：

$g_{l,ij}'=\Big[1-\min\Big(a_{l,i}^{\le t},a_{l-1,j}^{\le t}\Big)\Big] g_{l,ij}\qquad(2)$

其中 unit indices $i$ 和 $j$ 分别与输出层（$l$）和输入层（$l-1$）相对应。换言之，我们扩展了向量 $\boldsymbol{a}_l^{\le t}$ 和 $\boldsymbol{a}_{l-1}^{\le t}$，使之与 $l$ 层的梯度向量的维度匹配，然后进行 element-wise 的 minimum、减法和乘法（如 Fig. 1 所示）。如果输入数据是诸如图像或音频之类的复杂信号，我们不会对输入数据计算 attention。但是，在数据由单独或独立的特征组成的情况下，也可以将它们视为某一层的输出，并应用如上所述的方法。

请注意，我们使用 Eq. 2 创建 masks 来防止对于过去任务很重要的权重有过大的更新。这类似于 PackNet 的方法（Mallya＆Lazebnik，2017），在PackNet中，在进行启发式选择和重新训练之后，找到了 binary mask，随后使用该 mask 来冻结相应的网络权重。HAT在三个重要方面与PackNet不同。(1) 我们的 mask 是unit-based 的，weight-based 的 mask 会自动从中得出。因此，HAT 也存储和维护了一个轻量级结构。(2) 我们的 mask 是学习得到的，而不是启发式或规则驱动的。因此，HAT不需要预先分配压缩比，也不需要通过后训练步骤（post-training step）来确定参数的重要性。(3) 我们的 mask 不一定是二值的，允许是介于 0 到 1 之间的中间值。如果我们想以一些遗忘为代价，重新使用权重来学习其他任务，或者我们希望以更在线的方式工作——忘记最旧的任务而记住新的任务，那么这将很有用。

Hard Attention Training

为了获得一个完全二值的注意力向量 $\boldsymbol{a}_l^t$ ，我们可以使用 unit step function 作为 gate。但是，我们希望可以用后向传播来训练 embedding $\boldsymbol{e}_l^t$ （Fig. 1），所以我们更希望有一个可微的函数作为 gate。为了构建一个 pseudo-step function 以使得梯度流通，我们使用一个带有正标量参数 $s$ 的 sigmoid 函数（Eq. 1）。引入此缩放比例是为了控制 pseudo-step function 的极化（polarization）或“硬度（hardness）”及其结果输出 $\boldsymbol{a}_l^t$。我们的方法是在训练阶段 anneal $s$，产生梯度流，在测试阶段置 s=s_\max 。当使用 s_\max\gg 1 时，Eq. 1近似于一个 unit step function。注意当 $s\rarr \infin$ 时，我们有 $a_{l,i}^t\rarr\{0,1\}$ ；当 $s\rarr 0$ 时，我们有 $a_{l,i}^t\rarr 1/2$ 。我们将使用后者来开始一个 training epoch，所有网络单元均处于同等的活跃状态，并在该 epoch 中逐步将它们极化。

在一个 training epoch 中，我们线性递增 $s$ 的值：

$s=\cfrac{1}{s_{max}}+\Big(s_{max}-\cfrac{1}{s_{max}}\Big)\cfrac{b-1}{B-1}\qquad(3)$

论文中有对 Eq. 3 的一些阐释，但都比较直觉，这里就不展开了。

Embedding Gradient Compensation

论文发现 embedding $\boldsymbol{e}_l^t$ 的变化不大，并且这些权重的梯度幅值很小。论文认为问题的主要原因是由于引入了 annealing scheme（Eq. 3）。这一部分是论文提出的既能增大 $e_{l,i}^t$ 范围，又能增大其梯度幅值的方法，不详说了。

Promoting Low Capacity Usage

需要注意，hard attention values $a_{l,i}^t$ 处于活跃状态，也就是 $a_{l,i}^t\rarr 1$ ，直接决定了所对应的 units 将被分配给任务 $t$。所以，为了给将来的任务保留尽可能多的模型容量，我们需要对 attention vectors $A^t=\{\boldsymbol{a}_1^t,...,\boldsymbol{a}_{L-1}^t\}$ 进行稀疏化。为此，我们给损失函数 $\mathcal{L}$ 增加一个正则项，将至任务 $t-1$ 的累积注意力向量的集合 $A^{\lt t}=\{\boldsymbol{a}_1^{\lt t},...,\boldsymbol{a}_{L-1}^{\lt t}\}$ 也考虑进去：

$\mathcal{L}'(\boldsymbol{y},\hat{\boldsymbol{y}},A^t,A^{\lt t})=\mathcal{L}(\boldsymbol{y},\hat{\boldsymbol{y}})+cR(A^t,A^{\lt t}),\qquad\qquad(4)$

$R(A^t,A^{\lt t})=\cfrac {\sum_{l=1}^{L-1}\sum_{i=1}^{N_l}a_{l,i}^t\Big(1-a_{l,i}^{\lt t}\Big)} {\sum_{l=1}^{L-1}\sum_{i=1}^{N_l}1-a_{l,i}^{\lt t}} \qquad\qquad(5)$

Eq. 5 即正则项，$N_l$ 对应于第 $l$ 层的 units 的数目。注意 Eq. 5 是对 $A^t$ 的加权的 $L_1$ 正则。对过去任务的累积注意力 $A^{\lt t}$ 为当前任务定义了一个权重，也就是说当 $a_{l,t}^{\lt t}\rarr 1$，$a_{l,i}^t$ 得到一个趋近于0的权重。这就将参加了之前任务的 units 排除在正则之外，不会限制它们在当前任务中的复用。

角声寒，夜阑珊。怕人寻问，咽泪装欢。