On Reward Maximization and Distribution Matching for Fine-Tuning Language Models

1 INTRODUCTION ∇ ! 𝔼 '' ! 𝑅 ! ( 𝑥 ) ∇ ! 𝐷 # $ 𝑝 , 𝜋 ! Reward Maximization Distribution Matching ? 𝔼 ! ! ∇ '' 𝑅 '' 𝑥 + 𝔼 ! ! 𝑅 '' 𝑥 ∇ '' log 𝜋 '' ( 𝑥 ) Parametric reward Figure 1 : In this study we make connection between two popular paradigms for fine-tuning sequence generation models according to preferences Reward Maximization ( RM ) and Distribution Matching ( DM ) . Pre-trained language models ( Devlin et al. , 2019 ; Radford et al. , 2019 ) are changing the landscape of Machine Learning research and practice . Due to their strong generative capabilities many studies have found it sufficient to “ nudge ” these models to conform to global preferences defined over the generated sequences instead of training from scratch using annotated data . These preferences could include topic and sentiment ( Dathathri et al. , 2020 ) , valid musical notes and molecular structures ( Jaques et al. , 2017a ) , code compilability ( Korbak et al. , 2021 ) , balancing gender bias ( Khalifa et al. , 2021 ) , eval- uation metrics for Machine Translation and Summarization ( Ranzato et al. , 2016 ; Bahdanau et al. , 2016 ) , or direct human feedback ( Ziegler et al. , 2019 ; Stiennon et al. , 2020 ) . This large body of studies is driven by two paradigms : Reward Maximization ( RM ) and Distribution Matching ( DM ) . Reward Maximization RM exploits the intuitive notion that we can nudge pre-trained models towards some preferences by providing global sequence-level rewards when the model generates outputs that satisfy desired features . For instance , if the model is producing toxic content , we can apply Reinforcement Learning ( RL ) techniques to discourage it from producing similar content again in the future . However , the risk of naively applying RL is that the model can undergo catastrophic forgetting of its original distribution . For example , it can degenerate into producing a single nonsensical but at least nontoxic sequence . Although several studies have considered hand-crafting general rewards to ensure desirable features like fluency ( Liu et al. , 2016a ; Tambwekar et al. , 2019 ) , coming up with rewards that are not incomplete or imperfect is highly non-trivial ( Wu et al. , 2016 ; Vedantam et al. , 2015 ) . These challenges have sparked a wide discussion on the overall effectiveness of RM for some tasks such as Neural Machine Translation ( Choshen et al. , 2020 ; Kiegeland & Kreutzer , 2021 ) . Reward Maximization with KL-Control To tackle the aforementioned issues of “ catastrophic forgetting ” , several studies , still under an RM paradigm , have considered incorporating a distributional term inside the reward to be maximized . In particular Jaques et al . ( 2017b ; 2019 ) ; Ziegler et al . ( 2019 ) and Stiennon et al . ( 2020 ) have applied variations of KL-control ( Todorov , 2007 ; Kappen et al. , 2012 ) which adds a penalty term to the reward term so that the resulting policy does not deviate too much from the original one in terms of KL-divergence . The overall objective with the KL-penalty is maximized using an RL algorithm of choice including : PPO ( Schulman et al. , 2017a ) as in Ziegler et al . ( 2019 ) or Q-learning ( Mnih et al. , 2013 ) as in Jaques et al . ( 2017b ) . Adding this distributional KL-penalty to the reward raises some important questions : What effect does it have on the shape of the optimal policy ? Does this new objective have any interpretation from a distributional perspective ? Distribution Matching A different recent paradigm for fine-tuning language models to satisfy downstream preferences formulates the problem as Distribution Matching ( DM ) . This paradigm consists of two steps : first a target distribution incorporating the desired preferences is defined as an Energy-Based Model ( LeCun et al. , 2006 ) . Then the forward KL divergence is minimized between this target distribution and an auto-regressive policy using a family of algorithms referred to as Distributional Policy Gradients ( DPG ) ( Parshakova et al. , 2019b ; Khalifa et al. , 2021 ; Korbak et al. , 2021 ) . This approach capitalizes on the flexibility of EBMs in specifying the target distribution . For example , the EBM can be defined so that it conforms to all downstream preferences while its corresponding normalized distribution has a minimal KL divergence from the original , pre-trained language model , therefore tackling the problem of “ catastrophic forgetting ” ( Khalifa et al. , 2021 ) . Interestingly , this DM paradigm can also deal with distributional preferences , for instance , for debiasing language models by specifying that the generated sequences should be gender-balanced , i.e . that 50 % of generations contain female mentions . Such distributional constraints can not be defined in the RM paradigm where a reward is calculated for a single sequence . Contrasting these two paradigms for fine-tuning language models , we can notice the promises and limitations of each.1 RM approaches are equipped with a large arsenal of RL algorithms and optimization techniques that can be efficient in reward maximization , however they lack the distributional perspective that is necessary for avoiding catastrophic forgetting and imposing distributional preferences over LMs . DM approaches are better suited to tackle those limitations , however , the family of DPG algorithms currently used is not as rich as its RL counterpart . So far , the connections between these two seemingly distinct paradigms have not been thoroughly explored . By establishing such connections we might import ideas from one approach to the other . This is our goal in this paper , detailing the nuanced connections and applying them to a case-study in variance reduction . Overall , our contributions are the following : • We untangle the connections between the RM and DM paradigms for fine-tuning language models . We provide a detailed comparison between the family of DPG algorithms with Policy Gradients of standard RL . • We provide an interpretation of KL-control techniques from a distribution matching perspective , placing such techniques at an intermediate place between RM and DM . • We exploit these connections to theoretically justify applying baselines — a variance reduction technique from RL — to DPG and derive a particular choice of a baseline . On an array of controllable language generation experiments about constraining topic , sentiment , and gender distributions over a pre-trained language model , we show that adding baselines leads to superior performance on constraint satisfaction , stability on small batch sizes , and sample efficiency . 2 BACKGROUND . Standard Policy Gradients One popular method for adapting the behaviour of language models to certain preferences has been that of assigning a “ reward ” score R ( x ) for sequences x sampled 1See Appendix A for an extended discussion of related work . from an autoregressive language model ( policy ) πθ . Then , the simplest policy gradient algorithm in reinforcement learning , namely , REINFORCE ( Williams , 1992a ) , aims to find the policy πθ ( x ) that maximizes the average reward Ex∼πθR ( x ) , and this leads , via the so-called “ log derivative trick ” , to a gradient ascent algorithm based on : ∇θEx∼πθR ( x ) = Ex∼πθR ( x ) ∇θ log πθ ( x ) , ( 1 ) which iteratively samples x from πθ and update parameters by increments proportional to R ( x ) ∇θ log πθ ( x ) . KL-control Jaques et al . ( 2017b ; 2019 ) ; Ziegler et al . ( 2019 ) , inspired by KL-control ( Todorov , 2007 ; Kappen et al. , 2012 ) , add a KL penalty term to the reward objective to penalize large deviations from the original pretrained model . That is , they maximize the expectation Ex∼πθRzθ ( x ) , where : Rzθ ( x ) . = r ( x ) − β log πθ ( x ) a ( x ) , ( 2 ) and where β is a free hyperparameter balancing the wish to maximize r with that of not deviating from a. Distributional Policy Gradients Distributional Policy Gradients ( DPG ) ( Parshakova et al. , 2019b ) is a recent approach used to fit an autoregressive policy πθ to the distribution p ( x ) = P ( x ) /Z induced by the EBM P ( x ) , where Z = ∑ x P ( x ) is the normalization constant ( partition function ) . Given an arbitrary EBM P ( x ) , DPG optimizes the loss functionDKL ( p , πθ ) with respect to the parameters θ of an autoregressive model πθ , a loss which is minimized for πθ = p. The KL-divergence minimization objective leads to a gradient estimate of the form : ∇θDKL ( p , πθ ) =−∇θEx∼p log πθ ( x ) ( 3 ) =− ∑ x p ( x ) ∇θ log πθ ( x ) = − 1 Z ∑ x P ( x ) ∇θ log πθ ( x ) ( 4 ) =− 1 Z Ex∼πθ P ( x ) πθ ( x ) ∇θ log πθ ( x ) . ( 5 ) 3 REWARD MAXIMIZATION VS DISTRIBUTION MATCHING . In the previous section , we have summarized three approaches that have been suggested for fine-tuning language models . Two of them can be characterized as “ Reward Maximization ” ( RM ) : Standard Policy Gradients ( PG ) and KL Control . On the other hand , DPG clearly belongs to the realm of “ Distribution Matching ” ( DM ) as it first defines the target distribution and then optimizes a policy to match it . In the rest of this section , we will explore connections between these two seemingly distinct concepts and , in the following section , we will exploit them to improve DM-based methods . 3.1 STANDARD VS. PARAMETRIC REWARDS . Let us start with distinguishing between a “ parametric reward ” Rθ which depends on θ and a standard reward R , which does not . If we wished to maximize the expected parametric reward , EπθRθ ( x ) , we would follow its gradient , leading to the identities : ∇θEx∼πθRθ ( x ) = ∇θ ∑ x πθ ( x ) Rθ ( x ) ( 6 ) = ∑ x πθ ( x ) ∇θRθ ( x ) + ∑ x Rθ ( x ) ∇θπθ ( x ) ( 7 ) = ∑ x πθ ( x ) ∇θRθ ( x ) + ∑ x πθ ( x ) Rθ ( x ) ∇θ log πθ ( x ) ( 8 ) = Ex∼πθ∇θRθ ( x ) ︸ ︷︷ ︸ RG-term +Ex∼πθRθ ( x ) ∇θ log πθ ( x ) ︸ ︷︷ ︸ PG-term . ( 9 ) Equation ( 9 ) is the sum of two terms : the first one , the “ RG-term '' ( Reward Gradient term ) , involves the gradient of the reward . The second one , the “ PG-term ” ( Policy Gradient term ) , was obtained using the “ log derivative trick ” and involves the gradient of the policy stricto sensu . In standard RL , where the reward does not depend on θ , the RG-term disappears and the gradient of expected reward consists solely of the PG-term . However , when Rθ depends on θ , the gradients are distinct ( apart from specific cases where the RG-term evaluates to 0 , as we will see below ) .

Summary: The paper builds upon previous work on fine-tuning pre-train language models to conform to pre-specified "distributional constraints" (such as: "the model should generate the same proportion of sentences with female or male subjects", etc...). Specifically, the authors draw parallel between a previous approach (distribution matching, DM) and RL-based approaches a la REINFORCE. Based on this comparison, they propose to adapt variance reduction techniques from RL (in particular the use of a baseline) to DM in order to facilitate training. Experimental results show that the resulting improved approach boasts better sample-efficiency, improved stability and overall competitive performance.