Original: On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting
Revised: On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: The availability of large pre-trained models is changing the landscape of Machine
Revised: The availability of large pre-trained models is changing the landscape of Machine

Original: Learning research and practice, moving from a “training from scratch” to a “finetuning” paradigm.
Revised: Learning research and practice, moving from a “training from scratch” to a “ﬁnetuning” paradigm.

Original: While in some applications the goal is to “nudge” the pre-traineddistribution towards preferred outputs, in others it is to steer it towards a differentdistribution over the sample space.
Revised: While in some applications the goal is to “nudge” the pre-traineddistribution towards preferred outputs, in others it is to steer it towards a differentdistribution over the sample space.

Original: Two main paradigms have emerged to tackle thischallenge: Reward Maximization (RM) and, more recently, Distribution Matching(DM).
Revised: Two main paradigms have emerged to tackle thischallenge: Reward Maximization (RM) and, more recently, Distribution Matching(DM).

Original: RM applies standard Reinforcement Learning (RL) techniques, such as
Revised: RM applies standard Reinforcement Learning (RL) techniques, such as

Original: Policy Gradients, to gradually increase the reward signal.
Revised: Policy Gradients, to gradually increase the reward signal.

Original: DM prescribes to firstmake explicit the target distribution that the model is fine-tuned to approximate.
Revised: DM prescribes to ﬁrstmake explicit the target distribution that the model is ﬁne-tuned to approximate.

Original: Here we explore the theoretical connections between the two paradigms, andshow that methods such as KL-control developed for RM can also be construedas belonging to DM.
Revised: Here we explore the theoretical connections between the two paradigms, andshow that methods such as KL-control developed for RM can also be construedas belonging to DM.

Original: We further observe that while DM differs from RM, it cansuffer from similar training difficulties, such as high gradient variance.
Revised: We further observe that while DM differs from RM, it cansuffer from similar training difﬁculties, such as high gradient variance.

Original: We leverageconnections between the two paradigms to import the concept of baseline into DMmethods.
Revised: We leverageconnections between the two paradigms to import the concept of baseline into DMmethods.

Original: We empirically validate the benefits of adding a baseline on an array ofcontrollable language generation tasks such as constraining topic, sentiment, andgender distributions in texts sampled from a language model.
Revised: We empirically validate the beneﬁts of adding a baseline on an array ofcontrollable language generation tasks such as constraining topic, sentiment, andgender distributions in texts sampled from a language model.

Original: We observe superiorperformance in terms of constraint satisfaction, stability and sample efficiency.
Revised: We observe superiorperformance in terms of constraint satisfaction, stability and sample efﬁciency.

Original:  Introduction
Revised:  Introduction

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: Parametricreward

Original: -trained language models (Devlin et al.
Revised: -trained language models (Devlin et al.

Original: Radford et al., 2019) are changing thelandscape of Machine Learning research andpractice.
Revised: Radford et al., 2019) are changing thelandscape of Machine Learning research andpractice.

Original: 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 sequencesinstead of training from scratch using annotateddata.
Revised: Due to their strong generative capabilities many studies have found it sufﬁcient to“nudge” these models to conform to global preferences deﬁned over the generated sequencesinstead of training from scratch using annotateddata.

Original: These preferences could include topic andsentiment (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), evaluation 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).
Revised: These preferences could include topic andsentiment (Dathathri et al., 2020), valid musical notes and molecular structures (Jaques et al.,2017a), code compilability (Korbak et al., 2021),reducing distributional biases (Khalifa et al., 2021; Weidinger et al., 2021), evaluation metrics for Machine Translation and Summarization (Ranzato et al., 2016; Bahdanau et al., 2016), or directhuman feedback (Ziegler et al., 2019; Stiennon et al., 2020).

Original: This large body of studies is driven by two different paradigms: Reward Maximization (RM) and Distribution Matching (DM).
Revised: This large body of studies is driven bytwo different paradigms: Reward Maximization (RM) and Distribution Matching (DM).

Original: Reward Maximization RM exploits the intuitive notion that we can nudge pre-trained modelstowards some preferences by providing global sequence-level rewards when the model generatesoutputs that satisfy desired features.
Revised: Reward Maximization RM intuitively nudge pre-trained models towards some preferences byproviding global sequence-level rewards when the model generates outputs that satisfy desiredfeatures.

Original: 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 thefuture.
Revised: For instance, if the model is producing toxic content, we can apply Reinforcement Learning(RL) techniques to discourage it from producing similar content.

Original: However,the risk of naively applying RL is thatthe model can undergo catastrophic forgettingof its original distribution.
Revised: However, naively applying RL yieldmodels that undergo catastrophic forgetting of its original distribution.

Original: For example, it can degenerate into producing a single nonsensical but atleast nontoxic sequence.
Revised: For example, it can degenerateinto producing a single nonsensical but at least nontoxic sequence.

Original: Although several studies have considered hand-crafting general rewards toensure desirable features like fluency (Liu et al., 2016a; Tambwekar et al., 2019), coming up withrewards that are not incomplete or imperfect is highly non-trivial (Wu et al., 2016; Vedantam et al.,2015).
Revised: Although several studies haveconsidered hand-crafting general rewards to ensure desirable features like ﬂuency (Liu et al., 2016a; Tambwekar et al., 2019), coming up with complete or perfect rewards is highly non-trivial (Wu et al. ,2016; Vedantam et al., 2015).

Original: These challenges have sparked a wide discussion on the overall effectiveness of RM for sometasks such as Neural Machine Translation (Choshen et al., 2020; Kiegeland & Kreutzer, 2021).
Revised: This has sparked a wide discussion on the overall effectiveness of RMfor some tasks such as NMT (Choshen et al., 2020; Kiegeland & Kreutzer, 2021).

Original: Reward Maximization with KL-Control To tackle the aforementioned issues of “catastrophicforgetting”, several studies, still under an RM paradigm, have considered incorporating a distributionalterm inside the reward to be maximized.
Revised: Reward Maximization with KL-Control To tackle the aforementioned issues of “catastrophicforgetting”, several studies, still under an RM paradigm, have considered incorporating a distributionalterm inside the reward to be maximized.

Original: In particular Jaques et al.
Revised: In particular Jaques et al.

Original: (2017b, 2019); Ziegler et al.(2019) and Stiennon et al.
Revised: (2017b, 2019); Ziegler et al.(2019) and Stiennon et al.

Original: (2020) have applied variations of KL-control (Todorov, 2007; Kappenet al., 2012) which adds a penalty term to the reward term so that the resulting policy does not deviatetoo much from the original one in terms of KL-divergence.
Revised: (2020) have applied variations of KL-control (Todorov, 2007; Kappenet al., 2012) which adds a penalty term to the reward term so that the resulting policy does not deviatetoo much from the original one in terms of KL-divergence.

Original: The overall objective with the KL-penaltyis maximized using an RL algorithm of choice including: PPO (Schulman et al., 2017a)
Revised: The overall objective with the KL-penaltyis maximized using an RL algorithm of choice including: PPO (Schulman et al., 2017a)

Original: as in Ziegleret al.
Revised: as in Ziegleret al.

Original: (2019) or Q-learning (Mnih et al., 2013)
Revised: (2019) or Q-learning (Mnih et al., 2013)

Original: as in Jaques et al. (2017b).
Revised: as in Jaques et al. (2017b).

Original: Adding this distributional
Revised: Adding this distributional

Original: KL-penalty to the reward raises some important questions: What effect does it have on the shape ofthe optimal policy?
Revised: KL-penalty to the reward raises some important questions: What effect does it have on the shape ofthe optimal policy?

Original: Does this new objective have any interpretation from a distributional perspective?
Revised: Does this new objective have any interpretation from a distributional perspective?

Original: Distribution Matching A different recent paradigm for fine-tuning language models to satisfydownstream preferences formulates the problem as Distribution Matching (DM).
Revised: Distribution Matching A different recent paradigm for ﬁne-tuning language models to satisfydownstream preferences formulates the problem as Distribution Matching (DM).

Original: This paradigmconsists of two steps: first a target distribution incorporating the desired preferences is defined as an
Revised: This paradigmconsists of two steps: ﬁrst a target distribution incorporating the desired preferences is deﬁned as an

Original: Energy-Based Model (LeCun et al., 2006).
Revised: Energy-Based Model (LeCun et al., 2006).

Original: Then the forward KL divergence is minimized betweenthis target distribution and an auto-regressive policy using a family of algorithms referred to as
Revised: Then the forward KL divergence is minimized betweenthis target distribution and an auto-regressive policy using a family of algorithms referred to as

Original: Distributional Policy Gradients (DPG) (Parshakova et al., 2019b; Khalifa et al., 2021; Korbak et al.,2021).
Revised: Distributional Policy Gradients (DPG) (Parshakova et al., 2019b; Khalifa et al., 2021; Korbak et al.,2021).

Original: This approach capitalizes on the flexibility of EBMs in specifying the target distribution.
Revised: This approach capitalizes on the ﬂexibility of EBMs in specifying the target distribution.

Original: For example, the EBM can be defined so that it conforms to all downstream preferences while itscorresponding normalized distribution has a minimal KL divergence from the original, pre-trainedlanguage model, therefore tackling the problem of “catastrophic forgetting” (Khalifa et al., 2021).
Revised: For example, the EBM can be deﬁned so that it conforms to all downstream preferences while itscorresponding normalized distribution has a minimal KL divergence from the original, pre-trainedlanguage model, therefore tackling the problem of “catastrophic forgetting” (Khalifa et al., 2021).

Original: 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.
Revised: 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.

Original: Such distributional constraints cannot be definedin the RM paradigm where a reward is calculated for a single sequence.
Revised: Such distributional constraints cannot be deﬁnedin the RM paradigm where a reward is calculated for a single sequence.

Original: Contrasting these two paradigms for fine-tuning language models, we can notice the promises andlimitations of each.
Revised: We can notice the promises and limitations of these two paradigms for ﬁne-tuning language models .

Original: 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 distributionalperspective that is necessary for avoiding catastrophic forgetting and imposing distributional preferences over LMs.
Revised: RM approaches are equipped with an arsenal of RL algorithms and optimization techniques that canbe efﬁcient in reward maximization, however they lack the distributional aspect to avoid catastrophicforgetting and impose distributional preferences over LMs.

Original: DM approaches arebetter suited to tackle those limitations, however, the family of DPG algorithms currently used is not as rich as its RL counterpart.
Revised: DM approaches are suited to tackle thoselimitations, however, the family of DPG algorithms currently used is not as rich as its RL counterpart.

Original: So far, the connections betweenthese two seemingly distinct paradigms have  not been thoroughly explored.
Revised: While the connections between these two seemingly distinct paradigms have been noted (Parshakovaet al., 2019b), they have not been explored in detail.

Original: By establishing suchconnections we might import ideas from one approach to the other.
Revised: Clarifying such connections might help importideas from one approach to the other.

Original: This is our goal in this paper,detailing the nuanced connections and applying them to a case-study in variance reduction.
Revised: This is our goal in this paper, detailing the nuanced connectionsand applying them to a case-study in variance reduction.

Original: Overall,our contributions are the following:
Revised: Overall, our contributions are the following:

Original: [Figure] • We reveal connections between the RM and DM paradigms through a detailed comparisonbetween the family of DPG algorithms with Policy Gradients (Table 1).•
Revised: • We clarify relations between the RM and DM paradigms through a detailed comparison between the family of DPG algorithms and Policy Gradients (Table 1), stressing the differences between

Revised: parametric and non-parametric rewards that are important in this regard.

Original:  We introduce an interpretation of KL-control techniques from a distribution matching perspective,placing such techniques at an intermediate place between RM and DM (Theorem 1).• We show how these connections can enable cross-pollination between the two perspectives byapplying baselines — a variance reduction technique from RL — to DPG and derive a particularchoice of a baseline (Facts 1 and 2).
Revised: • We introduce an interpretation of KL-control techniques from a distribution matching perspective,placing such techniques at an intermediate place between RM and DM (Theorem 1).• We show how these connections can enable cross-pollination between the two perspectives byapplying baselines — a variance reduction technique from RL — to DPG and derive a particular choice of a baseline (Facts 1 and 2).

Original: On an array of controllable language generation experiments,we show that adding baselines leads to superior performance on constraint satisfaction (Figure 3),stability on small batch sizes, and sample efficiency (Figure 4).
Revised: On an array of controllable language generation experiments,we show that adding baselines leads to superior performance on constraint satisfaction (Figure 3), stability on small batch sizes, and sample efﬁciency (Figure 4).

Original:  Background
Revised:  Background

Original: Standard Policy Gradients One popular method for adapting the behaviour of language modelsto certain preferences has been that of assigning a “reward” score R ( x ) for sequences x sampledfrom an autoregressive language model (policy) π θ .
Revised: Standard Policy Gradients One popular method for adapting the behaviour of language modelsto certain preferences has been that of assigning a “reward” score R ( x ) for sequences x sampledfrom an autoregressive language model (policy) ⇡ ✓ .

Original: Then, the simplest policy gradient algorithm inreinforcement learning, namely, REINFORCE (Williams, 1992a), aims to find the policy π θ ( x ) thatmaximizes the average reward E x ∼ π θ R ( x ) , and this leads, via the so-called “log derivative trick”, toa gradient ascent algorithm that iteratively samples x from π θ and update parameters by incrementsproportional to R ( x ) ∇ θ log π θ ( x ) via the following identity:
Revised: Then, the simplest policy gradient algorithm inreinforcement learning, namely, REINFORCE (Williams, 1992a), aims to ﬁnd the policy ⇡ ✓ ( x ) thatmaximizes the average reward E x ⇠ ⇡ ✓ R ( x ) , and this leads, via the so-called “log derivative trick”, toa gradient ascent algorithm that iteratively samples x from ⇡ ✓ and update parameters by incrementsproportional to R ( x ) r ✓ log ⇡ ✓ ( x ) via the following identity:

Original: [Equation]
Revised: [Equation]

Original: KL-control (Todorov, 2007; Kappen et al., 2012), was leveraged by Jaques et al.
Revised: KL-control (Todorov, 2007; Kappen et al., 2012), was leveraged by Jaques et al.

Original: (2017b, 2019) and
Revised: (2017b, 2019) and

Original: Ziegler et al.
Revised: Ziegler et al.

Original: (2019) to include a KL penalty term in the reward function to penalize large deviationsfrom the original pretrained model a ( x ) , weighted by a free hyperparameter β to control the trade-offbetween the two goals.
Revised: (2019) to include a KL penalty term in the reward function to penalize large deviationsfrom the original pretrained model a ( x ) , weighted by a free hyperparameter  to control the trade-offbetween the two goals.

Original: That is, they maximize the expectation E x ∼ π θ R zθ ( x ) , where:
Revised: That is, they maximize the expectation E x ⇠ ⇡ ✓ R z ✓ ( x ) , where:

Original: [Equation]
Revised: [Equation]

Original: Distributional Policy Gradients (DPG) (Parshakova et al., 2019b) is a recent approach used to fitan autoregressive policy π θ to the distribution p ( x ) 
Revised: Distributional Policy Gradients (DPG) (Parshakova et al., 2019b) is a recent approach used to ﬁtan autoregressive policy ⇡ ✓ to the distribution p ( x ) =

Original: = P ( x ) /Z induced by the EBM P ( x ) , where
Revised: P ( x ) /Z induced by the EBM P ( x ) , where

Original: Z = (cid:80) x P ( x ) is the normalization constant (partition function).
Revised: Z = P x P ( x ) is the normalization constant (partition function).

Original: Given an arbitrary EBM P ( x ) ,
Revised: Given an arbitrary EBM P ( x ) ,

Original: DPG optimizes the loss function D KL ( p, π θ ) with respect to the parameters θ of an autoregressivemodel π θ , a loss which is minimized for π θ = p .
Revised: DPG optimizes the loss function D KL ( p, ⇡ ✓ ) with respect to the parameters ✓ of an autoregressivemodel ⇡ ✓ , a loss which is minimized for ⇡ ✓ = p .

Original: The KL-divergence minimization objective leads toa gradient estimate of the form:
Revised: The KL-divergence minimization objective leads toa gradient estimate of the form:

Original: [Equation]
Revised: [Equation]

Original:  Reward Maximization vs Distribution Matching
Revised:  Reward Maximization vs Distribution Matching

Original: In the previous section, we have summarized three approaches that have been suggested for fine-tuninglanguage models.
Revised: In the previous section, we have summarized three approaches that have been suggested for ﬁne-tuninglanguage models.

Original: Two of them can be characterized as “Reward Maximization” (RM): Standard
Revised: Two of them can be characterized as “Reward Maximization” (RM): Standard

Original: Policy Gradients (PG) and KL-control.
Revised: Policy Gradients (PG) and KL-control.

Original: On the other hand, DPG clearly belongs to the realm of“Distribution
Revised: On the other hand, DPG clearly belongs to the realm of“Distribution

Original: Matching” (DM) as it first defines the target distribution and then optimizes a policy tomatch it.
Revised: Matching” (DM) as it ﬁrst deﬁnes the target distribution and then optimizes a policy tomatch it.

Original: In the rest of this section, we will explore connections between these two seemingly distinctconcepts and, in the following section, we will exploit them to improve DM-based methods.
Revised: In the rest of this section, we will explore connections between these two seemingly distinctconcepts and, in the following section, we will exploit them to improve DM-based methods.

Original: 3.1 Standard vs. Parametric Rewards
Revised: 3.1 Standard vs. Parametric Rewards

Original: Let us start with distinguishing between a “parametric reward” R θ which depends on θ and a standardreward R , which does not.
Revised: Let us start with distinguishing between a “parametric reward” R ✓ which depends on ✓ and a standardreward R , which does not.

Original: If we wished to maximize the expected parametric reward, E π θ R θ ( x ) , wewould follow its gradient, leading to the identities:
Revised: If we wished to maximize the expected parametric reward, E ⇡ ✓ R ✓ ( x ) , wewould follow its gradient, leading to the identities:

Original: [Equation]
Revised: [Equation]

Original: Equation (8) is the sum of two terms: the first one, the “RG-term" (Reward Gradient term), involvesthe gradient of the reward.
Revised: Equation (8) is the sum of two terms: the ﬁrst one, the “RG-term" (Reward Gradient term), involvesthe gradient of the reward.

Original: 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 .
Revised: 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 .

Original: In standard RL,where the reward does not depend on θ , the RG-term disappears and the gradient of expected rewardconsists solely of the PG-term.
Revised: In standard RL,where the reward does not depend on ✓ , the RG-term disappears and the gradient of expected rewardconsists solely of the PG-term.

Original: However, when R θ depends on θ , the gradients are distinct (apartfrom specific cases where the RG-term evaluates to 0 , as we will see below).
Revised: However, when R ✓ depends on ✓ , the gradients are distinct (apartfrom speciﬁc cases where the RG-term evaluates to 0 , as we will see below).

Original: 3.2 KL-control as Distribution Matching
Revised: 3.2 KL-control as Distribution Matching

Original: Adding a KL-penalty term to the reward (as in the case of KL-control) leads to a parametric reward.
Revised: Adding a KL-penalty term to the reward (as in the case of KL-control) leads to a parametric reward.

Original: However, due to the particular form of its objective, the RG-term actually vanishes , 2 leaving onlythe PG-term E x ∼ π θ R zθ ( x ) ∇ θ log π θ ( x ) and simplifying the tuning procedure to a standard Policy
Revised: However, due to the particular form of its objective, the RG-term actually vanishes , 2 leaving onlythe PG-term E x ⇠ ⇡ ✓ R z ✓ ( x ) r ✓ log ⇡ ✓ ( x ) and simplifying the tuning procedure to a standard Policy

Original: Gradient.
Revised: Gradient.

Original: While this algorithm falls under the RM paradigm, here we argue that is it multifaceted, andexplore deeper connections with the DM paradigm.
Revised: While this algorithm falls under the RM paradigm, here we argue that is it multifaceted, andexplore deeper connections with the DM paradigm.

Original: More precisely, the maximization of the rewardwith the KL penalty term is equivalent to a distributional matching with an underlying emergentsequential EBM, a remark that already reveals some similarities with DPG.
Revised: More precisely, the maximization of the rewardwith the KL penalty term is equivalent to a distributional matching with an underlying emergentsequential EBM, a remark that already reveals some similarities with DPG.

Original: Theorem 1.
Revised: Theorem 1.

Original: Consider the following EBM:
Revised: Consider the following EBM:

Original: [Equation]
Revised: [Equation]

Original: 9) and let p z be the normalized distribution p z ( x )
Revised: and let p z be the normalized distribution p z ( x )

Original: = 

Revised: =

Original: Z P z ( x ) , with Z = (cid:80) x P z ( x ) .
Revised: Z P z ( x ) , with Z = P x P z ( x ) .

Original: Then:
Revised: Then:

Original: (i) arg max [Equation] ∼ [Equation]
Revised: (i) arg max ⇡ ✓ E x ⇠ ⇡ ✓ R

Original: ( x ) =

Original: [Equation]

Original: π

Original: [Equation] D KL ( π [Equation] p z ) 
Revised: z ✓ ( x ) = arg min ⇡ ✓ D KL ( ⇡ ✓ , p z ) ;

Original: [Equation] (ii) arg max π ∈D ( X ) E x ∼ π R zπ ( x ) =
Revised: (ii) arg max ⇡ 2 D ( X ) E x ⇠ ⇡ R z ⇡ ( x ) =

Original: p z , where D ( X ) is the family of all distributions over 
Revised: p z , where D ( X ) is the family of all distributions over X ,

Original: X , and R zπ ( x ) .
Revised: and R z ⇡ ( x ) .

Original: = r ( x ) − β log π ( x ) a ( x ) .
Revised: = r ( x )   log ⇡ ( x ) a ( x ) .

Original: Proof.
Revised: Proof.

Original: A simple way to prove this is to notice that the expectation of the reward R zθ has a monotonically decreasing relationship with the reverse KL divergence between π θ and p z
Revised: A simple way to prove this is to notice that the expectation of the reward R z ✓ has a monotonically decreasing relationship with the reverse KL divergence between ⇡ ✓ and p z

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: so that the arg min π θ D KL ( π θ , p z ) coincides with the arg max π θ E x ∼ π θ R zθ ( x ) , proving (i).
Revised: so that the arg min ⇡ ✓ D KL ( ⇡ ✓ , p z ) coincides with the arg max ⇡ ✓ E x ⇠ ⇡ ✓ R z ✓ ( x ) , proving (i).

Original: Onthe other hand, arg min π ∈D ( X ) D KL ( π, p z ) , which also corresponds to arg max π ∈D ( X ) E x ∼ π R zπbecause of (i) applied to a family π θ ′ covering D ( X ) in full, is just p z , concluding the proof.
Revised: Onthe other hand, arg min ⇡ 2 D ( X ) D KL ( ⇡ , p z ) , which also corresponds to arg max ⇡ 2 D ( X ) E x ⇠ ⇡ R z ⇡because of (i) applied to a family ⇡ ✓ 0 covering D ( X ) in full, is just p z , concluding the proof.

Original: Overall, we can conclude that the addition of the distributional term (KL-penalty) to the rewarddoes indeed provide a DM interpretation, namely in terms of minimizing the reverse KL divergencewith an emergent underlying distribution p z ( x ) .
Revised: Overall, we can conclude that the addition of the distributional term (KL-penalty) to the rewarddoes indeed provide a DM interpretation, namely in terms of minimizing the reverse KL divergencewith an emergent underlying distribution p z ( x ) .

Original: We note that p z ( x ) does not correspond to a freeand explicit choice of EBM (e.g. one that balances the gender and topic distributions of a languagemodel).
Revised: We note that p z ( x ) does not correspond to a freeand explicit choice of EBM (e.g. one that balances the gender and topic distributions of a languagemodel).

Original: Instead equation (9) appears in a restrictive format, which is implicitly defined by the reward
Revised: Instead equation (9) appears in a restrictive format, which is implicitly deﬁned by the reward

Original: R zθ , along with a β hyperparameter without a clear meaning.
Revised: R z ✓ , along with a  hyperparameter without a clear meaning.

Original: By contrast, the DPG algorithmsare designed to perform DM on any EBM specification, corresponding to an explicit distributionalobjective.
Revised: By contrast, the DPG algorithmsare designed to perform DM on any EBM speciﬁcation, corresponding to an explicit distributionalobjective.

Original: 3.3 Similarities and Differences between DPG and Policy Gradients
Revised: 3.3 Similarities and Differences between DPG and Policy Gradients

Original: In the previous subsection, we have connected KL-control, a method designed under a RM paradigm,to DM.
Revised: In the previous subsection, we have connected KL-control, a method designed under a RM paradigm,to DM.

Original: Now, we turn to the converse question of whether DPG, a DM method, can be connected to
Revised: Now, we turn to the converse question of whether DPG, a DM method, can be connected to

Original: RM.
Revised: RM.

Original: We begin by noting that after defining R θ = P ( x ) π θ ( x ) , the DPG gradient E x ∼ π θ P ( x ) π θ ( x ) ∇ θ log π θ ( x )acquires the format of the PG-term E π θ R θ ∇ θ log π θ ( x ) .
Revised: We begin by noting that after deﬁning R ✓ = P ( x ) ⇡ ✓ ( x ) , the DPG gradient E x ⇠ ⇡ ✓ P ( x ) ⇡ ✓ ( x ) r ✓ log ⇡ ✓ ( x ) acquires the format of the PG-term E ⇡ ✓ R ✓ r ✓ log ⇡ ✓ ( x ) .

Original: However, the DM objective of DPG cannot be considered as maximizing the average “reward”
Revised: However, the DM objective of DPG cannot be considered as maximizing the average “reward”

Original: R θ ( x )= P ( x ) π θ ( x ) , as this would requireadding also the RG-term E π θ ∇ θ P ( x ) π θ ( x ) into the gradient, which case does not vanish [Equation] [Equation] [Equation]
Revised: ( x ) would require also the RG-term E r

Original: [Equation] [Equation] [Equation] [Equation]
Revised: P

Original: [Equation] [Equation]
Revised: ✓ ( x the gradient,

Original: DPG
Revised: DPG

Original: Reward

Original: [Equation]

Original: Baseline

Original: [Equation]

Original: Z ∇ θ with Baseline

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: Nonetheless, the analogy behind this gradient term is more fruitful than it first appears.
Revised: Nonetheless, the analogy behind this gradient term is more fruitful than it ﬁrst appears.

Original: As a matterof fact, DPG gradient estimates suffer from the same high-variance problems as with standard PG.
Revised: As a matterof fact, DPG gradient estimates suffer from the same high-variance problems as with standard PG.

Original: While the objective of DPG (distribution matching) is different from that of Policy Gradients (rewardmaximization), DPG also needs to estimate the PG-term E π θ R θ ( x ) ∇ θ log π θ ( x ) at a given value ofθ , using a batch of samples x .
Revised: While the objective of DPG (distribution matching) is different from that of Policy Gradients (rewardmaximization), DPG also needs to estimate the PG-term E ⇡ ✓ R ✓ ( x ) r ✓ log ⇡ ✓ ( x ) at a given value of✓ , using a batch of samples x .

Original: For such a fixed θ , we can define provisionally set R ( x ) .
Revised: For such a ﬁxed ✓ , we can deﬁne provisionally set R ( x ) .

Original: = R θ and theproblem of gradient estimation for this fixed θ is identical to the estimation E x ∼ π θ R ( x ) ∇ θ log π θ ( x )based on a set of samples x in standard RL.
Revised: = R ✓ and theproblem of gradient estimation for this ﬁxed ✓ is identical to the estimation E x ⇠ ⇡ ✓ R ( x ) r ✓ log ⇡ ✓ ( x )based on a set of samples x in standard RL.

Original: Therefore, the techniques that have been developedto reduce the variance of the gradients estimates in RL can be ported to DPG insofar as we arecomputing the gradient estimates at a given θ .In
Revised: Therefore, the techniques that have been developedto reduce the variance of the gradients estimates in RL can be ported to DPG insofar as we arecomputing the gradient estimates at a given ✓ .

Original:  Section 4, we show how one can import one suchvariance reduction technique to the DPG: baselines.
Revised: In Section 4, we show how one can import one suchvariance reduction technique to the DPG: baselines.

Original:  A Case Study on Variance Reduction
Revised:  A Case Study on Variance Reduction

Original: [Figure] Baselines are a standard variance reduction technique in thecontext of Policy Gradients (Sutton & Barto, 2018).
Revised: Baselines are a standard variance reduction technique in thecontext of Policy Gradients (Sutton & Barto, 2018).

Original: The idea isto subtract from the reward R ( x ) a value B that does not introduce bias to the gradients but may change variance.
Revised: The idea isto subtract from the reward R ( x ) a value B that does not introduce bias to the gradients but may change variance.

Original: After theintroduction of baseline, equation (1) then takes the followingform:
Revised: After theintroduction of baseline, equation (1) then takes the followingform:

Original: [Equation]
Revised: [Equation]

Original: In standard RL, the simplest form of baseline B is just theaverage of the rewards for the policy: 

Revised: In standard RL, the simplest form of baseline B is just theaverage of the rewards for the policy:

Original: [Equation]
Revised: [Equation]

Original: Following the same methodology of taking the baseline as the expectation of the reward term, we canobtain a remarkably simple form of a baseline for DPG: 

Revised: Following the same methodology of taking the baseline as the expectation of the reward term, we canobtain a remarkably simple form of a baseline for DPG:

Original: [Equation]
Revised: [Equation]

Original: Fact 1.
Revised: Fact 1.

Original: Subtracting B from R θ ( x ) does not introduce bias into DPG gradient estimates.
Revised: Subtracting B from R ✓ ( x ) does not introduce bias into DPG gradient estimates.

Original: Proof.
Revised: Proof.

Original: Let us rewrite the DPG gradient in (5) with the added baseline B =
Revised: Let us rewrite the DPG gradient in (5) with the added baseline B =

Original: Z :
Revised: Z :

Original: [Equation]
Revised: [Equation]

Revised: ⇠ ⇡ ✓

Original:  Here, the second term does not introduce bias because Z (cid:104) (cid:80) x ∇ θ π θ ( x ) (cid:105) = 0 , leaving us with theexact same form of gradient as in the original DPG algorithm.
Revised: [Equation] Here, the second term does not introduce bias because Z [Equation] i = 0 , leaving us with theexact same form of gradient as in the original DPG algorithm.

Revised: In the scope of this paper, our focus is on importing to DPG simple constant baselines.

Revised: The advantage is that this is a technique that is not impacted by the fact that R ✓ depends on ✓ : it can be applied “ ✓ -locally” to provide a more accurate estimate of E x ⇠ ⇡ ✓ R ✓ ( x ) r ✓ log ⇡ ✓ ( x ) for a ﬁxed ✓ , irrespective of the values of R ✓ 0 elsewhere, while variance reduction techniques that involve several ✓ 0 s simultaneously raise additional challenges for parametric rewards.

Original: Require: P , initial generative model a 1: π θ ← a , q← a 2: for each iteration do 3: for each episode do 4: sample x from q ( · ) 5: θ ← θ + α ( θ ) (cid:104) P ( x ) q ( x ) 
Revised: Require: P , initial generative model a 1: ⇡ ✓ a , q a 2: for each iteration do 3: for each episode do 4: sample x from q ( · ) 5: ✓ ✓ + ↵ ( ✓ ) h

Original: − Z π θ ( x ) q ( x ) (cid:105) ∇ θ log π θ ( x ) 6: if D KL ( p || π θ ) <
Revised: P ( x ) q ( x )  Z ⇡ ✓ ( x ) q ( x ) i r ✓ log ⇡ ✓ ( x ) 6: if D KL ( p || ⇡ ✓ ) <

Original: D KL ( p || q ) then 7: q ← π θ
Revised: D KL ( p || q ) then 7: q ⇡ ✓

Original: Ensure: π θ
Revised: Ensure: ⇡ ✓

Original: Note that since B RL depends on θ , it has tobe be re-estimated after each gradient update.
Revised: Note that since B RL depends on ✓ , it has tobe be re-estimated after each gradient update.

Original: On the other hand, B does not depend onθ , which is an advantage because B couldbe now estimated by averaging over samplesfrom all the different θ ’s without introducingbias, leading to a more accurate estimation.
Revised: On the other hand, B does not depend on✓ , which is an advantage because B couldbe now estimated by averaging over samplesfrom all the different ✓ ’s without introducingbias, leading to a more accurate estimation.

Original: See Table 1 for a comparison of these twoforms of baselines.
Revised: See Table 1 for a comparison of these twoforms of baselines.

Original: The off-policy DPG version introducedin (Parshakova et al., 2019b) and its KLadaptive variant (Khalifa et al., 2021) samplea proposal distribution q instead of the policy π θ .
Revised: The off-policy DPG version introducedin (Parshakova et al., 2019b) and its KLadaptive variant (Khalifa et al., 2021) samplea proposal distribution q instead of the policy ⇡ ✓ .

Original: Then, the baseline takes the form
Revised: Then, the baseline takes the form

Original: [Equation]
Revised: [Equation]

Original: (15) where the π θ ( x ) q ( x ) term is an importance weight correcting for the bias introduced by sampling from q .
Revised: (15) where the ⇡ ✓ ( x ) q ( x ) term is an importance weight correcting for the bias introduced by sampling from q .

Original: Similarly to the DPG case, we can prove the following (see Appendix C):
Revised: Similarly to the DPG case, we can prove the following (see Appendix C):

Original: Fact 2.
Revised: Fact 2.

Original: Subtracting B off ( x ) from R θ ( x ) does notintroduce bias into off-policy DPG gradientestimates.
Revised: Subtracting B off ( x ) from R ✓ ( x ) does not bias the off-policy DPG gradient estimates.

Original: In practice, as shown on Figure 2, adding a baseline to KL-adaptive DPG (Algorithm 1) centers theadvantage (defined as A .
Revised: In practice, as shown on Figure 2, adding a baseline to KL-adaptive DPG (Algorithm 1) centers theadvantage (deﬁned as A .

Original: =
Revised: =

Original: P ( x ) q ( x ) − Z π θ ( x ) q ( x ) ) around 0 leading to better performance on: convergence(section 4.3), stability on small batch sizes (section 4.4), and variance reduction (section 4.5).
Revised: P ( x ) q ( x )  Z ⇡ ✓ ( x ) q ( x ) ) around 0 leading to better performance on: convergence(section 4.3), stability on small batch sizes (section 4.4), and variance reduction (section 4.5).

Original: 4.1 Generation with Distributional Control
Revised: 4.1 Generation with Distributional Control

Original: We investigate the benefits of adding a baseline to the DPG algorithm, on the Generation with
Revised: We investigate the beneﬁts of adding a baseline to the DPG algorithm, on the Generation with

Original: Distributional Control (GDC) (Khalifa et al., 2021) framework.
Revised: Distributional Control (GDC) (Khalifa et al., 2021) framework.

Original: GDC makes use of DPG to controlthe properties of pre-trained language models to satisfy certain constraints.
Revised: GDC makes use of DPG to controlthe properties of pre-trained language models to satisfy certain constraints.

Original: In our experiments, followtarget distribution form of Parshakova et al.
Revised: In our experiments, followtarget distribution form of Parshakova et al.

Original: (2019a)
Revised: (2019a)

Original: and Khalifa et al.
Revised: and Khalifa et al.

Original: (2021), in which the EBM
Revised: (2021), in which the EBM

Original: P ( x ) is defined so that its normalized variant p ( x ) matches a set of desired moments constraints ongiven features ϕ i ( x ) , while having a minimal KL divergence D KL ( p, a ) from an original pretrainedlanguage model a , to avoid catastrophic forgetting.
Revised: P ( x ) is deﬁned so that its normalized variant p ( x ) matches a set of desired moments constraints ongiven features  i ( x ) , while having a minimal KL divergence D KL ( p, a ) from an original pretrainedlanguage model a , to avoid catastrophic forgetting.

Original: These constraints are expressed as conditions ¯ µ i =
Revised: These constraints are expressed as conditions ¯ µ i =

Original: E x ∼ p
Revised: E x ⇠ p

Original: ϕ i ( x ) , for i ∈ { 1 , . . . , n } , by which themoments (expectations) under the distribution p of each feature ϕ i ( x ) are required to take certaindesired values ¯ µ i .
Revised:  i ( x ) , for i 2 { 1 , . . . , n } , by which themoments (expectations) under the distribution p of each feature  i ( x ) are required to take certaindesired values ¯ µ i .

Original: For instance, let ϕ 1 ( x ) = 1{+
+} iff the topic of x is science and ϕ 2 ( x ) = 

Revised: For instance, let  1 ( x ) = {+
+} iff the topic of x is science and  2 ( x ) =

Original: iff xmentions a female person, then imposing moments ¯ µ 1 = 1 and ¯ µ 2 = 0 .
Revised: iff xmentions a female person, then imposing moments ¯ µ 1 = 1 and ¯ µ 2 = 0 .

Original: 5 constrains the languagemodel p to only generate sequences about science, half of which mention females.
Revised: 5 constrains the languagemodel p to only generate sequences about science, half of which mention females.

Original: P ( x ) is uniquelydetermined by the following form: 

Revised: P ( x ) is uniquelydetermined by the following form:

Original: [Equation]
Revised: [Equation]

Original: (16) where λ i terms control the moments µ i of the associated features, which can be estimated throughself-normalized importance sampling (Owen, 2013); and then, to make the moments match thedesired values, the λ i terms can be optimized through SGD (Parshakova et al., 2019a).
Revised: [Equation] where  i terms control the moments µ i of the associated features, which can be estimated throughself-normalized importance sampling (Owen, 2013); and then, to make the moments match thedesired values, the  i terms can be optimized through SGD (Parshakova et al., 2019a).

Original: 4.2 Experimental setup
Revised: 4.2 Experimental setup

Original: We demonstrate the benefits of adding a baseline to the family of Distributional Policy Gradientsalgorithms on an array of  controlled language generation tasks for English. For this, we modifythe GDC framework Khalifa et al. (2021) namely its KL-DPG algorithm to include a baseline asshown in Algorithm 1. We refer to this method in all the following experiments as GDC ++ .
Revised: We evaluate our method on an array of 10 controlled text generation tasks.

Original: Inaddition to comparing GDC ++ with GDC , for tasks for which the comparison is meaningful (i.e. for“pointwise constraints", see below), we compare with two reward maximization baselines: Reinforce(Williams, 1992b) and Ziegler (Ziegler et al., 2019). Reinforce tries to maximize the expectedreward E x ∼ π θ R ( x ) , where R ( x ) [-= 
 iff the pointwise constraints are met.-]
Revised: For each, given apre-trained language model a ( x ) , and a set of constraints, the objective of each ﬁne-tuning method isto obtain a ﬁne-tuned language model ⇡ ✓ that satisﬁes the imposed constraints while deviating asminimally as possible from the original language model a ( x ) [-= 
 iff the pointwise constraints are met.-].

Original: Ziegler instantiates the 
Revised: Constraints are deﬁned as a set of binary features {  i } and their corresponding desired percentages(moments) { ¯ µ i } within the generations of the target language model.

Original: KL-control approach: its objective includes a KL penalty term for departures from a . Tasks We evaluate GDC, GDC ++ and the two baselines on 10 tasks , which we construct throughsets of moment constraints { ¯ µ i } for binary features { ϕ i } . These include 6 sets of purely pointwiseconstraints (for which ¯ µ i = 1 ) and 4 sets including distributional constraints ( 0 < ¯ µ i < 1 ). Weconsider the following constraint types:
Revised: Based on the value of themoment constraints these 10 tasks are divided into 6 tasks of pointwise constraints (for which ¯ µ i = 1 ),2 tasks of distributional constraints ( 0 < ¯ µ i < 1 ) and 2 tasks of mixed type constraints (hybrid):

Original: (a) Single-word constraints, where ϕ ( x ) = 1 iff the a given word appears in the sequence x .
Revised: (a) Single-word constraints, where  ( x ) = 1 iff the a given word appears in the sequence x .

Original: We experimentboth with frequent words (task 1: “amazing” , original frequency:− and(task
Revised: We experiment with frequent words (task 1: “amazing”, original frequency: 10  4 ) and (task2: “WikiLeaks”, original frequency:

Original: “WikiLeaks”, original frequency

Original: 10 − 5 ) rare words,(
Revised: 10  rare words

Original: b) Wordlist constraints, where ϕ ( x ) = 

Revised: (b) Wordlist constraints, where  ( x ) =

Original: iff x contains at least one word from a given list.
Revised: iff x contains at least one word from a given list.

Original: Weconsider lists of word associated with politics (task 3) and science (task 4) published by Dathathri et al.
Revised: Weconsider lists of word associated with politics (task 3) and science (task 4) published byDathathri et al.

Original: (c) Sentiment classifier constraints, where ϕ ( x ) = 

Revised: (2020),(c) Sentiment classiﬁer constraints, where  ( x ) =

Original: if x is classified as positive (task 5), or negative (task 6) by a pre-trained classifier published by Dathathri et al.
Revised: if x is classiﬁed as positive (task 5), ornegative (task 6) by a pre-trained classiﬁer published by Dathathri et al.

Original: (d) A single distributional constraint where ϕ ( x ) = [-
-] iff x contains a female figure mention,and ¯ µ = 0 .
Revised: (2020).(d) A single distributional constraint where  ( x ) = [-
-]1 iff x contains a female ﬁgure mention,and ¯ µ = 0 .

Original: 5 (task 8),(e)
Revised: 5 (task 8),

Original:  A set of four distributional constraints: ϕ i ( x ) = 1{+
+} iff x contains at least one of the words inthe “science", “art", “sports" and “business" wordlists (compiled by Dathathri et al.
Revised: (e) A set of four distributional constraints:  i ( x ) = {+
+} iff x contains at least one of the words inthe “science", “art", “sports" and “business" wordlists (compiled by Dathathri et al.

Original: (2020)),respectively.
Revised: (2020)),respectively.

Original: For each i , ¯ µ i = 0 .
Revised: For each i , ¯ µ i = 0 .

Original: 25 (task 8),(f)
Revised: 25 (task 8),(f)

Original: Hybrid constraints where ϕ 1 ( x )
Revised: Hybrid constraints where  1 ( x )

Original: = 1{+
+} iff x contains more female than male pronouns,¯ µ 1 = 0 .
Revised: = {+
+} iff x contains more female than male pronouns,¯ µ 1 = 0 .

Original: 5 and ϕ 2 ( x ) = [-
-] iff x contains at least one of the words from the “sports" wordlist(task 9) or “politics” wordlist, ¯ µ 2 ( x ) = 1 (task 10).
Revised: 5 and  2 ( x ) = [-
-]1 iff x contains at least one of the words from the “sports" wordlist(task 9) or “politics” wordlist, ¯ µ 2 ( x ) = 1 (task 10).

Original: Following (Khalifa et al., 2021), for hybrid and distributional constraints (tasks 8-10) we compareonly GDC and GDC ++ because the RM objective of Ziegler and Reinforce is not equipped to handledistributional constraints.
Revised: Methods We modify the GDC framework Khalifa et al. (2021) namely its KL-DPG algorithmto include a baseline as shown in Algorithm 1. We refer to this method as GDC ++ . In addition tocomparing GDC ++ with GDC we compare with two reward maximization baselines: Reinforce(Williams, 1992b) and Ziegler (Ziegler et al., 2019).

Revised: Reinforce tries to maximize the expected reward

Revised: E x ⇠ ⇡ ✓ R ( x ) , where R ( x ) = 1 iff the pointwise constraints are met.

Revised: Ziegler instantiates the KLcontrol approach: its objective includes a KL penalty term for departures from a .

Revised: Following (Khalifaet al., 2021), for hybrid & distributional constraints (tasks 8-10) we compare only GDC and GDC ++because the RM objective of Ziegler and Reinforce is not equipped to handle them.

Original: Metrics We report the following metrics evaluated over batches of samples from π θ at eachvalidation step:
Revised: Metrics We report the following metrics at each validation step over batches of samples from ⇡ ✓ :

Original: E x ∼ π θ ϕ i ( x ) , measuring the ability to reach the target moment of the i -th feature.2.
Revised: E x ⇠ ⇡ ✓  i ( x ) , measuring the ability to reach the target moment of the i -th feature.2.

Original: D KL ( p, π θ ) , the forward KL divergence from the optimal target distribution p , 73.
Revised: D KL ( p, ⇡ ✓ ) , the forward KL divergence from the optimal target distribution p , 73.

Original: D KL ( π θ , a ) , the reverse KL divergence from the original pretrained language model a .4.
Revised: D KL ( ⇡ ✓ , a ) , the reverse KL divergence from the original pretrained language model a .4.

Original: Distinct-n score, a measure of text diversity in terms of the frequency of repetitions within asingle sample x , proposed by (Li et al., 2016a).5.
Revised: Distinct-n score, a measure of text diversity in terms of the frequency of repetitions within asingle sample x , proposed by (Li et al., 2016a).5.

Original: Self-BLEU-n, a measure of text diversity on a distributional level across samples proposedby (Zhu et al., 2018), ensuring that policies don’t converge into limited number of sequencesthat satisfy the imposed constraints Caccia et al.
Revised: Self-BLEU-n, a measure of text diversity on a distributional level across samples proposedby (Zhu et al., 2018), ensuring that policies don’t converge into limited number of sequencesthat satisfy the imposed constraints Caccia et al.

Original: Training details For tasks 1-6, we use a pre-trained GPT-2 small with 117M parameters (Radfordet al., 2019) as the original language model a .
Revised: Training details For tasks 1-6, we use a pre-trained GPT-2 small with 117M parameters (Radfordet al., 2019) as the original language model a .

Original: For tasks 7-10, a is the same pre-trained modeladditionally fine-tuned on the WikiBio (Lebret et al., 2016) dataset.
Revised: For tasks 7-10, a is the same pre-trained modeladditionally ﬁne-tuned on the WikiBio (Lebret et al., 2016) dataset.

Original: See Appendix E for more details.
Revised: See Appendix E for more details.

Original: 4.3 Results
Revised: 4.3 Results

Original: We present the evolution of our metrics through training epochs in Figure 3 (aggregated over tasks1-6) and Figure 6 in the Appendix (aggregated over tasks 7-10).
Revised: We present the evolution of our metrics through training epochs in Figure 3 (aggregated over tasks1-6) and Figure 6 in the Appendix (aggregated over tasks 7-10).

Original: Results for each task are presentedseparately on Figures 7-10 in the Appendix.
Revised: Results for each task are presentedseparately on Figures 7-10 in the Appendix.

Original: Consistent with prior work (Khalifa et al., 2021), we observe that Reinforce is able to quickly achievehigh levels of constraint satisfaction, but at the cost of large deviations from a , which translates intosignificantly decreased diversity of generated samples (in terms of Self-BLEU-5 and Distinct-1).
Revised: Consistent with prior work (Khalifa et al., 2021), we observe that Reinforce is able to quickly achievehigh levels of constraint satisfaction, but at the cost of large deviations from a , which translates intosigniﬁcantly decreased diversity of generated samples (in terms of Self-BLEU-5 and Distinct-1).

Original: The
Revised: The

Original: KL penalty term in Ziegler imposes an upper bound on deviation from a but the deviation is stillsignificant enough to result in a drop in diversity.
Revised: KL penalty term in Ziegler imposes an upper bound on deviation from a but the deviation is stillsigniﬁcant enough to result in a drop in diversity.

Original: Moreover, we have observed Ziegler’s objective toresult in very unstable training.
Revised: Moreover, we have observed Ziegler’s objective toresult in very unstable training.

Original: GDC and GDC ++ are the only fine-tuning methods that address constraint satisfaction based on aclear formal objective, i.e. reducing the divergence from p .
Revised: GDC and GDC ++ are the only ﬁne-tuning methods that address constraint satisfaction based on aclear formal objective, i.e. reducing the divergence from p .

Original: The approach translates into significantlysmaller deviations from a and maintaining diversity within and across samples.
Revised: The approach translates into signiﬁcantlysmaller deviations from a and maintaining diversity within and across samples.

Original: The addition ofa baseline indeed reduces the variance.
Revised: The addition ofa baseline indeed reduces the variance.

Original: We analyze that extensively in Appendix 4.5 while herefocusing on the downstream effects of variance reduction.
Revised: We analyze that extensively in Appendix 4.5 while herefocusing on the downstream effects of variance reduction.

Original: One is that π θ is now able to compound
Revised: One is that ⇡ ✓ is now able to compound

Original: [Figure]

Original: [Figure]
Revised: [Figure]

Original: staying closer to p and a at the same time , while achieving slightly better constraint satisfaction.
Revised: staying closer to p and a at the same time , while achieving slightly better constraint satisfaction.

Original: Wehave also observed that baseline stabilizes training, leading to smoother curves.
Revised: Wehave also observed that baseline stabilizes training, leading to smoother curves.

Original: 4.4 The effect of baseline across batch sizes
Revised: 4.4 The effect of baseline across batch sizes

Original: We expect that reducingvariance in the gradient estimates  can allow to trainthe models with lowerbatch sizes, performing gradient updates on estimates based on smaller batch sizes can increasethe sample efficiency.
Revised: We expect that reducing gradient estimates variance can allow to train with lower batch sizes,performing gradient updates on estimates based on smaller batch sizes can increase the sampleefﬁciency.

Original: To test this hypothesis, we rerun tasks 1 (a pointwise constraint on the word“amazing") and 8 (a set of distributional constraints on topics) with four batch sizes (256, 512, 1024,2048).
Revised: To test this, we rerun tasks 1 (a pointwise constraint on the word “amazing") and 8 (distributional constraints on topics) with four batch sizes (256, 512, 1024, 2048).

Original: We present the results on Figures 4a and 4b. The benefits of adding a baseline — higherconstraint satisfaction, lower divergence from p , more stable training — are especially evident withlower batch sizes.
Revised: Figures 4a and 4bshow the beneﬁts of adding a baseline — higher constraint satisfaction, lower divergence from p ,more stable training — and is especially evident with lower batch sizes.

Original: For instance, with batch size 256, GDC ++ obtains a significantly higher constraintsatisfaction rate and lower divergence from p .
Revised: For instance, with batch size256, GDC ++ obtains a signiﬁcantly higher constraint satisfaction rate and lower divergence from p .

Original: Furthermore, stable training with smaller batch sizes translates into better sample efficiency.
Revised: Furthermore, stable training with smaller batch sizes translates into better sample efﬁciency.

Original: For instance, in task 1 (Figure 4a), GDC ++ with batch size 256 needs 1M samples to achieve294 E x ∼ π θ ϕ
Revised: For instance, in task 1 (Figure 4a), GDC ++ with batch size 256 needs 1M samples to achieve

Original:  x ) = 0 .
Revised: E x ⇠ ⇡ ✓  ( x ) = 0 .

Original: 5 while GDC ++ with batch size 2048 needs 4M.
Revised: 5 while GDC ++ with batch size 2048 needs 4M.

Original: In contrast, GDC with batch size256 does not achieve E x ∼ π θ ϕ x ) = 0 .
Revised: In contrast, GDC with batch size256 does not achieve E x ⇠ ⇡ ✓  ( x ) = 0 .

Original: 5 at all, confirming the importance of adding the baseline.
Revised: 5 at all, conﬁrming the importance of adding the baseline.

Original: 4.5 Empirical Evaluation of Variance Reduction
Revised: 4.5 Empirical Evaluation of Variance Reduction

Original: Next, we evaluate empirically the effect of the baseline for variance reduction.
Revised: Next, we evaluate empirically the effect of the baseline for variance reduction.

Original: We select two tasks:task 1 (a pointwise constraint) and task 7 (distributional constraints) described in Section 4.2, eachwith 3 different seeds, while monitoring the following variance measures:
Revised: We select two tasks:task 1 (a pointwise constraint) and task 7 (distributional constraints) described in Section 4.2, eachwith 3 different seeds, while monitoring the following variance measures:

Original: Gradient Variance
Revised: Gradient Variance

Original: The gradient estimate is defined as: G θ ( x ) .
Revised: The gradient estimate is deﬁned as: G ✓ ( x ) .

Original: =
Revised: =

Original: A ( x ) ∇ θ log π θ ( x ) , where
Revised: A ( x ) r ✓ log ⇡ ✓ ( x ) , where

Original: G θ ( x ) ∈ R | θ | is an unbiased estimate of the gradient of the forward KL loss ∇ θ D KL ( p, π θ ) with respect to the parameters θ .
Revised: G ✓ ( x ) 2 R | ✓ | is an unbiased estimate of the gradient of the forward KL loss r ✓ D KL ( p, ⇡ ✓ ) with respect to the parameters ✓ .

Original: We then have, with µ ( G θ ) .
Revised: We then have, with µ ( G ✓ ) .

Original: = E x ∼ q G θ ( x ) :
Revised: [Equation]

Original: [Figure] [Equation] [Figure]
Revised: [Equation] [Equation]

Original: Variance of the advantage This is definedby:
Revised: Variance of the advantage deﬁned by:

Original: [Equation]
Revised: [Equation]

Original: [Equation] where, µ A ≡ E x ∼ q A( x ) is the mean of theadvantage, which we showed above to be nullafter the addition of the baseline [Figure]
Revised: where, µ A ⌘ E x ⇠ q A( x ) is the mean of the advantage, which we showed above to be null after theaddition of the baseline.

Original: Expected absolute value of the advantage This metric is defined as:
Revised: Expected absolute value of the advantage Thismetric is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: It directly provides a standard measure of distributional discrepancy between p and π θ , in termsof TVD (Total Variation Distance).
Revised: It directly provides a standard measure of distributional discrepancy between p and ⇡ ✓ , in terms of TVD (Total Variation Distance).

Original: Indeed wehave:
Revised: We have:

Revised: E x ⇠

Original: [Equation]
Revised: [Equation]

Original: Results  Figure 5shows that GDC ++ obtains lower variance in the gradient estimates Var( G θ ) and the variance of the advantage
Revised: Results In Figure 5 GDC ++ obtains lower variancein the gradient estimates Var( G ✓ ) and the varianceof the advantage

Original: Var (A) in both pointwise and distributional experiments compared to its non-baseline counterpart GDC.
Revised: Var (A) in both pointwise and distributional experiments compared to its non-baselinecounterpart GDC.

Original: We further observe a decreasingtrend in the mean absolute value of the advantage µ | A | which is correlated with a decreasingtrend in the TVD distance between the trainedpolicy π θ and the optimal distribution p .
Revised: We further observe a decreasingtrend in the mean absolute value of the advantageµ | A | which is correlated with a decreasing trend inthe TVD distance between the trained policy ⇡ ✓ andthe optimal distribution p .

Original: Overall, these results shows that adding a baseline to DPGreduces the variance during training and yields better convergence towards the optimal distribution p .
Revised: Overall, these results shows that adding a baseline to DPG reduces thevariance during training and yields better convergence towards the optimal distribution p .

Original:  Conclusion
Revised:  Conclusion

Revised: Fine-tuning large language models has become an active area of research, due to its importance inadapting large language models to satisfy task-level preferences, or in combating their social riskssuch as “distributional” stereotyping (Weidinger et al., 2021;

Revised: Welbl et al., 2021) 9 .

Original: In this paper, we analyzed the nuanced relation between  RM and DM approaches to fine-tuninglanguage models: we demonstrated that KL-control can be seen as a form of DM and showedthat while DPG and PG have different goals, some similarities (similar forms of gradient estimatesdespite different objectives) can be exploited.
Revised: In this paper, weanalyzed in depth the nuanced relation between two popular ﬁne-tuning paradigms: RM and DM. Wedemonstrated that KL-control can be seen as a form of DM and showed that while DPG and PG havedifferent goals, some similarities (similar forms of gradient estimates despite different objectives) canbe exploited.

Original: We used these insights to inform an extension of DPG,consisting in adding a baseline to reduce the variance of gradient estimates.
Revised: We used these insights to inform an extension of DPG, consisting in adding a baselineto reduce the variance of gradient estimates.

Original: The connections we established suggest that despite fundamental differences between DPG and RL, at least some of the theoretical results and algorithmic techniques from RL can be adapted to a DM framework without losing their formal guarantees.
Revised: The connections we established suggest that despite fundamental differences between DPG and RL,some of the theoretical results and algorithmic techniques from RL can be adapted to a DM frameworkwithout losing their formal guarantees.

Original: In this paper, we have focused on variancereduction using baselines, but the space of possible enhancements is vast.
Revised: In this paper, we focus on variance reduction using baselines,but the space of possible enhancements is vast.

Original: Promising candidatesinclude further reducing the variance using a learned value function (Konda & Tsitsiklis, 2000) andpreventing detrimentally large policy updates by maintaining a trust region in the policy space –akin to techniques such as trust-region policy optimisation (Schulman et al., 2015) and proximalpolicy optimisation (Schulman et al., 2017b).
Revised: Promising candidates include further reducing thevariance using a learned value function (Konda & Tsitsiklis, 2000) and preventing detrimentally largepolicy updates by maintaining a trust region in the policy space – akin to techniques such as TRPO(Schulman et al., 2015) and PPO (Schulman et al., 2017b).

Original: Another future direction could consist in analyzing therelation between explicit EBMs in DPG and implicit EBMs arising in KL-control and characterizingthe space of EBMs that could be reached through KL-control.Checklist For all authors...
Revised: Another future direction could consist inanalyzing the relation between explicit EBMs in DPG and implicit EBMs arising in KL-control andcharacterizing the space of EBMs that could be reached through KL-control.

Original: [Table]

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?

Original: [Yes]

Original: [Table]

Original: (b) Did you describe the limitations of your work?

Original: [Yes]

Original: In Section 5.(c)

Original: Did you discuss any potential negative societal impacts of your work?

Original: [Yes] SeeAppendix A.(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] In Appendix C wepresent proofs of all mathematical facts referred to in the paper.

Original: [Table]

Original: If you ran experiments...

Original: [Table]

Original: (a) Did you include the code, data, and instructions needed to reproduce the main exper-

Original: [Table]

Original: imental results (either in the supplemental material or as a URL)?

Original: [Yes] The code is

Original: included as supplementary material available to the reviewers and area chairs and will

Original: be made publicly available alongside the camera ready version of the paper.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] In Appendix E we provide the hyperparameters used throughoutour experiments and report our hardware configuration.

Original: In Appendix D, we describe

Original: in detail how D KL ( p, π θ ) and TVD( p, π θ ) were estimated and provide an extended

Original: pseudocode for our training loop in Algorithm 2.

Original: [Table]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [No] However, we found the variance across random seeds

Original: to be negligible and not comparing across random seeds is a standard practice whenworking with large language models where the cost of a single run is significant.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g.,

Original: type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] In Appendix E we report ourhardware configuration.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] In Appendix E.

Original: (b) Did you mention the license of the assets?

Original: [Yes] In Appendix E.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?

Original: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]