On Robust Prefix-Tuning for Text Classification

1 INTRODUCTION . Large-scale pretrained language models ( LMs ) ( Peters et al. , 2018 ; Devlin et al. , 2019 ; Radford et al. , 2019 ; Liu et al. , 2019 ; Yang et al. , 2019 ; Raffel et al. , 2020 ; Lewis et al. , 2020 ; Brown et al. , 2020 ; Xue et al. , 2021 ) have proven effective for downstream NLP tasks . While finetuning a pretrained model for a specific task has been the common practice , it comes at the cost of maintaining a full copy of the LM with the parameters entirely modified . The prohibitively huge memory demand poses a severe challenge for the deployment of practical NLP systems , which motivates the development of low-storage adaptation methods ( Houlsby et al. , 2019 ; Li & Liang , 2021 ) . Recently , increasing interest has been focused on prompt-based tuning approaches for pretrained language models ( Wallace et al. , 2019 ; Puri & Catanzaro , 2019 ; Shin et al. , 2020 ; Jiang et al. , 2020b ; Zhong et al. , 2021 ; Gao et al. , 2021 ; Hu et al. , 2021 ; Liu et al. , 2021 ) . By prepending several elaborately-selected tokens to the given input sequences , the LM is triggered to respond with appropriate outputs without updating its parameters . Prefix-tuning ( Li & Liang , 2021 ) introduces the idea of replacing the discrete prompt tokens at the input with the virtual ones at the start of each layer in the LM . By optimizing the layerwise continuous prefix embedding instead of selecting candidates in the vocabulary list , the expressive ability of prompts is further enhanced with a rather small amount of parameters to be updated . As a result , prefix-tuning requires near 1000× fewer task-specific parameters than finetuning the entire pretrained model ( Bommasani et al. , 2021 ) . Despite being lightweight and modular , prefix-tuning is still lacking in robustness . In the NLP community , a variety of techniques for generating adversarial examples have been proposed to attack a text classifier by perturbing inputs ( Zhang et al. , 2020 ) . Conventional attack techniques include character-level ( Eger et al. , 2019 ; He et al. , 2021 ) , word-level ( Alzantot et al. , 2018 ; Ren et al. , 2019 ; Garg & Ramakrishnan , 2020 ) , sentence-level modification ( Iyyer et al. , 2018 ; Ribeiro et al. , 2018 ; Xu et al. , 2021 ) , or a mixture of them ( Ebrahimi et al. , 2018 ; Li et al. , 2019 ) . Instead of perturbing each input sentence separately , recently , universal adversarial triggers ( UAT ) ( Wallace et al. , 2019 ) becomes powerful by prepending the same adversarial tokens to all test inputs . UAT prompts the model to generate malicious outputs , which shares the same spirit with the promptbased tuning approaches . It remains a mystery whether prefix-tuning , a variant of prompt-based tuning techniques , can defend against UAT as well as other different kinds of attacking techniques . In defense of adversarial attacks , different types of defense techniques are developed , including model functional improvement ( Li & Sethy , 2019 ; Jones et al. , 2020 ) , certification ( Jia et al. , 2019 ; Huang et al. , 2019 ; Shi et al. , 2020 ; Xu et al. , 2020 ; Ye et al. , 2020 ) , adversary detection ( Pruthi et al. , 2019 ; Zhou et al. , 2019 ) , and adversarial training ( Miyato et al. , 2017 ; 2019 ; Zhu et al. , 2020 ; Jiang et al. , 2020a ; Liu et al. , 2020 ; Wang et al. , 2021 ; Dong et al. , 2021 ; Zhou et al. , 2021 ) . While these approaches have enhanced model robustness , difficulties emerge when fitted to prefix-tuning . Most of the techniques require modification to the architecture and the parameters of the LM or additional maintenance of adversary detectors . Directly applying such techniques necessitates auxiliary model update and storage , which will inevitably hamper the modularity of prefix-tuning . Moreover , The excessively long time for adversarial training is also a hindrance to the efficient use of prefix-tuning . We ask the following question : Can we improve the robustness of prefix-tuning while preserving its efficiency and modularity , without modifying the pretrained model parameters ? In this work , we propose a robust prefix-tuning framework for text classification . The main idea of our framework is to add an extra batch-level prefix tuned for each batch to the original prefix embedding during test time for robustness enhancement . We first record the layerwise activations in the LM at the position of generating label prediction with correctly classified training data . We project the collected activation matrices of each layer onto low-level canonical manifolds as the charac- terization of “ correct ” model behavior . In this way , the correctness of any layerwise activations at the position of prediction generation can be estimated by projecting to the canonical manifolds and measuring the distance between them . For each test batch during inference , the added extra prefix is tuned on the fly with the original prefix fixed to minimize the calculated distance . Triggered by the summed prefix , the LM is prone to generating correct label predictions . We conduct extensive experiments on three text classification benchmarks and show that the proposed framework substantially improves model robustness against five strong textual attack approaches including input perturbation attack of different levels as well as the UAT attack . To the best of our knowledge , we are the first to propose the defense approach for prefix-tuning while keeping its lightweightness and modularity . Moreover , we provide an interpretation of our robust prefix-tuning framework from the optimal control perspective and pose several directions for future research . 2 PREFIX-TUNING FOR TEXT CLASSIFICATION . Prefix-tuning is a lightweight alternative to finetuning when using large-scale pretrained language models to solve downstream NLP tasks . The intuition of prefix-tuning follows prompt-based methods that a proper context prepended to input sentences triggers the desired response of the LM without changing the large amount of LM parameters . Instead of instantiating the prepended context with discrete tokens , prefix-tuning uses trainable prefix embeddings as a replacement , which is also known as soft prompts . The continuous prefix embeddings enable continuous optimization and are prepended to all Transformer layers to improve expressiveness . Following the notation of Li & Liang ( 2021 ) , the activation at the i-th position of the j-th layer in an L-layer autoregressive Transformer LM is denoted as h ( j ) i . hi = [ h ( 0 ) i ; · · · ; h ( L−1 ) i ] represents the stacked activations : hi = { Pθ [ i , : ] , if i ∈ Pidx , LMϕ ( zi , h < i ) , otherwise . ( 1 ) where Pidx is the sequence of prefix indices and zi is the i-th token in the input sequence . The activations of the first |Pidx| positions are directly calculated by Pθ . All of the activations at the following positions depend on the prefix as the autoregressive LM follows the left-to-right calculation process . To stabilize the optimization , the prefix embedding matrix Pθ is reparameterized as Pθ [ i , : ] = MLPθ ( P̂θ [ i , : ] ) by a feedforward network MLPθ with a smaller matrix P̂θ . While prefix-tuning is proposed for conditional generation tasks , in this work , we use prefix-tuning for text classification . As shown in Figure 1 , following the protocol of decaNLP ( McCann et al. , 2018 ) , we frame the samples in classification tasks into a SQuAD-like scheme consisting of context , question , and label : the context and the label part refer to the text sequence to be classified and the ground-truth label , while the question part is a prescribed task description sentence fixed for all samples . We denote x = [ context , question , [ ANS ] ] , where [ ANS ] is a special token that separates question and label . We let y = [ label ] and |y| = 1 as the label is one token . At the position that [ ANS ] is inputted , the LM generates the prediction of the next label token , and we denote this position as the output position o . While o can be different for different input x ’ s , in this paper , we omit the relation o = o ( x ) for simplicity . Prefix-tuning aims to steer the LM to maximize the probability of the label . We use all samples in the training setDtr to optimize the prefix Pθ [ i , : ] . The objective is min θ E ( x , y ) ∼DtrL ( y|x ; θ ) = max θ E ( x , y ) ∼Dtr log [ W ( h ( L ) o ) ] y , ( 2 ) where W in the LM transforms the top-layer output h ( L ) o to a probability vector over the vocabulary . With continuous optimization on training samples , prefix-tuning is expected to steer the LM to generate correct label predictions for test data . With the large-scale LM parameters fixed , the obtained task-specific prefix is lightweight and modular . However , prefix-tuning is still vulnerable to text attacks . With the context part perturbed by text attack techniques , the LM can be fooled to generate erroneous label prediction at the output position . Figure 1 shows an example of perturbation : by modifying a single character m in the word remember with k , the prediction of the LM is shifted from positive to negative . Therefore , it remains under exploration how to robustify prefix-tuning without hampering its modularity or introducing additional large model updates and storage . 3 ROBUST PREFIX-TUNING . We propose a robust prefix-tuning framework for text classification . Our intuition follows prefixtuning that proper prefix embeddings prepended to inputs can steer a LM with correct responses . When the inputs are adversarially perturbed , the LM activations at the output position fail to be steered in the correct way by the original prefix Pθ [ i , : ] . Inspired by Khoury & Hadfield-Menell ( 2018 ) that the perturbed data often deviates from the low-dimensional data manifold , our robust prefix-tuning framework uses the layerwise activations by correctly classified training data to construct canonical manifoldsM . When provided with perturbed inputs during inference , we add an extra prefix P ′ψ [ i , : ] tuned for each test batch to Pθ [ i , : ] that aims to rectify the erroneous activations at the output position so that they stay close to the canonical manifolds . In this way , we expect the summed prefix to steer the LM with correct label generation against input perturbations . As shown in Figure 1 , our robust prefix-tuning framework consists of three steps . The first step is collecting correct LM activations at the output position o triggered by Pθ [ i , : ] . We denote SC as the set of correctly classified training examples . For the j-th layer , the collected activation matrix H ( j ) C stacks the j-th layer LM activation at the output position o with the input of all c ∈ SC : H ( j ) C = [ h ( j ) o , c ] ∈ R|SC |×d . ( 3 ) The d represents the dimension of the LM hidden state . In practice , we always have |SC | > > d. The second step is constructing canonical manifolds . We project the collected j-th layer activation matrixH ( j ) C onto a low-level manifoldM ( j ) as the characterization of the correct j-th layer behavior . We use PCA ( Pearson , 1901 ) to get the projectionQ ( j ) onto the canonical manifold of the j-th layer : H̃ ( j ) C = U ( j ) Σ ( j ) V ( j ) T , ( 4 ) Q ( j ) = V ( j ) p T V ( j ) p , ( 5 ) where H̃ ( j ) C = H ( j ) C −11 TH ( j ) C /|SC | normalizes the rows ofH ( j ) C to mitigate the randomness among samples in SC before projection . V ( j ) p consists of the first p singular vectors and Q ( j ) ∈ Rd×d . The third step is tuning P ′ψ [ i , : ] to robustify prefix-tuning during inference . Here the vector P ′ ψ [ i , : ] is not reparameterized by MLP . With the additional prefix P ′ψ [ i , : ] , the token-wise activations become hi = { Pθ [ i , : ] + P ′ ψ [ i , : ] , if i ∈ Pidx , LMϕ ( zi , h < i ) , otherwise . ( 6 ) For the j-th layer at the output position , the LM activation matrix triggered by Pθ [ i , : ] +P ′ψ [ i , : ] with the potentially perturbed test input batch ST is stacked as H ( j ) T = [ h ( j ) o , t ] ∈ R|ST |×d ( 7 ) for all t ∈ ST . We use the distance from H ( j ) T to the j-th canonical manifoldM ( j ) as the loss for the tuning of P ′ψ [ i , : ] for each batch . Projecting H ( j ) T toM ( j ) yields H ( j ) T Q ( j ) , thus the objective is min ψ N−1∑ j=0 L ( ψ ( j ) , Q ( j ) ) = N−1∑ j=0 ∥∥∥H ( j ) T ( I −Q ( j ) ) ∥∥∥ 2 . ( 8 ) We also replace the H ( j ) T in Eq . ( 8 ) with H ( j ) T − 11 TH ( j ) T /|ST | as normalization before projection to mitigate randomness among test samples when |ST | > 1 . After tuning P ′ψ [ i , : ] , the activated H ( j ) T is closer to M ( j ) . As the manifold characterizes the correct behavior of the j-th layer activation , by regulating the layerwise activations at the output position , the summed prefix Pθ [ i , : ] +P ′ψ [ i , : ] is prone to steering the LM to generate correct label predictions . Our framework is also applicable to other soft prompt-based tuning methods ( Qin & Eisner , 2021 ; Hambardzumyan et al. , 2021 ; Lester et al. , 2021 ; Cho et al. , 2021 ; Tsimpoukelli et al. , 2021 ) by recording the activations of correctly classified training data , constructing canonical manifolds for the soft prompts , and tuning additional soft prompts for robustness during inference . In this work , we conduct experiments on prefix-tuning . Remark . From the optimal control ( OC ) perspective , prefix-tuning can be formalized as seeking the OC of the pretrained LM for downstream tasks , and our robust prefix-tuning can be interpreted as seeking the close-loop control for robust downstream tasks . We attach the details in Appendix G .

Summary:  This paper introduces a tweak to Prefix-Tuning to make it more resilient to adversarial perturbations of the input. The idea is to add a batch-level prefix at inference to the original one which enhances robustness. Critically, Robust Prefix-Tuning (RPT) does not require auxiliary model updates or storage, in contrast with other robustness methods.Thus, this approach makes prefix-tuning more robust while preserving its modularity and low storage requirements.  The authors conduct experiments on 3 text classification tasks, 5 textual attacks and different training regimes (normal training, adversarial training and adversarial data augmentation). In nearly all instances, their method improves robustness (sometimes considerably so) while preserving the accuracy on the original text.    The authors also present RPT from an optimal control perspective and conduct a qualitative study that shows how RPT impacts attention weights.