Source-linked AI summary
Continual Pre-training of Language Models
Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, Bing Liu
TL;DR
Continual DAP-training asks how to adapt a language model across sequential unlabeled domains without losing prior or general knowledge. DAS uses unit-level soft-masking, a robustness-based importance proxy, and contrastive integration; experiments report improved average performance and knowledge transfer.
Problem
Continual DAP-training must adapt a language model to emerging domains while preventing catastrophic forgetting of previously learned domains and general language knowledge.
Method
DAS uses importance-based soft-masking, a proxy for general-knowledge importance, and complementary representations to preserve and integrate knowledge during continual DAP-training.
Results
DAS outperforms all baselines on average and achieves the best knowledge transfer, reflected by a negative forgetting rate.
Takeaways & Limitations
Continual DAP-training can retain prior knowledge while transferring knowledge across domains, improving end-task performance within the evaluated setting.
Takeaways & Limitations
The approach uses attention heads as its unit of importance, and reported results can exhibit high variance because of very small test sets.
Abstract
from arXiv · showhide
Language models (LMs) have been instrumental for the rapid advance of natural language processing. This paper studies continual pre-training of LMs, in particular, continual domain-adaptive pre-training (or continual DAP-training). Existing research has shown that further pre-training an LM using a domain corpus to adapt the LM to the domain can improve the end-task performance in the domain. This paper proposes a novel method to continually DAP-train an LM with a sequence of unlabeled domain corpora to adapt the LM to these domains to improve their end-task performances. The key novelty of our method is a soft-masking mechanism that directly controls the update to the LM. A novel proxy is also proposed to preserve the general knowledge in the original LM. Additionally, it contrasts the representations of the previously learned domain knowledge (including the general knowledge in the pre-trained LM) and the knowledge from the current full network to achieve knowledge integration. The method not only overcomes catastrophic forgetting, but also achieves knowledge transfer to improve end-task performances. Empirical evaluation demonstrates the effectiveness of the proposed method.
1 INTRODUCTION
The paper formulates continual domain-adaptive pre-training as updating a language model across sequential unlabeled domains while retaining prior and general knowledge. DAS addresses this problem with soft-masking, a robustness-based importance proxy, and knowledge integration.
- Problem: Continual DAP-training incrementally adapts a pre-trained language model to sequential unlabeled domain corpora whose data becomes unavailable after training.End-task fine-tuning evaluates the resulting model on each domain.
- Problem: The system must prevent catastrophic forgetting of both previously learned domains and the general language knowledge in the original model.The paper also identifies knowledge transfer across domains as a desired property.
- Method: DAS uses soft-masking to assign importance values between 0 and 1 to units and control their backward gradient updates without masking the forward pass.The full network remains available for end-task fine-tuning and cross-domain transfer.
- Method: DAS computes unit importance for general knowledge before continual learning using a robustness-based proxy that does not access the original pre-training data.Its continual-learning function then trains each domain while addressing forgetting and transfer.
- Contributions: The contributions combine continual DAP-training, soft-masking, contrastive knowledge integration, and a proxy for preserving general knowledge.The paper reports experimental evidence for DAS’s effectiveness.
2 RELATED WORK
Prior work covers domain adaptation, continual learning, pruning, and contrastive representation learning, but does not directly solve continual DAP-training with full-network knowledge integration. DAS differs by updating the language model while controlling forgetting at the unit level.
- DAP-training: Domain adaptation can directly update a language model or train adapters and prompts, but DAS belongs to the direct-update family.The paper notes that continual learning is especially challenging for direct updates because of catastrophic forgetting.
- Continual learning: Existing continual-learning methods address catastrophic forgetting through approaches such as regularization, while DAS controls attention heads or neurons rather than individual parameters.The paper contrasts this unit-level control with EWC.
- Continual learning: NLP continual-learning research has covered several end-task settings, but the cited work does not address domain-adaptive pre-training.Some language-model systems use separate adapters, which limits their knowledge transfer and forgetting properties.
- Neural network pruning: Neural-network pruning methods are not directly applicable because continual DAP-training must preserve both individual domain knowledge and general knowledge.The paper frames this as different from simply discarding unimportant parameters.
- Contrastive learning: DAS uses contrastive learning to integrate representations by contrasting previous-domain and general knowledge with knowledge from the full network.The paper states that contrasting only domain-specific knowledge can split rather than integrate the representations.
3 PROPOSED DAS TECHNIQUE
DAS continually adapts a language model to new domains by preserving important general and previously learned knowledge while encouraging complementary representations for knowledge integration. It combines importance-based soft-masking, a robustness proxy for initialization, and contrastive learning during domain training.
- Continual DAP-training: DAS preserves important general and domain knowledge with soft-masks while leaving the forward pass unrestricted to facilitate cross-domain knowledge transfer.Soft-masks control backward gradient flow using accumulated unit importance, but are not applied during the forward pass.
- Overall procedure: DAS consists of initialization before continual learning, followed by domain training and importance computation for each new domain.Initialization estimates importance for general language knowledge; continual learning uses accumulated importance and current-domain data, then computes current-domain importance.
- Importance computation: Importance is computed for Transformer attention heads and neurons using virtual parameters whose gradients estimate each unit’s importance without updating network parameters.The method treats multi-head attention, intermediate, and output layers similarly; importance values correspond to individual units.
- Initialization: Because original pre-training data are inaccessible, DAS uses model robustness as a proxy for estimating importance to general knowledge.The proxy feeds the same unlabeled input through the model twice with different dropout masks and uses the divergence between representations.
- Knowledge integration: DAS adds a contrastive loss that encourages complementary representations between accumulated knowledge and the full network, alongside the masked-language-modeling loss.The final DAP-training objective combines MLM loss after soft-masking with the contrastive loss; the full network represents both prior and current-domain knowledge.
4 EXPERIMENTS
Experiments evaluate DAS on six domain corpora and corresponding end-task datasets against non-continual and continual baselines. DAS generally outperforms the baselines, combining forgetting prevention, knowledge transfer, and robust adaptation of the full language model.
- Experimental Setup: The evaluation uses RoBERTa, six unlabeled domain corpora, corresponding end-task classification datasets, and 16 adapted baselines.The corpora cover restaurant, phone, and camera reviews plus ACL, AI, and PubMed papers; final models are evaluated after all domains are incrementally learned.
- Main Results: DAS outperforms all baselines on average and achieves the best knowledge transfer, measured by a negative forgetting rate.The main results report end-task macro-F1, accuracy, and forgetting rate averaged over five random seeds.
- Main Results: DAS achieves both forgetting prevention and knowledge transfer, whereas baselines emphasizing only one objective sacrifice accuracy or obtain weaker transfer.KD, EWC, and DER++ focus on forgetting prevention, while BCL, CLASSIC, and DEMIX provide weaker knowledge transfer than DAS.
- Ablation Study: Directly updating the full LM outperforms adapter- and prompt-based methods, while full-network learning is more effective than HAT-style sub-networks.The paper attributes the adapter and prompt gap to limited randomly initialized trainable parameters and reports poor performance for HAT and HAT-All.
- Ablation Study: The KL-divergence proxy performs similarly to DAS and better than the Wiki-plus-MLM variant for estimating general-knowledge importance.The proxy uses dropout-based robustness, which is less tied to a specific domain than adapting with MLM on Wikipedia.
5 CONCLUSION
DAS is presented as a method for continual DAP-training that preserves prior knowledge, estimates unit importance, and supports knowledge integration. The approach uses attention heads as the relevant Transformer units.
- 5 CONCLUSION: DAS preserves important previous knowledge through soft-masking units according to their importance, addressing catastrophic forgetting and facilitating knowledge transfer.Its other key ideas are estimating general-knowledge importance with a proxy and learning complementary representations for knowledge integration.
- 5 CONCLUSION: DAS uses a novel proxy to compute the importance of units associated with general knowledge in the language model.
- 5 CONCLUSION: The approach learns complementary representations for knowledge integration and was evaluated extensively, with experiments showing DAS’s effectiveness.
- 5 CONCLUSION: The authors identify combining the two learning functions as future work to further improve results.
- 5 CONCLUSION: The method uses attention heads as units because they are considered an important Transformer component.
A DATASETS DETAILS
The evaluation covers four types of end-tasks across product reviews, scientific citations, artificial intelligence relations, and biomedical chemical-protein interactions.
- A DATASETS DETAILS: Aspect sentiment classification uses an aspect and review sentence to classify positive, negative, or neutral sentiment, with Phone and Camera limited to positive and negative labels.
- A DATASETS DETAILS: Citation intent classification assigns a citing sentence one of five citation functions, including background, motivation, uses, extension, and comparison or contrast future.
- A DATASETS DETAILS: Relation classification labels a within-sentence span containing two entities with one of seven relations, such as feature of, used for, or part of.
- A DATASETS DETAILS: Chemical-protein interaction classification assigns spans containing chemical-protein pairs to interaction categories including activator, inhibitor, upregulator, and antagonist.
B BASELINE DETAILS
The baselines span independent task models, non-continual DAP-training approaches, shared or expanding adapters, regularization, replay, distillation, and parameter isolation.
- BASELINE DETAILS: Non-continual baselines build a separate model for each task independently, avoiding both knowledge transfer and catastrophic forgetting.
- BASELINE DETAILS: RoBERTa performs end-task fine-tuning without DAP-training, whereas DAP-RoBERTa applies masked language model loss for domain-adaptive pre-training.
- BASELINE DETAILS: Adapter-tuning adds two-layer fully connected bottleneck modules while freezing the Transformer during DAP-training; prompt-tuning instead trains virtual tokens.The adapter bottleneck size is 128, while prompt-tuning initializes 100 tokens with learning rate 0.3.
- BASELINE DETAILS: Naive continual learning incrementally trains one model across domains with masked language model loss and no mechanism for catastrophic forgetting.NCL-Adapter shares one adapter set across all domains rather than adding a new adapter for each domain.
- BASELINE DETAILS: DEMIX adds a new adapter for each domain, while BCL uses capsule networks and CLASSIC uses contrastive learning to support continual learning.
- BASELINE DETAILS: HAT isolates parameters with task embeddings and requires task information during end-task fine-tuning, including when applied to Transformer or adapter layers.
- BASELINE DETAILS: KD uses representational KL-divergence distillation, EWC regularizes parameter changes with L2 penalties, and DER++ replays stored domain tokens with distillation.DER++ stores 16.4K tokens per learned domain, the largest memory usable for the system to run.
C IMPLEMENTATION DETAILS
The implementation uses RoBERTaBASE with masked-language-model DAP-training, fixed sequence-length and optimizer settings, and task-specific fine-tuning schedules. Additional tables summarize averaged performance and metric variability.
- C IMPLEMENTATION DETAILS: RoBERTaBASE is the backbone, with a masked language model head used for DAP-training and standard RoBERTa end-task fine-tuning.For the three aspect sentiment tasks, the aspect and review sentence are concatenated using </s>.
- C IMPLEMENTATION DETAILS: The maximum input length is 164, and Adam is used for both DAP-training and end-task fine-tuning.
- C IMPLEMENTATION DETAILS: DAP-training uses learning rate 1e-4, batch size 256, and 2.5K steps per domain, approximately one full pass through domain data.
- C IMPLEMENTATION DETAILS: The importance-computation subset contains 1.64 Million tokens, with λ set to 1 and τ set to 0.05.
- C IMPLEMENTATION DETAILS: End-task fine-tuning uses learning rate 1e-5 and batch size 16, with 5, 10, or 15 epochs depending on the dataset.The schedules are 5 epochs for Restaurant, 10 for ACL, AI, and PubMed, and 15 for Phone and Camera.
- C IMPLEMENTATION DETAILS: Table 5 reports DAS performance averaged across all domains after final DAP-training over 5 random seeds, while Table 6 reports standard deviations for DAS and baselines.
D DAP-TRAINING IN DIFFERENT ORDERS
Across five domain orders, DAS consistently outperforms NCL, indicating that its effectiveness is not tied to one particular ordering.
- DAS is better than NCL across all five evaluated domain orders.The orders vary the sequence of Restaurant, ACL, AI, Phone, PubMed, and Camera domains.
E STANDARD DEVIATIONS
Standard-deviation results over five random-seed runs indicate that DAS and its variants are stable, while some baselines show substantially greater variability.
- DAS results are stable over five runs with random seeds.The corresponding standard deviations are reported for the main-paper results.
- Some baselines exhibit quite large standard deviations, including RoBERTa in AI and DAP-RoBERTa in Camera.
- DAS and its variants also show stable results for the metrics corresponding to Table 4.Table 7 reports these standard deviations for the proposed model and ablation variants.