Decouple knowledge from paramters for plug-and-play language modeling (2023)

Abstract: Pre-trained language models (PLM) have made impressive results in a wide range of NLP tasks and it has been revealed that one of the key factors to their success is the parameters of these models implicitly learn various types of knowledge in the pre-training corpus. However, encoding knowledge implicitly in the model parameters has two fundamental drawbacks. First, the knowledge is neither editable nor scalable once the model is trained, which is especially problematic in that knowledge is consistently evolving. Second, it lacks interpretability and prevents us from understanding what kind of knowledge PLM needs to solve a certain task. In this paper, we introduce {pasted macro ‘MODEL’}, a pre-training model with differentiable plug-in memory (DPM). The key intuition behind is to decouple the knowledge storage from model parameters with an editable and scalable key-value memory and leverage knowledge in an explainable manner by knowledge retrieval in the {pasted macro ‘MEMORY’}. We conduct extensive experiments under various settings to justify this design choice. In domain adaptation setting, {pasted macro ‘MODEL’} could be easily adapted to different domains with pluggable in-domain memory—obtaining 3.95 F1 improvements across four domains, without any in-domain training. {pasted macro ‘MODEL’} could also keep absorbing new knowledge after pre-training is done by knowledge updating operation in the {pasted macro ‘MEMORY’} without re-training. Finally, we show that by incorporating training samples into {pasted macro ‘MEMORY’} with knowledge prompting, {pasted macro ‘MODEL’} could further be improved by the instruction of in-task knowledge.

## introduction large pre-trained language models (plm) @xcite @xcite have become a revolutionary breakthrough in nlp area. optimized by carefully designed self-supervised objectives on unlabeled corpus and fine-tuned on downstream tasks, plms perform remarkably well in a wide range of nlp benchmarks. recent studies @xcite have revealed that one of the key factors to the success of plms is that the parameters of these models implicitly learn various types of knowledge in the pre-training corpus. owing to these learned syntactic, semantic, factual and commonsense knowledge, plms show great understanding, generalization and reasoning abilities in multiple downstream tasks @xcite . @xcite pointed out, the feed-forward layers (ffn), constituting two-thirds of a transformer model's parameters, are essentially key-value memories and store all kinds of knowledge of plm. the first linear layer of ffn acts like a set of sparsely activated keys detecting input patterns while the second is the corresponding value. to aggressively capture more knowledge, larger plms are continuously proposed, from 110m bert @xcite to 530b mt-nlg @xcite , yet plm has not reached upper bound @xcite . however, a fundamental question still remains: for plm, is it the optimal way to implicitly encode knowledge in its parameters? we argue that the implicit knowledge encoding approach has two fundamental drawbacks. first, the learned knowledge is neither editable nor scalable once the model is trained (e.g., bert doesn't know what is a bert). nevertheless, world knowledge is actually infinite and evolving. we thus would never expect an ever-large model to capture all the knowledge in its parameters and to be continuously retrained for the newly coming one. second, the current plms lack interpretability at the knowledge level. implicit knowledge encoding fails to provide provenance for model's prediction and makes plm a black box preventing humans from understand-ing which knowledge plm requires for a certain problem. in this work, we propose a novel architecture of plm, pluglm, which decouples the knowledge storage from model parameters and explicitly leverages the knowledge in an explainable manner. as shown in figure 1 , we balance the functionality of ffn layer with a differentiable plug-in key-value memory (dpm), which is highly scalable as well as editable. each slot of dpm encodes the knowledge to a pair of key and value, and thus we can explicitly retrieve the required knowledge in natural language from dpm rather than unnamed vectors in ffn. to justify the design choice of decoupling the knowledge from parameters, we conduct extensive evaluations under different settings. in the domain adaptation setting, pluglm could be easily adapted to different domains with pluggable indomain memory-obtaining 3.95 f1 improvements across four domains on average and up to 11.55 f1 improvement on acl-arc citation intent classification dataset, without any in-domain pre-training. in the knowledge update setting, pluglm could absorb new knowledge after pre-training is done in a training-free way by knowledge updating operation in the dpm, with an improvement up to 4 f1 scores in linnaeus ner dataset. pluglm could further be improved by incorporating training samples into dpm with knowledge prompting as a kind of in-task knowledge. ## preliminary feed-forward layers transformer @xcite , the backbone for all plms, is made of stacked self-attention (self-attn) and feedforward (ffn) layers. the former captures the contextual interaction among inputs and the latter process each input independently. let @xmath0, the ffn could be formulated as: ## pluglm the overall architecture of pluglm is illustrated in figure 1 . because ffn is essentially a key-value memory network @xcite @xcite , pluglm creatively decouples the knowledge storage from model parameters by replacing 2 ffn with a differential plug-in key-value memory, dpm ( §4.1) and conducting knowledge retrieval in dpm with knowledge attention ( §4.2) for explicit knowledge usage instead of storing all knowledge implicitly in the model parameters. in §4.3, we detailedly explain how pluglm is trained in both pre-training and finetuning stages. ## experiments pluglm mainly tries to decouple the knowledge storage from parameters and leverage knowledge in an explainable way. we conduct comprehensive experiments to show the superiority of this novel architecture: we could easily adapt the model to different domains without in-domain pre-training by switching dpm ( §5.1.1 and §5.1.2), alleviate catastrophic forgetting by storing dpm ( §5.1.1), inject new knowledge into the model by enlarging dpm ( §5.2), further enhance the model by injecting in-task knowledge into dpm ( §5.3) and unveil the black-box plm with direct access to the knowledge retrieved from dpm (appendix d). we also carefully examine each key design in pluglm and point the direction for future work in §5.4. ## conclusion for the first time, we challenge the current implicit knowledge encoding mechanism for plms with two fundamental drawbacks and insightfully propose to decouple knowledge storage from model parameters with an editable and scalable key-value memory. inspired by the findings that ffn stores all kinds of knowledge and is essentially a keyvalue memory network, we transform ffn archi-tecture into deep retrieval with a differentiable plugin memory (dpm), which makes the knowledge encoding of plms more flexible and interpretable. extensive experimental results in different scenarios including domain adaptation, knowledge update and in-task knowledge learning verify the design choice of pluglm. we believe this architectural design would pave a new direction for future research on plm, especially for super-large plm.