Selective Token Generation for Few-shot Language Modeling

1 INTRODUCTION . Natural language processing ( NLP ) have recently achieved great progress using advanced neural language models ( Radford et al. , 2018 ; Devlin et al. , 2019 ; Yang et al. , 2019 ; Liu et al. , 2019 ; Clark et al. , 2020 ; Raffel et al. , 2019 ; Lan et al. , 2019 ; Lewis et al. , 2020 ) . However , these neural models typically require large-scale training data for each individual task , and solving a new NLP task that has only a few examples is still challenging problem ( Yin , 2020 ) . Especially , natural language generation ( NLG ) with limited training data is an important yet more difficult task due to its fast adaptation of sequential prediction models in a wide range of applications including text summarization , question answering , data-to-text generation , machine translation , etc ( Peng et al. , 2020 ; Chen et al. , 2020 ; Xu et al. , 2021 ; Schick & Schütze , 2020 ; Chang et al. , 2021 ; Radford et al. , 2019 ; Lewis et al. , 2020 ; Brown et al. , 2020 ) . More recently , pretrained language models ( PLMs ) have shown great generalization ability when combined with large-scale data and big transformer-based models ( Devlin et al. , 2019 ; Radford et al. , 2019 ; Lewis et al. , 2020 ; Brown et al. , 2020 ; Subramanyam Kalyan et al. , 2021 ) . Therefore , transfer learning from transformer PLMs has been popularly used for few-shot NLG tasks with promising results . In specific , the use of PLM for few-shot NLG can be categorized into three approaches : 1 ) prompt-based , 2 ) finetuning , and 3 ) additive learning . Prompt-based approaches encode a task description and task-specific examples as a natural language prompt for few-shot text generation ( Radford et al. , 2019 ; Brown et al. , 2020 ; Zheng & Huang , 2021 ; Schick & Schütze , 2020 ; Li & Liang , 2021 ) . While these approaches can take full advantage of the universal natural language understanding and generation capabilities of large-scale PLMs without further training of the main model , these have some limitations in dealing with a large domain shift from the pretraining corpus data , tuning suitable task-specific prompts , and covering an increased size of conditioning examples . On the other hand , finetuning of the PLM is able to explicitly impart task-specific knowledge to the model and hence lift the above limitations ( Ziegler et al. , 2019 ; Xu et al. , 2021 ; Chen et al. , 2020 ) . However , these finetuned models are prone to overfitting when only a small amount of training data is available . In order to alleviate such an overfitting problem , additive learning has been extensively exploited by incorporating task-specific adapters into the PLM ( Zeldes et al. , 2020 ; Stickland & Murray , 2019 ) . In general , task-specialized adapters for few-shot NLG are trained by maximum likelihood estimation ( MLE ) . While MLE is efficient in learning , it suffers from the exposure bias problem due to the difference in the training and inference mechanisms ( He et al. , 2019 ) , and this problem can be severe with limited training data . Reinforcement learning ( RL ) is capable of resolving this exposure bias problem by sequential output sampling during training ( Ranzato et al. , 2015 ; Keneshloo et al. , 2019 ; Shi et al. , 2021 ) . Moreover , it allows to leverage the target-specific sequence-level objectives such as BLEU and ROUGE ( Wu et al. , 2018 ; Guo et al. , 2021 ) . However , in the task of NLG , the exponentially large space of output sequences restricts the use of RL since it leads to high variance and unstable training which is more serious in the few-shot setting . In this work , we develop a novel RLbased additive learning algorithm on the transformer-based PLM to overcome these shortcomings and to improve the performance of few-shot NLG . In particular , we first convert the NLG task to the sequential token generation task based on the transformer language model , and then propose a selective token generation between the PLM and the task-specific adapter , during both RL-based training and inference . The proposed output token selection enables to not only explicitly maintain a general prior knowledge from the frozen PLM but also focus only on the task-relevant parts in sequence generation . In addition , in few-shot learning this partial token generation makes the task-specific adapter more resilient to overfitting and furthermore reduces the overall output space which leads to stable RL training . Here , in order to make the two token generators ( policies ) complement each other as well as to realize the robust output selection at the token level on the fly , we exploit a separate token-level policy selector . It is noted that both the policy selector and the task-specific adapter are simultaneously learned by the RL algorithm . Experimental results on various few-shot NLG tasks show that the proposed selective token generation outperforms the previous PLM-based additive learning algorithms with the comprehensive ( non-selective ) token generation . Our main contributions can be summarized as follows . • A novel selective token generation between the PLM and the task-specific adapter is proposed for transformer-based few-shot NLG . • A separate selecting module is exploited to adaptively determine each output token in a sequence both at training and testing time . • RL is applied to train both the policy selector and the task-specific adapter that is complementary to the PLM in text generation . • An extensive empirical validation on few-shot NLG tasks demonstrates that the proposed selective token generation performs better in comparison to the previous PLM-based additive learning algorithms . 2 BACKGROUND . 2.1 NATURAL LANGUAGE GENERATION . The goal of NLG is to generate a text sequence y = [ y0 , ... , yT ] for a given task , where yt is the tth output token from a vocabulary V , and T is the output sequence length . For this generation , we aims to model the distribution of y that is autoregressively factorized as pθ ( y ) = ∏T t=0 pθ ( yt|y < t ) , where θ denotes the model parameters and y < t = [ y0 , ... , yt−1 ] . Here , the conditional distribution to sample a token for each step , pθ ( yt|y < t ) , is defined by the softmax function on the output logits fθ ( yt|y < t ) . Note that in general , the language generation is conditioned on input context according to a given task . Here , we encode the conditioning context by the same sequential model for generating an output sequence , and for simplicity we omit it . In this work , we utilize the autoregressive transformer for our generative model . 2.2 ADDITIVE LEARNING FOR FEW-SHOT GENERATION . To effectively leverage the general linguistic knowledge , θ is first initialized by the PLM parameters , θLM , for NLG . Given N task-specific training instances , D = { yn∗ } Nn=1 , where yn∗ is the nth ground-truth output sequence , directly finetuning θLM usingD can incur the severe overfitting problem when N is small in the few-shot scenario . Therefore , we add the task-specific adapter , gθa parameterized by θa , on top of the PLM , and optimize only θa ( Zeldes et al. , 2020 ; Stickland & Murray , 2019 ) . In specific , we reformulate f ( ·|y < t ; θ ) = WTh ( y < t ; θh ) where W ∈ RH×|V| and h ∈ RH denote the weight matrix and the penultimate representations , respectively , and θ = { W , θh } . Then , we define the task-specific conditional distribution as follows : p ( yt|y < t ; θLM , θa ) = softmax ( WLM ThLM ( y < t ) +Wa T g ( hLM ( y < t ) ; θg ) ) , ( 1 ) where hLM ( y < t ) = h ( y < t ; θh , LM ) and θa = { Wa , θg } . Here , the summation of the PLM logits and the adapter logits is motivated by auxiliary training1 . It is noted that in our additive learning θa is updated while θLM is kept frozen . Hence , in the following we omit θLM such that pθa ( yt|y < t ) = p ( yt|y < t ; θLM , θa ) for simplicity . 2.3 MAXIMUM LIKELIHOOD ESTIMATION ( MLE ) . Given a small amount of training data D = { yn∗ } Nn=1 , MLE optimizes θ by maximizing the data log-likelihood as follows : θ̂ = argmax θ N∑ n=1 T∑ t=0 log pθ ( y n∗ t |yn∗ < t ) . ( 2 ) Here , the output token at each step is conditioned on not the previous sampled tokens from the current model but the previous ground-truth tokens yn∗ < t . Namely , tokens are drawn from the data distribution during training , which is opposed to that tokens are drawn from the model distribution at test time . This discrepancy , also known as the exposure bias , leads that the errors will be accumulated along the generated sequence at test time since the model is biased to only perform well on the ground-truth history distribution . Especially , this bias problem can be more severe in the few-shot training . In addition , the token-level cross-entropy loss in MLE training is different from the sequence-level test metrics such as BLEU and ROUGE that are commonly used in the tasks of NLG . 2.4 REINFORCEMENT LEARNING ( RL ) . As an alternative to MLE , RL is able to overcome the exposure bias problem of MLE by sequencelevel sampling from the model distribution during training ( Ranzato et al. , 2015 ) . Also , RL can improve the performance by directly optimizing the evaluation metrics ( Guo et al. , 2021 ) . In order to use RL for our additive learning , we reformulate our text generation as an RL problem : at each time step t , the agent takes the current state st = y < t as an input and performs an action at that outputs a token yt by a policy πθ ( at|st ) corresponding to pθ ( yt|y < t ) . Then , the agent receives a reward rt = r ( st , at ) and deterministically transitions to the next state st+1 . Here , note that the token-level intermediate reward rt = 0 , ∀t < T when we use the delayed reward associated with the sequencelevel evaluation metric between the two full sequences , y and y∗ . Let τ = { ( st , at , rt ) } Tt=0 be the trajectory generated by πθ . The RL objective for the optimal agent is to maximize the expected sum of future discounted rewards : J ( πθ ) = Eτ∼πθ [ T∑ t=0 γtrt ] , ( 3 ) where γ ∈ [ 0 , 1 ) is the discount factor . 1Although the auxiliary training is particularly designed for maximizing the likelihood of the target task output , it also can take an advantage for RL since the adapter logits are nearly zero before training is advanced . Namely , it lets the task-specific conditional distribution start learning from the distribution of PLM , not a uniform distribution . Among a number of algorithms to approximately optimize the RL objective , we employ an actorcritic algorithm ( Bahdanau et al. , 2017 ) since it explicitly optimizes the policy network and it can also alleviate the delayed reward problem . The actor-critic algorithm requires the additional critic network to estimate the value of a state , V π ( st ) = Eπ [ ∑T t′=t γ t′−trt′ |st ] = ∑ at π ( at|st ) Qπ ( st , at ) where the state-action value function Qπ ( st , at ) = Eπ [ ∑T t′=t γ t′−trt′ |st , at ] = rt + V π ( st+1 ) . We use the following policy gradient loss to learn the policy parameters θ : L = − T∑ t=0 Aπθ ( st , at ) log πθ ( at|st ) , ( 4 ) where Aπθ ( st , at ) = Qπθ ( st , at ) − V πθ ( st ) is the advantage function that quantifies how an action at is better than the average action in state st . In few-shot text generation , the extremely large action space ( |V|T ) as well as the small amount of training data often make it difficult to perform RL with degraded performances , even though we conduct the additive learning from the PLM . Furthermore , it commonly has a delayed reward function ( e.g . BLEU ) which is defined after an entire sequence generated . It is hard to decide which token and how much contributes to the reward . This problem is known as the credit assignment . Therefore , in this work , we propose a selective token generation for improving the RL-based additive learning .

Summary: The authors proposes a method to improve large pre-trained language model's ability in few-shot language generation. The main idea is to freeze the large pre-trained language model (PLM), fine-tune a copy of this PLM on the task at hand, and use a selector to switch between generating the next token from the freezes PLM or the fine-tuned PLM. Instead of the commonly used maximum likelihood estimation, the authors cast the optimization problem as a reinforcement learning (RL) problem using the task evaluation metric such as BLEU or ROUGE as the reward. The authors perform experiments on few-shot data-to-text and text summarization settings, comparing to a few baselines (fine-tuned PLM and a few variants of the proposed method).