PEMA : An Offsite-Tunable Plug-in External Memory Adaptation for Language Models (2024)

Abstract: Pre-trained language models (PLMs) show impressive performance in various downstream NLP tasks. However, pre-training large language models demands substantial memory and training compute. Furthermore, due to the substantial resources required, many PLM weights are confidential. Consequently, users are compelled to share their data with model owners for fine-tuning specific tasks. To overcome the limitations, we introduce Plug-in External Memory Adaptation (PEMA), a Parameter-Efficient Fine-Tuning (PEFT) method, enabling PLM fine-tuning without requiring access to all the weights. PEMA integrates with context representations from test data during inference to perform downstream tasks. It uses external memory to store PLM-generated context representations mapped with target tokens. Our method utilizes weight matrices of LoRA-like bottlenecked adapter in the PLM’s final layer to enhance efficiency. Our approach also includes Gradual Unrolling, a novel interpolation strategy to improve generation quality. We validate PEMA’s effectiveness through experiments on syntactic and real datasets for machine translation and style transfer. Our findings show that PEMA outperforms other PEFT approaches in memory and latency efficiency for training, and also excels in maintaining sentence meaning and generating appropriate language and styles.

## introduction pre-trained language models (plms) are widely used in downstream nlp tasks @xcite . recent developments in large language models have shown remarkable performance in zero-shot and few-shot learning scenarios @xcite @xcite @xcite . however, fine-tuning is still required to optimize the performance of the nlp tasks such as machine translation @xcite @xcite . the most straightforward approach to fine-tuning is full finetuning @xcite , which involves fine-tuning all parameters in a plm. yet, this approach requires substantial resources regarding memory and training compute @xcite @xcite . to over-come this limitation, researchers have proposed parameter-efficient fine-tuning (peft) methods to fine-tune a full model efficiently. adapter tuning @xcite @xcite utilizes small, additional parameters known as adapters inserted between layers within a plm. on the other hand, lora @xcite uses trainable low-rank matrices that incrementally update the pre-trained weights. these fine-tuning methods require access to all the weights of plms. however, proprietary plms such as chat-gpt (openai, 2022), bard @xcite , and claude (anthropicai, 2023) are confidential. hence, the owners of these plms do not reveal all the model weights. consequently, data owners possessing their datasets and wishing to fine-tune proprietary plms for specific downstream tasks must provide their datasets to the plm owners for fine-tuning (openai, 2023a). however, this process can be challenging due to the confidential nature of the datasets, which may involve privacy concerns @xcite . figure 1a shows problems for fine-tuning proprietary plms. to overcome this situation, @xcite proposes the offsite-tuning approach that uses one-third of the middle layers of a plm, referred to as the emulator. nevertheless, this approach still needs a large parameter size, and compressing the full model into an emulator requires a computationally intensive distillation process. to address the challenges mentioned above, we introduce a novel peft method named plug-in external memory adaptation (pema) designed for efficient fine-tuning of proprietary plms in machine translation tasks. pema utilizes weight matrices of lora-like bottlenecked adapter designed for learning downstream tasks with accessible features provided by openai api @xcite and minimal part of plm's weight (language model head). in the training phase, the data owner begins the process by providing a prompt with initial input to the plm owner, which includes an instruction and a source sentence from a parallel corpus. the plm owner receives this initial input to generate a context representation (i.e., a hidden representation from plm) and predict the next-token. then, it iteratively processes subsequent inputs containing the predicted next-tokens. this approach avoids the need for the full dataset from the data owner. throughout this process, the data owner builds an external memory comprised of context representa-tions and corresponding desired target tokens. they train pema by reconstructing the stored context representations and predicting target tokens based on these representations. figure 1b shows the training phase process of pema. during the inference phase, the data owner uses a prompt to request a context representation for test data from the plm owner. the plm owner then outputs a context representation and a nexttoken probability given the prompt. pema also outputs a next-token probability based on a context representation. these probabilities are interpolated to compute a final next-token probability. we propose gradual unrolling (gu ), an interpolation strategy that initially emphasizes pema's distribution, gradually shifts to the plm's context-based predictions as the sentence progresses. figure 1c illustrates the inference phase process of pema. we evaluate pema by comparing it with other peft methods. pema shows better resource efficiency, consuming less gpu memory and running faster. additionally, pema outperforms other baselines in translating english sentences into german and paraphrasing informal sentences into formal ones while preserving the original meaning. lastly, we conduct ablation studies to assess the effectiveness of each component of pema. pema is publicly available for further exploration into offsite-tunable efficient fine-tuning. foot_0 2 related work ## plug-in external memory adaptation this section describes plug-in external memory adaptation (pema), which aims to fine-tune a plm without requiring a full model during training. pema integrates its output with that of the plm (i.e., next-token probability) during inference to facilitate downstream nlp tasks. at training, pema utilizes context representations of the plm and its lora-like bottlenecked adapter. for inference, pema requires context representation, the language model head (lm head) from the plm, and the lora-like bottlenecked adapter. it uses external memory to build a context representation f (c i ), mapped with the desired target token y i . using the external memory, we train pema in two phases. the first phase involves reconstruction training to reconstruct f (c i ) with b rct a, resulting in the output of a reconstruction loss. sub- lm head 𝑾 𝒉𝒅 pre-trained language model 𝒄 𝟏 [𝑰𝒏𝒔𝒕] i love you (𝒄 𝟐 ) 𝑰𝒏𝒔𝒕 i love you, ෝ 𝝎 𝟏 training context 𝒄 𝒊 representations 𝒇(𝒄 𝒊 ) (𝒚 𝟏 ) ich (𝒚 𝟐 ) liebe targets 𝒚 𝒊 𝑷 𝐋𝐌 ( ෝ 𝝎 𝟏 ) ❶ 𝒄 𝟏 [𝑰𝒏𝒔𝒕] i love you ❷ 𝒇(𝒄 𝟏 ) ❸ 𝒄 𝟐 [𝑰𝒏𝒔𝒕] i love you, ෝ 𝝎 𝟏 building external memory representations 𝒇(𝒄 𝒊 ) ich liebe targets 𝒚 𝒊 𝑷 𝐏𝐄𝐌𝐀 ( ෝ 𝝎 𝒊 ) target token representations 𝒉 𝒑𝒅 reconstructed representations 𝒉 𝒓𝒄𝒕 𝑨𝒇(𝒄 𝒊 ) 𝑩 𝒑𝒅 𝑩 𝒓𝒄𝒕 𝑨 𝓛 𝒕𝒐𝒕𝒂𝒍 = 𝜿𝓛 𝒓𝒄𝒕 +(𝟏 -𝜿)𝓛 𝒑𝒅 lm head 𝑾 𝒉𝒅 train [𝑰𝒏𝒔𝒕] hello dear colleagues [𝑰𝒏𝒔𝒕] hello dear colleagues test context 𝒙 gradual unrolling interpolation 𝝀𝑷 𝐏𝐄𝐌𝐀 ෝ 𝝎 + 𝟏 -𝝀 𝑷 𝐋𝐌 ( ෝ 𝝎) target token representations 𝒉 𝒑𝒅 𝑨𝒇(𝒄 𝒊 ) 𝑩 𝒑𝒅 𝑨 lm head 𝑾 𝒉𝒅 pre-trained language model inference final distribution of words 𝑷( ෝ 𝝎 𝒊 ) plm owner data owner 𝒇(𝒄 𝒊 ) 𝓛 𝒑𝒅 𝓛 𝒓𝒄𝒕 𝒇(𝒙) external memory pema 𝒇(𝒄 𝟐 ) 𝒇(𝒄 𝟏 ) 𝑬𝒒. 𝟓 𝑬𝒒. 𝟔 hallo predicted token 𝑨𝒇(𝒄 𝒊 ) 𝑩 𝒑𝒅 𝑨 @xcite given to a plm. the plm outputs the representation f (c i ) and predicts the next-token distribution p lm ( ŵi ). the representation f (c i ) is then aligned with its target y i . in the training phase, pema uses external memory for two tasks: preserving the original representation via reconstruction training with b rct and generating a target token probability distribution using b pd . for inference, the model inputs a test data representation to generate two probability distributions: p lm ( ŵi ) and p p em a ( ŵi ). these are then interpolated using gradual unrolling to obtain the final token distribution. ## gradual unrolling interpolation given that an adaptation model trained with only a limited number of parameters may lack the contextawareness and language-generation capabilities of pre-trained language models, it is more effective to use the adaptation model to guide the generation of tokens of the desired task at the beginning of the sentence, and rely on a pre-trained language model to provide context for the rest of the sentence. to achieve this, we suggest the gradual unrolling strategy, which aims for strong p p em a ( ŵ|x) interpolation at the beginning of generation and gradually decreases the interpolation. as the sentence progresses, the pre-trained language model increasingly contributes to providing the necessary context, as shown in figure 3 . in the context of sentence generation, we define sl as the input sentence length, excluding instruction and user-defined variables λ max . λ represents the proportion of the adaptation model's interpolation (0 ≤ λ ≤ 1). we also have the dependent variables of the current step (cs) and the step size (ss). the step size is computed as @xmath0, and cs is initialized to λ max at the start of sentence generation. at each token generation step, cs decreases by ss until the end of the sentence (i.e., cs @xmath1. then, we calculate the current interpolation proportion λ cur (i.e., λ at equation 6 ) as λ @xmath2. ## experiments this section describes the experiments and results to show both the computational efficiency and performance in downstream tasks of pema. first, we perform an experiment on the computational efficiency of pema. subsequently, we evaluate pema across two downstream tasks: the wmt22 @xmath3. lastly, we conduct an ablation study to show the gradual improvement by incorporating each idea of pema. ## conclusion in this paper, we present pema, a novel parameterefficient fine-tuning approach for language models. unlike existing peft methods, pema utilizes minimal pre-trained model parameters during training, making it an efficient and adaptable method for offsite-tuning. pema includes a token prediction decoder, gradual unrolling, and a reconstruction decoder to improve model performance. our comprehensive evaluations on translation and style transfer tasks demonstrate pema's effectiveness in generating text that more closely follows target domain distributions. additionally, pema proves its computational efficiency by utilizing minimal training memory and achieving faster training latency with a syntactic dataset. overall, pema offers efficient fine-tuning and presents a promising direction for an offsite-tunable peft approach in downstream nlp tasks.