F-MALLOC: Feed-forward Memory Allocation for Continual Learning in Neural Machine Translation

TL;DR

F-MALLOC is inspired by recent insights highlighting that feed-forward layers emulate neural memories and encapsulate crucial translation knowledge and decomposes feed-forward layers into discrete memory cells and allocates these memories to different tasks.

Abstract

In the evolving landscape of Neural Machine Translation (NMT), the pretrain-then-finetune paradigm has yielded impressive results. However, the persistent challenge of Catastrophic Forgetting (CF) remains a hurdle. While previous work has introduced Continual Learning (CL) methods to address CF, these approaches grapple with the delicate balance between avoiding forgetting and maintaining system extensibility. To address this, we propose a CL method, named $\textbf{F-MALLOC}$ ($\textbf{F}$eed-forward $\textbf{M}$emory $\textbf{ALLOC}ation)$. F-MALLOC is inspired by recent insights highlighting that feed-forward layers emulate neural memories and encapsulate crucial translation knowledge. It decomposes feed-forward layers into discrete memory cells and allocates these memories to different tasks. By learning to allocate and safeguard these memories, our method effectively alleviates CF while ensuring robust extendability. Besides, we propose a comprehensive assessment protocol for multi-stage CL of NMT systems. Experiments conducted following this new protocol showcase the superior performance of F-MALLOC, evidenced by higher BLEU scores and almost zero forgetting.

Figures (6)

• Figure 1: Illustration of F-MALLOC. For simplification, we depict a decomposed feed-forward layer. (a) The Original General Domain Model: Highlighting the general domain task in green. (b) Pruned General Domain Model: Post-pruning, pruned memories are 'writable' (depicted in white), while others are designated as 'read-only.' (c) Learning a New Task: The model learns to allocate some memories to the new task and mark them 'read-only' (depicted in yellow). 'read-only' memories remain available for future tasks' forward propagation. However, backward propagation through them is prohibited. (d) Multi-task Model: After learning all tasks, each task occupies a share of memory capacity. The forward pass of the last task is shown.
• Figure 2: Illustration of estimating feed-forward memory importance via JS divergence.
• Figure 3: Illustration of new domain learning: forward (the left) and backward (the right) propagation. Here, we show the inner structure of the feed-forward layer.
• Figure 4: Forgetting rate and saturation rate across different training stages.
• Figure 5: Feed-forward memory capacity usage in the training process of task sequence 0. Vertical dash lines indicate task switches.