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)

Original: 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 “ﬁnetuning” 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-trained distribution towards preferred outputs, in others it is to steer it towards a different distribution 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 this challenge: 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 ﬁrstmake explicit the target distribution that the model is ﬁne-tuned to approximate.
Revised: DM prescribes to ﬁrst make 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, and show that methods such as KL-control developed for RM can also be construed as belonging to DM.

Original: We further observe that while DM differs from RM, it cansuffer from similar training difﬁculties, such as high gradient variance.
Revised: We further observe that while DM differs from RM, it can suffer 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 leverage connections between the two paradigms to import the concept of baseline into DM methods.

Original: 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.
Revised: We empirically validate the beneﬁts of adding a baseline on an array of controllable language generation tasks such as constraining topic, sentiment, and gender distributions in texts sampled from a language model.

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

Original:  Introduction
Revised:  Introduction

Original: [Equation]

Original: Parametricreward

Original: -trained language models (Devlin et al. Radford et al., 2019) are changing thelandscape of Machine Learning research andpractice.
Revised: Pre-trained language models (Devlin et al., 2019; Radford et al., 2019) are changing the landscape of Machine Learning research and practice.

Original: 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.
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 sequences instead of training from scratch using annotated data.

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),reducing distributional biases (Khalifa et al., 2021;
Revised: 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), reducing distributional biases (Khalifa et al., 2021;

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

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

Original: Reward Maximization RM intuitively nudge pre-trained models towards some preferences byproviding global sequence-level rewards when the model generates outputs that satisfy desiredfeatures.
Revised: Reward Maximization RM intuitively nudges pre-trained models towards certain preferences by providing global sequence-level rewards when the model generates outputs that satisfy desired features.

Original: For instance, if the model is producing toxic content, we can apply Reinforcement Learning(RL) techniques to discourage it from producing similar content.
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, naively applying RL yieldmodels that  undergo catastrophic forgetting of its original distribution.
Revised: However, naively applying RL yields a model that can undergo catastrophic forgetting of its original distribution.

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

Original: 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).
Revised: Although several studies have considered 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 nontrivial (Wu et al., 2016; Vedantam et al., 2015).

Original: This has sparked a wide discussion on the overall effectiveness of RMfor some tasks such as NMT (Choshen et al., 2020; Kiegeland & Kreutzer, 2021).
Revised: This has sparked a wide discussion on the overall effectiveness of RM for some tasks such as machine translation (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 “catastrophic forgetting”, several studies, still under an RM paradigm, have considered incorporating a distributional term 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) and Ziegler et al. (2019) or more recently Stiennon et al.

Revised: (2020), Ouyang 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: (2022), Bai et al. and Perez et al. (2022) 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.

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-penalty is maximized using an RL algorithm of choice including: PPO (Schulman et al., 2017a)

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

Original: (2019) or  Q-learning (Mnih et al., 2013)
Revised: (2019) or Bai et al. (2022) 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 of the 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 ﬁne-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 satisfy downstream preferences formulates the problem as Distribution Matching (DM).

Original: This paradigmconsists of two steps: ﬁrst a target distribution incorporating the desired preferences is deﬁned as an Energy-Based Model (LeCun et al., 2006).
Revised: This paradigm consists of two steps: ﬁrst a target distribution incorporating the desired preferences is deﬁned as an 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 Distributional Policy Gradients (DPG) (Parshakova et al., 2019b; Khalifa et al., 2021; Korbak et al.,2021).
Revised: 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, 2022a).

Original: This approach capitalizes on the ﬂexibility 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 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).
Revised: For example, the EBM can be deﬁned 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).

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 de-biasing 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 deﬁnedin the RM paradigm where a reward is calculated for a single sequence.
Revised: Such distributional constraints cannot be deﬁned in the RM paradigm where a reward is calculated for a single sequence.

Original: We can notice the promises and limitations of these two paradigms for ﬁne-tuning language models .
Revised: We can notice the promises and limitations of these two paradigms for ﬁne-tuning language models.

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

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

Original: While the connections between these two seemingly distinct paradigms have been noted (Parshakovaet al., 2019b), they have not been explored in detail.
Revised: While the connections between these two seemingly distinct paradigms have been noted (Parshakova et al., 2019b; Korbak et al., 2022b), they have not been explored in detail.

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

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

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

Original: • 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 parametric and non-parametric rewards that are important in this regard.
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 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 particular choice 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 by applying 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 efﬁciency (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 models to certain preferences has been that of assigning a “reward” score R ( x ) for sequences x sampled from an autoregressive language model (policy) ⇡ ✓ .

Original: 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:
Revised: Then, the simplest policy gradient algorithm in reinforcement learning, namely, REINFORCE (Williams, 1992a), aims to ﬁnd the policy ⇡ ✓ ( x ) that maximizes the average reward E x ⇠ ⇡ ✓ R ( x ) , and this leads, via the so-called “log derivative trick”, to a gradient ascent algorithm that iteratively samples x from ⇡ ✓ and update parameters by increments proportional 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 deviations from the original pretrained model a ( x ) , weighted by a free hyperparameter  to control the trade-off between 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 ﬁtan autoregressive policy ⇡ ✓ to the distribution p ( x ) = P ( x ) /Z induced by the EBM P ( x ) , where
Revised: Distributional Policy Gradients (DPG) (Parshakova et al., 2019b) is a recent approach used to ﬁt an autoregressive policy ⇡ ✓ to the distribution p ( x ) = P ( x ) /Z induced by the EBM P ( x ) , where

Original: Z =
Revised: Z =

Original: P x P ( x ) is the normalization constant (partition function).
Revised: 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 autoregressive model ⇡ ✓ , 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 to a gradient estimate of the form:

Original: [Equation]
Revised: [Equation]

Original: Reward Maximization vsDistribution Matching
Revised: vs

Original: In the previous section, we have summarized three approaches that have been suggested for ﬁne-tuninglanguage models.
Revised: In the previous section, we have summarized three approaches that have been suggested for ﬁne-tuning language 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 Matching” (DM) as it ﬁrst deﬁnes the target distribution and then optimizes a policy tomatch it.
Revised: On the other hand, DPG clearly belongs to the realm of “Distribution Matching” (DM) as it ﬁrst deﬁnes the target distribution and then optimizes a policy to match 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 distinct concepts 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 standard reward 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 ) , we would follow its gradient, leading to the identities:

Original: [Equation]
Revised: [Equation]

Original: Equation (8) is the sum of two terms: the ﬁrst 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), involves the 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 reward consists solely of the PG-term.

Original: 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).
Revised: However, when R ✓ depends on ✓ , the gradients are distinct (apart from 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 ) r ✓ log ⇡ ✓ ( x ) and simplifying the tuning procedure to a standardPolicy
Revised: However, due to the particular form of its objective, the RG-term actually vanishes , 3 leaving only the PG-term E x ⇠ ⇡ ✓ R z ✓ ( x ) r ✓ log ⇡ ✓ ( x ) and simplifying the tuning procedure to a standard

Original:  Gradient.
Revised: Policy 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 its nature is multifaceted, and explore 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 reward with the KL penalty term is equivalent to a distributional matching with an underlying emergent sequential 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: 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 =
Revised: Z P z ( x ) , with Z =

Original: P x P z ( x ) .
Revised: P x P z ( x ) .

Original: Then:
Revised: Then:

Original: (i) arg max ⇡ ✓ E x ⇠ ⇡ ✓ R
Revised: (i) arg max ⇡ ✓ E x ⇠ ⇡ ✓ R

Original: z ✓ ( x ) = arg min ⇡ ✓ D KL ( ⇡ ✓ , p z ) ; (ii) arg max ⇡ 2 D ( X ) E x ⇠ ⇡ R z ⇡ ( x ) =
Revised: z ✓ ( x ) = arg min ⇡ ✓ D KL ( ⇡ ✓ , p z ) ; (ii) arg max ⇡ 2 D ( X ) E x ⇠ ⇡ R z ⇡ ( x ) =

Original: p z , where D ( X ) is the family of all distributions over X , and R z ⇡ ( x ) .
Revised: p z , where D ( X ) is the family of all distributions over X , 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 : [Equation] i [Equation]

Revised:

Original: [Equation]
Revised: [Equation]

Revised: ⇡ ✓ r ✓ log ⇡ ✓

Original: [Equation]
Revised: [Equation]

Revised: x ) = 0 , via the identity E ⇡ ✓ r ✓ log ⇡ ✓

Revised: [Equation]

Revised: P

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Original:  so that the arg min ⇡ ✓ D KL ( ⇡ ✓ , p z ) coincides with the arg max ⇡ ✓ E x ⇠ ⇡ ✓ R z ✓ ( x ) , proving (i).
Revised: [Equation] [Table] [Equation] 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 ⇡ 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.
Revised: On the 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 reward does indeed provide a DM interpretation, namely in terms of minimizing the reverse KL divergence with 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 free and explicit choice of EBM (e.g. one that balances the gender and topic distributions of a language model).

Original: Instead equation (9) appears in a restrictive format, which is implicitly deﬁned 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 speciﬁcation, corresponding to an explicit distributionalobjective.
Revised: By contrast, the DPG algorithms are designed to perform DM on any EBM speciﬁcation, corresponding to an explicit distributional objective.

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 deﬁning R ✓ = P ( x )
Revised: We begin by noting that after deﬁning R ✓ = P ( x )

Original: ⇡ ✓ ( x ) , the DPG gradient E x ⇠ ⇡ ✓ P ( x ) ⇡ ✓ ( x ) r ✓ log ⇡ ✓ ( x ) acquires the format of the PG-term E ⇡ ✓ R ✓ r ✓ log ⇡ ✓ ( x ) .
Revised: ⇡ ✓ ( 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: ( x ) would require also the RG-term E r

Original: P ✓ ( x the gradient,
Revised: R ✓ ( x ) =

Original: DPG [Table]
Revised: P ( x ) ⇡ ✓ ( x ) , as this would require adding also the RG-term E ⇡ ✓ r ✓ P ( x ) ⇡ ✓ ( x ) into the gradient, which in this case does not vanish.

Original: Nonetheless, the analogy behind this gradient term is more fruitful than it ﬁrst 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 matter of 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 ) r ✓ 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 (reward maximization), 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 ﬁxed ✓ , we can deﬁne 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 ﬁxed ✓ is identical to the estimation E x ⇠ ⇡ ✓ R ( x ) r ✓ log ⇡ ✓ ( x )based on a set of samples x in standard RL.
Revised: = R ✓ and the problem 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 ✓ .
Revised: Therefore, techniques that have been developed to reduce the variance of the gradients estimates in RL can be ported to DPG insofar as we are computing the gradient estimates at a given ✓ .

Original: In 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 such variance reduction technique to the DPG: baselines.

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

Original: Baselines are a standard variance reduction technique in thecontext of Policy Gradients (Sutton & Barto, 2018).
Revised: Baselines are a standard variance reduction technique in the context 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 is to 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 the introduction of baseline, equation (1) then takes the following form:

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 the average 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 to be the expectation of the reward term, we can obtain 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 :

Revised: E

Original: [Equation]
Revised: [Equation]

Original: ⇠ ⇡ ✓
Revised: ✓

Original: [Equation]
Revised: [Equation]

Original: Here, the second term does not introduce bias because Z 
Revised: Here, the second term does not introduce bias because Z h P x r ✓ ⇡ ✓ ( x

Original: [Equation]
Revised: [Equation]

Original: i = 0 , leaving us with theexact same form of gradient as in the original DPG algorithm.
Revised: i = 0 , leaving us with the exact same form of gradient as in the original DPG algorithm.

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

Original: 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: ✓ ✓ + ↵ ( ✓ ) h
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: P ( x ) q ( x )  Z ⇡ ✓ ( x ) q ( x ) i r ✓ 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 to be 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 could be now estimated by averaging over samples from all the different ✓ ’s without introducing bias, 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 two forms 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 introduced in (Parshakova et al., 2019b) and its KLadaptive variant (Khalifa et al., 2021) sample a 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: 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 not bias the off-policy DPG gradient estimates.
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 (deﬁned as A .
Revised: In practice, as shown on Figure 2, adding a baseline to KL-adaptive DPG (Algorithm 1) centers the advantage (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 beneﬁts 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 control the 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, follow target distribution form of Parshakova et al.

Original: (2019a) and Khalifa et al.
Revised: (2019a), Khalifa et al.

Original: (2021), in which the EBM
Revised: (2021) and Korbak et al.

Original:  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.
Revised: (2022a), in which the EBM P ( x ) is deﬁned so that its normalized variant p ( x ) matches a set of desired moments constraints on given features  i ( x ) , while having a minimal KL divergence D KL ( p, a ) from an original pretrained language model a , to avoid catastrophic forgetting.

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

Original: E
Revised: E

Original: x
Revised: x

Original: ⇠ p
Revised: ⇠ p

Original: 
Revised:

Original: 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 .
Revised: i ( x ) , for i 2 { 1 , . . . , n } , by which the moments (expectations) under the distribution p of each feature  i ( x ) are required to take certain desired values ¯ µ i .

Original: For instance, let  1 ( x ) = 

Revised: For instance, let  1 ( x ) =

Original: iff the topic of x is science and  2 ( x ) = 

Revised: 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 x mentions 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 language model 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 uniquely determined by the following form:[-
-]

Original: [Equation]
Revised: [Equation]

Original: [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).
Revised: where  i terms control the moments µ i of the associated features, which can be estimated through self-normalized importance sampling (Owen, 2013); and then, to make the moments match the desired 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 evaluate our method on an array of 10 controlled text generation tasks.
Revised: We evaluate our method on an array of 10 controlled text generation tasks.

Original: 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 ) .
Revised: For each, given a pretrained language model a ( x ) , and a set of constraints, the objective of each ﬁne-tuning method is to obtain a ﬁne-tuned language model ⇡ ✓ that satisﬁes the imposed constraints while deviating as minimally as possible from the original language model a ( x ) .

Original: 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.
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: 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):
Revised: Based on the value of the moment 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 experiment with frequent words (task 1: “amazing”, original frequency: 10  4 ) and (task2: “WikiLeaks”, original frequency: 10   rare words (b) Wordlist constraints, where  ( x ) = 

Revised: We experiment with frequent words ( task 1: “amazing”, original frequency: 10  4 ) and (task 2: “WikiLeaks”, original frequency: 10  5 ) rare words, (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 byDathathri et al.
Revised: We consider lists of word associated with politics (task 3) and science (task 4) published by Dathathri et al.

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

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

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

Original: (2020).(d) A single distributional constraint where  ( x ) = 1{+
+} iff x contains a female ﬁgure mention,and ¯ µ = 0 .
Revised: (d) A single distributional constraint where  ( x ) = {+
+} iff x contains a female ﬁgure mention, and ¯ µ = 0 .

Original: 5 (task 8), (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.
Revised: 5 (task 8), (e) A set of four distributional constraints:  i ( x ) = [-
-]1 iff x contains at least one of the words in the “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) Hybrid constraints where  1 ( x ) {+= 
+}
Revised: 25 (task 8), (f) Hybrid constraints where  1 ( x ) {+= 
+}

Original: [-= 
-] 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 ) = 1 iff x contains at least one of the words from the “sports" wordlist(task 9) or “politics” wordlist, ¯ µ 2 ( x ) = 1 (task10).
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

Original: Methods We modify the GDC framework Khalifa et al.
Revised: Methods We modify the GDC framework Khalifa et al.

Original: (2021) namely its KL-DPG algorithmto include a baseline as shown in Algorithm 1.
Revised: (2021), namely its KL-DPG algorithm, to include a baseline as shown in Algorithm 1.

Original: We refer to this method as GDC ++ .
Revised: We refer to this method as GDC ++ .

Original: In addition tocomparing GDC ++ with GDC we compare with two reward maximization baselines: Reinforce(Williams, 1992b) and Ziegler (Ziegler et al., 2019).
Revised: In addition to comparing GDC ++ with GDC we compare with two reward maximization baselines: Reinforce (Williams, 1992b) and Ziegler (Ziegler et al., 2019).

Original: Reinforce tries to maximize the expected reward 
Revised: Reinforce tries to maximize the expected reward E x

Original: E x ⇠ ⇡ ✓ R ( x ) , where R ( x ) = 1 iff{+
 if and only if+} the pointwise constraints are met.
Revised: ⇠ ⇡ ✓ R ( x ) , where R ( x ) = {+
 if and only if+} the pointwise constraints are met.

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

Original: 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.
Revised: Following (Khalifa et al., 2021), for hybrid and 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 at each validation step over batches of samples from ⇡ ✓ :
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 -thfeature.2.
Revised: E x ⇠ ⇡ ✓  i ( x ) , measuring the ability to reach the target moment of the i -th

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

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 a single sample x , proposed by (Li et al., 2016a).

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 proposed by (Zhu et al., 2018), ensuring that policies don’t converge into limited number of sequences that 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 (Radford et al., 2019) as the original language model a .

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

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

Revised: The code for all the experiments presented in the paper will be available at github.com/naver/gdc.

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 tasks 1-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 presented separately 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 intosigniﬁcantly decreased diversity of generated samples (in terms of Self-BLEU-5 and Distinct-1).The
Revised: Consistent with prior work (Khalifa et al., 2021; Korbak et al., 2022a), we observe that Reinforce is able to quickly achieve high levels of constraint satisfaction, but at the cost of large deviations from a , which translates into signiﬁcantly decreased diversity of generated samples (in terms of Self-BLEUand Distinct-1).

Original:  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.
Revised: The KL penalty term in Ziegler imposes an upper bound on deviation from a but the deviation is still signiﬁ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 to result in very unstable training.

Original: 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 .
Revised: GDC and GDC ++ are the only ﬁne-tuning methods that address constraint satisfaction based on a clear formal objective, i.e. reducing the divergence from p .

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

Original: The addition ofa baseline indeed reduces the variance.
Revised: The addition of a 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 here focusing on the downstream effects of variance reduction.

Original: One is that ⇡ ✓ is now able to compound[Figure] staying closer to p and a at the same time , while achieving slightly better constraint satisfaction.
Revised: One is that ⇡ ✓ is now able to compound 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: We have 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 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.
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 sample [Figure]

Original:  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).
Revised: efﬁciency. 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: 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.
Revised: Figures 4a and 4b show 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 size256, GDC ++ obtains a signiﬁcantly higher constraint satisfaction rate and lower divergence from p .
Revised: For instance, with batch size 256, 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 efﬁciency.
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 achieve
Revised: For instance, in task 1 (Figure 4a), GDC ++ with batch size 256 needs 1M samples to achieve

Original: E x ⇠ ⇡ ✓  ( 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 size 256 does not achieve E x ⇠ ⇡ ✓  ( x ) = 0 .

Original: 5 at all, conﬁrming 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, each with 3 different seeds, while monitoring the following variance measures:

Original: Gradient Variance
Revised: Gradient Variance

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

Original: =
Revised: =

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

Original: 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 ✓ .
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: [Equation]
Revised: =

Revised: E x ⇠ q G ✓

Original: [Equation]
Revised: [Equation]

Revised: x ) :

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: where, µ A ⌘ E x ⇠ q A( x ) is the mean of the advantage, which we showed above to be null after the addition of the baseline.

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

Original: [Equation]
Revised: [Equation]

Original: It directly provides a standard measure of distributional discrepancy between p and ⇡ ✓ , in terms of
Revised: It directly provides a standard measure of distributional discrepancy between p and ⇡ ✓ , in terms of

Original: TVD (Total Variation Distance).
Revised: TVD (Total Variation Distance).

Original: We have:
Revised: We have:

Original: E x ⇠

Original: [Equation]
Revised: [Equation]

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

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

Original: 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 .
Revised: We further observe a decreasing trend in the mean absolute value of the advantage µ | A | which is correlated with a decreasing trend in the TVD distance between the trained policy ⇡ ✓ and the optimal distribution p .

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

Revised: The idea of posing control problems as distribution matching has resurfaced numerous times in the RL literature (Kappen et al., 2012; Friston et al., 2010; Levine, 2018; Hafner et al., 2020; Buckley et al., 2017).

Original: Conclusion
Revised: KL-control can be seen as a generalisation of maximum entropy RL (MaxEnt RL) (Haarnoja et al., 2017, 2018) to informed priors. If in (2) we chose a ( x ) to be a uniform distribution (assuming right now ﬁniteness of X ) instead of a pretrained LM distribution, then the KL penalty D KL ( ⇡ ✓ , a ) would reduce to an entropy bonus. Both KL-control and MaxEnt RL can be derived from a general framework of control-as-inference (Levine, 2018) which poses control as minimising KL from a certain target distribution. However, most practical algorithms in the MaxEnt RL family minimise KL from a target policy which changes throughout training; in contrast, DPG’s target distribution p and

Revised: KL-control implicit target distribution p z are deﬁned at trajectory level and ﬁxed throughout training.

Revised: Perhaps the closest method to KL-control and DPG in the larger family of inference-based RL (Furuta et al., 2021) is AWR (Peng et al., 2019) which minimises the forward KL from an off-policy target distribution.

Revised: Yet another approach with apparent similarity to KL-control and DPG is state marginal matching (SMM) (Hazan et al., 2018; Lee et al., 2019).

Revised: SMM poses exploration as learning a policy that induces a state marginal distribution that matches a target state distribution.

Revised: While SMM’s target distribution is ﬁxed, it is deﬁned for individual states, while in the controllable language generation tasks we consider, the target distribution is deﬁned over a complete trajectory considered as a unit.

Revised: See Appendix B for an extended discussion of related work.

Original: 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: Fine-tuning large language models has become an active area of research, due to its importance in adapting large language models to satisfy task-level preferences, or in combating their social risks such as “distributional” stereotyping (Weidinger et al., 2021;

Original: Welbl et al., 2021) 9 .
Revised: Welbl et al., 2021).

Original: In this paper, weanalyzed in depth the nuanced relation between two popular ﬁne-tuning paradigms: RM and DM.
Revised: In this paper, we analyzed in depth the nuanced relation between two popular ﬁne-tuning paradigms: RM and DM.

Original: 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.
Revised: We demonstrated that KL-control can be seen as a form of DM and showed that while DPG and PG have different goals, some similarities (similar forms of gradient estimates despite different objectives) can be exploited.

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

Original: 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.
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 framework without losing their formal guarantees.

Original: In this paper, we focus on variance reduction 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 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).
Revised: Promising candidates include further reducing the variance using a learned value function (Konda & Tsitsiklis, 2000) and preventing detrimentally large policy 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 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.
Revised: Another future direction could consist in analyzing the relation between explicit EBMs in DPG and implicit EBMs arising in KL-control and characterizing the space of EBMs that could be reached through KL-control.