Year: 2021

Abstract:
The availability of large pre-trained models is changing the landscape of Machine Learning research and practice, moving from a “training from scratch” to a “finetuning” paradigm. 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. Two main paradigms have emerged to tackle this challenge: Reward Maximization (RM) and, more recently, Distribution Matching (DM). RM applies standard Reinforcement Learning (RL) techniques, such as Policy Gradients, to gradually increase the reward signal. DM prescribes to first make explicit the target distribution that the model is fine-tuned to approximate. Here we explore the intimate connections between the two paradigms, and show that methods such as KL-control developed in the RM paradigm can also be construed as belonging to DM. We further observe that while DM differs from RM, it can suffer from similar training difficulties, such as high gradient variance. We leverage connections between the two paradigms to import the concept of baseline into DM methods. We empirically validate the benefits 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. We observe superior performance in terms of constraint satisfaction, stability and sample efficiency.

text: 1 INTRODUCTION
∇!𝔼"!𝑅!(𝑥) ∇!𝐷#$ 𝑝, 𝜋!
Reward Maximization
Distribution Matching
?
𝔼!!∇"𝑅" 𝑥 + 𝔼!!𝑅" 𝑥 ∇" log 𝜋"(𝑥)
Parametric reward
Figure 1: In this study we make connection between two popular paradigms for fine-tuning sequence generation models according to preferences Reward Maximization (RM) and Distribution Matching (DM).
Pre-trained language models (Devlin et al., 2019; Radford et al., 2019) are changing the landscape of Machine Learning research and practice. Due to their strong generative capabilities many studies have found it sufficient to “nudge” these models to conform to global preferences defined over the generated sequences instead of training from scratch using annotated data. These preferences could include topic and sentiment (Dathathri et al., 2020), valid musical notes and molecular structures (Jaques et al., 2017a), code compilability (Korbak et al., 2021), balancing gender bias (Khalifa et al., 2021), eval-
uation metrics for Machine Translation and Summarization (Ranzato et al., 2016; Bahdanau et al., 2016), or direct human feedback (Ziegler et al., 2019; Stiennon et al., 2020). This large body of studies is driven by two paradigms: Reward Maximization (RM) and Distribution Matching (DM).
Reward Maximization RM exploits the intuitive notion that we can nudge pre-trained models towards some preferences by providing global sequence-level rewards when the model generates outputs that satisfy desired features. For instance, if the model is producing toxic content, we can apply Reinforcement Learning (RL) techniques to discourage it from producing similar content again in the future. However, the risk of naively applying RL is that the model can undergo catastrophic forgetting of its original distribution. For example, it can degenerate into producing a single nonsensical but at least nontoxic sequence. Although several studies have considered hand-crafting general rewards to ensure desirable features like fluency (Liu et al., 2016a; Tambwekar et al., 2019), coming up with
rewards that are not incomplete or imperfect is highly non-trivial (Wu et al., 2016; Vedantam et al., 2015). These challenges have sparked a wide discussion on the overall effectiveness of RM for some tasks such as Neural Machine Translation (Choshen et al., 2020; Kiegeland & Kreutzer, 2021).
Reward Maximization with KL-Control To tackle the aforementioned issues of “catastrophic forgetting”, several studies, still under an RM paradigm, have considered incorporating a distributional term inside the reward to be maximized. In particular Jaques et al. (2017b; 2019); Ziegler et al. (2019) and Stiennon et al. (2020) have applied variations of KL-control (Todorov, 2007; Kappen et al., 2012) which adds a penalty term to the reward term so that the resulting policy does not deviate too much from the original one in terms of KL-divergence. The overall objective with the KL-penalty is maximized using an RL algorithm of choice including: PPO (Schulman et al., 2017a) as in Ziegler et al. (2019) or Q-learning (Mnih et al., 2013) as in Jaques et al. (2017b). Adding this distributional KL-penalty to the reward raises some important questions: What effect does it have on the shape of the optimal policy? Does this new objective have any interpretation from a distributional perspective?
Distribution Matching A different recent paradigm for fine-tuning language models to satisfy downstream preferences formulates the problem as Distribution Matching (DM). This paradigm consists of two steps: first a target distribution incorporating the desired preferences is defined as an Energy-Based Model (LeCun et al., 2006). Then the forward KL divergence is minimized between this target distribution and an auto-regressive policy using a family of algorithms referred to as Distributional Policy Gradients (DPG) (Parshakova et al., 2019b; Khalifa et al., 2021; Korbak et al., 2021). This approach capitalizes on the flexibility of EBMs in specifying the target distribution. For example, the EBM can be defined so that it conforms to all downstream preferences while its corresponding normalized distribution has a minimal KL divergence from the original, pre-trained language model, therefore tackling the problem of “catastrophic forgetting” (Khalifa et al., 2021). Interestingly, this DM paradigm can also deal with distributional preferences, for instance, for debiasing language models by specifying that the generated sequences should be gender-balanced, i.e. that 50% of generations contain female mentions. Such distributional constraints cannot be defined in the RM paradigm where a reward is calculated for a single sequence.
Contrasting these two paradigms for fine-tuning language models, we can notice the promises and limitations of each.1 RM approaches are equipped with a large arsenal of RL algorithms and optimization techniques that can be efficient in reward maximization, however they lack the distributional perspective that is necessary for avoiding catastrophic forgetting and imposing distributional preferences over LMs. DM approaches are better suited to tackle those limitations, however, the family of DPG algorithms currently used is not as rich as its RL counterpart. So far, the connections between these two seemingly distinct paradigms have not been thoroughly explored. By establishing such connections we might import ideas from one approach to the other. This is our goal in this paper, detailing the nuanced connections and applying them to a case-study in variance reduction. Overall, our contributions are the following:
• We untangle the connections between the RM and DM paradigms for fine-tuning language models. We provide a detailed comparison between the family of DPG algorithms with Policy Gradients of standard RL.
• We provide an interpretation of KL-control techniques from a distribution matching perspective, placing such techniques at an intermediate place between RM and DM.
• We exploit these connections to theoretically justify applying baselines — a variance reduction technique from RL — to DPG and derive a particular choice of a baseline. On an array of controllable language generation experiments about constraining topic, sentiment, and gender distributions over a pre-trained language model, we show that adding baselines leads to superior performance on constraint satisfaction, stability on small batch sizes, and sample efficiency.

heading: 2 BACKGROUND
text: Standard Policy Gradients One popular method for adapting the behaviour of language models to certain preferences has been that of assigning a “reward” score R(x) for sequences x sampled
1See Appendix A for an extended discussion of related work.
from an autoregressive language model (policy) πθ. Then, the simplest policy gradient algorithm in reinforcement learning, namely, REINFORCE (Williams, 1992a), aims to find the policy πθ(x) that maximizes the average reward Ex∼πθR(x), and this leads, via the so-called “log derivative trick”, to a gradient ascent algorithm based on:
∇θEx∼πθR(x) = Ex∼πθR(x)∇θ log πθ(x), (1)
which iteratively samples x from πθ and update parameters by increments proportional to R(x)∇θ log πθ(x).
KL-control Jaques et al. (2017b; 2019); Ziegler et al. (2019), inspired by KL-control (Todorov, 2007; Kappen et al., 2012), add a KL penalty term to the reward objective to penalize large deviations from the original pretrained model. That is, they maximize the expectation Ex∼πθRzθ(x), where:
Rzθ(x) . = r(x)− β log πθ(x)
a(x) , (2)
and where β is a free hyperparameter balancing the wish to maximize r with that of not deviating from a.
Distributional Policy Gradients Distributional Policy Gradients (DPG) (Parshakova et al., 2019b) is a recent approach used to fit an autoregressive policy πθ to the distribution p(x) = P (x)/Z induced by the EBM P (x), where Z = ∑ x P (x) is the normalization constant (partition function). Given an arbitrary EBM P (x), DPG optimizes the loss functionDKL(p, πθ) with respect to the parameters θ of an autoregressive model πθ, a loss which is minimized for πθ = p. The KL-divergence minimization objective leads to a gradient estimate of the form:
∇θDKL(p, πθ) =−∇θEx∼p log πθ(x) (3) =− ∑ x p(x)∇θ log πθ(x) = − 1 Z ∑ x P (x)∇θ log πθ(x) (4)
=− 1 Z Ex∼πθ P (x) πθ(x) ∇θ log πθ(x). (5)

heading: 3 REWARD MAXIMIZATION VS DISTRIBUTION MATCHING
text: In the previous section, we have summarized three approaches that have been suggested for fine-tuning language models. Two of them can be characterized as “Reward Maximization” (RM): Standard Policy Gradients (PG) and KL Control. On the other hand, DPG clearly belongs to the realm of “Distribution Matching” (DM) as it first defines the target distribution and then optimizes a policy to match it. In the rest of this section, we will explore connections between these two seemingly distinct concepts and, in the following section, we will exploit them to improve DM-based methods.

heading: 3.1 STANDARD VS. PARAMETRIC REWARDS
text: Let us start with distinguishing between a “parametric reward” Rθ which depends on θ and a standard reward R, which does not. If we wished to maximize the expected parametric reward, EπθRθ(x), we would follow its gradient, leading to the identities:
∇θEx∼πθRθ(x) = ∇θ ∑ x πθ(x)Rθ(x) (6)
= ∑ x πθ(x)∇θRθ(x) + ∑ x Rθ(x)∇θπθ(x) (7)
= ∑ x πθ(x)∇θRθ(x) + ∑ x πθ(x)Rθ(x)∇θ log πθ(x) (8) = Ex∼πθ∇θRθ(x)︸ ︷︷ ︸ RG-term +Ex∼πθRθ(x)∇θ log πθ(x)︸ ︷︷ ︸ PG-term . (9)
Equation (9) is the sum of two terms: the first one, the “RG-term" (Reward Gradient term), involves the gradient of the reward. The second one, the “PG-term” (Policy Gradient term), was obtained using the “log derivative trick” and involves the gradient of the policy stricto sensu. In standard RL, where the reward does not depend on θ, the RG-term disappears and the gradient of expected reward consists solely of the PG-term. However, when Rθ depends on θ, the gradients are distinct (apart from specific cases where the RG-term evaluates to 0, as we will see below).

heading: 3.2 KL CONTROL AS DISTRIBUTION MATCHING
text: Adding a KL-penalty term to the reward (as in the case of KL-control) leads to a parametric reward. However, due to the particular form of its objective, the RG-term actually vanishes,2 leaving only the PG-term Ex∼πθRzθ(x)∇θ log πθ(x) and simplifying the tuning procedure to a standard Policy Gradient. While this algorithm falls under the RM paradigm, here we argue that is it multifaceted, and explore deeper connections with the DM paradigm. More precisely, the maximization of the 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. Actually, if we consider the following EBM:
Pz(x) = a(x)e r(x)/β , (10) then its normalized distribution, pz(x) = 1Z Pz(x), with Z = ∑ x Pz(x), is, among all possible distributions, the one that obtains maximal expected reward Rzθ(x). A simple way to prove this fact is to notice that the expectation of the reward Rzθ has a monotonically decreasing relationship with the reverse KL divergence between πθ and pz:
DKL(πθ, pz) = Ex∼πθ log πθ(x)
pz(x) = Ex∼πθ
[ log πθ(x)− log 1
Z a(x)er(x)/β ] = logZ − 1
β Ex∼πθ
[ r(x)− β log πθ(x)
a(x)
] = logZ − 1
β Ex∼πθRzθ(x), (11)
so that the arg minπθ DKL(πθ, pz) coincides with the arg maxπθ Ex∼πθR z θ(x). Provided that the family of distributions πθ is large enough to cover all distributions over X , arg minπθ DKL(πθ, pz) is just pz , which concludes the proof.3
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 pz(x). We note that pz(x) does not correspond to an explicit choice of EBM (e.g. the one that balances the gender and topic distributions of a language model). Instead equation (10) has a limited form implicitly defined by the reward Rzθ , along with a β hyperparameter without a clear meaning. By contrast, the DPG algorithms are designed to perform DM on any EBM specification, corresponding to an explicit distributional objective.

heading: 3.3 SIMILARITIES AND DIFFERENCES BETWEEN DPG AND POLICY GRADIENTS
text: In the previous subsection, we have connected KL-control, a method designed under a RM paradigm, to DM. Now, we turn to the converse question of whether DPG, a DM method, can be connected to RM. We begin by noting that after defining Rθ =
P (x) πθ(x) , the DPG gradient Ex∼πθ P (x) πθ(x) ∇θ log πθ(x) acquires the format of the PG-term EπθRθ∇θ log πθ(x). However, the DM objective of DPG cannot be considered as maximizing the average “reward” Rθ(x) =
P (x) πθ(x) , as this would require adding also the RG-term Eπθ∇θ P (x) πθ(x) into the gradient, which in this case does not vanish.
Nonetheless, the analogy behind this gradient term is more fruitful than it first appears. As a matter of fact, DPG gradient estimates suffer from the same high-variance problems as with standard PG. While the objective of DPG (distribution matching) is different from that of Policy Gradients (reward
2This is because Eπθ∇θR z θ(x) = −β Eπθ∇θ log πθ(x) = 0, via the often used identity Eπθ∇θ log πθ(x) = ∑ x πθ(x)∇θ log πθ(x) = ∑ x∇θπθ(x) = ∇θ ∑ x πθ(x) = 0.
3The optimal policy pz is briefly mentioned in (Ziegler et al., 2019) without reference or derivation. The proof, which we believe to clarify important underlying connections, is ours.
maximization), DPG also needs to estimate the PG-term EπθRθ(x)∇θ log πθ(x) at a given value of θ, using a batch of samples x. For such a fixed θ, we can define provisionally R(x) .= Rθ and the problem of gradient estimation for this fixed θ is identical to the estimation Ex∼πθR(x)∇θ log πθ(x) based on a set of samples x in standard RL. Therefore, the 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 θ.4 In Section 4, we show how one can import one such variance reduction technique to the DPG: baselines.

heading: 4 A CASE STUDY ON VARIANCE REDUCTION
text: Baselines are a standard variance reduction technique in the context of Policy Gradients (Sutton & Barto, 2018). 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. After the introduction of baseline, equation (1) then takes the following form:
∇θEπθR(x) = Eπθ [R(x)−B]∇θ log πθ(x). (12) In standard RL, the simplest form of baseline B is just the average of the rewards for the policy:5
BRL = Ex∼πθR(x). (13) Following the same methodology of taking the baseline as the expectation of the reward term, we can obtain a remarkably simple form of a baseline for DPG:
B = Ex∼πθ P (x) πθ(x) = ∑ x πθ(x) P (x) πθ(x) =
∑
x P (x) = Z. (14)
To confirm that B does not introduce bias to the gradients, let us rewrite the DPG gradient in (5) with the added baseline B = Z:
Ex∼πθ [ Rθ(x)− Z ] ∇θ log πθ(x) = Ex∼πθRθ(x)∇θ log πθ(x)− Z Ex∼πθ ∇θ log πθ(x)
= Ex∼πθRθ(x)∇θ log πθ(x)− Z [∑
x
∇θπθ(x) ] (15)
Here, the second term does not introduce bias because Z [∑ x∇θπθ(x) ]
= 0, leaving us with the same exact form of gradient as in the DPG algorithm.
Note that since BRL depends on θ, it has to be be re-estimated after each gradient update. 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. See Table 1 for a comparison of these two forms of baselines.
4To avoid confusion, note that variance reduction methods for SGD focus on the objective of better estimating the true gradient expectation Eπθ [...], for a fixed θ. The fact that θ will then be updated based on such estimates is orthogonal to this objective.
5While this choice of the baseline is not optimal (see Appendix B.1 for a proof), it is widely used in practice.
Algorithm 1 KL-Adaptive DPG with baseline 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: θ ← θ + α(θ) [ P (x) q(x) − Z πθ(x) q(x) ] ∇θ log πθ(x)
6: if DKL(p||πθ) < DKL(p||q) then 7: q ← πθ
Ensure: πθ
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 πθ. Then, the baseline takes the form
Boff(x) = Z πθ(x)
q(x) , (16)
where the πθ(x)q(x) term is an importance weight correcting for the bias introduced by sampling from q. Similarly to the DPG case, this baseline does not introduce bias (see Appendix B for details about this version of the baseline). In practice, as shown on Figure 2, adding a baseline to KL-adaptive DPG (Algorithm 1) centers the advantage (defined as A .= P (x)q(x) − Z πθ(x) q(x) ) around 0 leading to better performance on: convergence (section 5.3), stability on small batch sizes (section 5.4), and variance reduction (section 5.5).

heading: 5 EXPERIMENTS AND RESULTS

heading: 5.1 GENERATION WITH DISTRIBUTIONAL CONTROL
text: We investigate the benefits of adding a baseline to the DPG algorithm, on the Generation with Distributional Control (GDC) (Khalifa et al., 2021) framework. GDC makes use of DPG to control the properties of pre-trained language models to satisfy certain constraints. In our experiments, follow target distribution form of Parshakova et al. (2019a) and Khalifa et al. (2021), in which the EBM P (x) is defined 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 DKL(p, a) from an original pretrained language model a, to avoid catastrophic forgetting. These constraints are expressed as conditions µ̄i = Ex∼pφi(x), for i ∈ {1, . . . , n}, by which the moments (expectations) under the distribution p of each feature φi(x) are required to take certain desired values µ̄i. For instance, let φ1(x) = 1 iff the topic of x is science and φ2(x) = 1 iff x mentions a female person, then imposing moments µ̄1 = 1 and µ̄2 = 0.5 constrains the language model p to only generate sequences about science, half of which mention females. P (x) is uniquely determined by the following form:6
P (x) = a(x)e ∑n i=1 λiφi(x), (17)
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).

heading: 5.2 EXPERIMENTAL SETUP
text: We demonstrate the benefits of adding a baseline to the family of Distributional Policy Gradients algorithms on an array of controlled language generation tasks for English. For this, we modify the GDC framework Khalifa et al. (2021) namely its KL-DPG algorithm to include a baseline as shown in Algorithm 1. We refer to this method in all the following experiments as GDC++. In addition 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 expected reward Ex∼πθR(x), where R(x) = 1 iff the pointwise constraints are met. Ziegler instantiates the 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 through sets of moment constraints {µ̄i} for binary features {φi}. These include 6 sets of purely pointwise constraints (for which µ̄i = 1) and 4 sets including distributional constraints (0 < µ̄i < 1). We consider the following constraint types:
6For a more precise formulation of this EBM, see (Khalifa et al., 2021).
(a) Single-word constraints, where φ(x) = 1 iff the a given word appears in the sequence x. We experiment both 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) = 1 iff x contains at least one word from a given list. We consider lists of word associated with politics (task 3) and science (task 4) published by Dathathri et al. (2020),
(c) Sentiment classifier constraints, where φ(x) = 1 if x is classified as positive (task 5), or negative (task 6) by a pre-trained classifier published by Dathathri et al. (2020).
(d) A single distributional constraint where φ(x) = 1 iff x contains a female figure mention, and µ̄ = 0.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. (2020)), respectively. For each i, µ̄i = 0.25 (task 8), (f) Hybrid constraints where φ1(x) = 1 iff x contains more female than male pronouns, µ̄1 = 0.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).
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 distributional constraints.
Metrics We report the following metrics evaluated over batches of samples from πθ at each validation step:
1. Ex∼πθφi(x), measuring the ability to reach the target moment of the i-th feature. 2. DKL(p, πθ), the forward KL divergence from the optimal target distribution p,7 3. DKL(πθ, a), the reverse KL divergence from the original pretrained language model a. 4. 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). 5. 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. (2020).
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. For tasks 7-10, a is the same pre-trained model additionally fine-tuned on the WikiBio (Lebret et al., 2016) dataset. See Appendix E for more details.

heading: 5.3 RESULTS
text: 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). Results for each task are presented separately on Figures 7-10 in the Appendix.
Consistent with prior work (Khalifa et al., 2021), 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 significantly decreased diversity of generated samples (in terms of Self-BLEU-5 and Distinct-1). The KL penalty term in Ziegler imposes an upper bound on deviation from a but the deviation is still significant enough to result in a drop in diversity. Moreover, we have observed Ziegler’s objective to result in very unstable training.
GDC and GDC++ are the only fine-tuning methods that address constraint satisfaction based on a clear formal objective, i.e. reducing the divergence from p. The approach translates into significantly smaller deviations from a and maintaining diversity within and across samples. The addition of a baseline indeed reduces the variance. We analyze that extensively in Appendix 5.5 while here focusing on the downstream effects of variance reduction. One is that πθ is now able to compound staying closer to p and a at the same time, while achieving slightly better constraint satisfaction. We have also observed that baseline stabilizes training, leading to smoother curves.8
7See Appendix D for a detailed description of how DKL(p, πθ) is computed. 8The interested reader can compare the large fluctuations of the Ziegler objective to more stable training
curves of GDC , and even more of GDC++ , in the disaggregated curves in Figures 7-10 of the Appendix.

heading: 5.4 THE EFFECT OF BASELINE ACROSS BATCH SIZES
text: We expect that reducing variance in the gradient estimates can allow to train the models with lower batch sizes, performing gradient updates on estimates based on smaller batch sizes can increase the sample efficiency. 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). We present the results on Figures 4a and 4b. The benefits of adding a baseline — higher constraint satisfaction, lower divergence from p, more stable training — are especially evident with lower batch sizes. For instance, with batch size 256, GDC++ obtains a significantly higher constraint satisfaction rate and lower divergence from p.
Furthermore, stable training with smaller batch sizes translates into better sample efficiency. For instance, in task 1 (Figure 4a), GDC++ with batch size 256 needs 1M samples to achieve Ex∼πθφ(x) = 0.5 while GDC++ with batch size 2048 needs 4M. In contrast, GDC with batch size 256 does not achieve Ex∼πθφ(x) = 0.5 at all, confirming the importance of adding the baseline.

heading: 5.5 EMPIRICAL EVALUATION OF VARIANCE REDUCTION
text: Next, we evaluate empirically the effect of the baseline for variance reduction. We select two tasks: task 1 (a pointwise constraint) and task 7 (distributional constraints) described in Section 5.2, each with 3 different seeds, while monitoring the following variance measures:
Gradient Variance The gradient estimate is defined as: Gθ(x) . = A(x)∇θ log πθ(x), where Gθ(x) ∈ R|θ| is an unbiased estimate of the gradient of the forward KL loss ∇θDKL(p, πθ) with
respect to the parameters θ. We then have, with µ(Gθ) . = Ex∼qGθ(x):
Var(Gθ) . = Ex∼q ‖Gθ(x)− µ(Gθ)‖22 (18) = Ex∼q||Gθ(x)||22 − ||µ(Gθ)||22. (19)
1000 2000 epochs
102
103
104
Va r(G
)
Task1: Pointwise
1000 2000 epochs
106
6 × 105
2 × 106
Va r(G
)
Task7: Distributional
method GDC++ GDC
1000 2000 epochs
10 4
3 × 10 5 4 × 10 5
6 × 10 5
2 × 10 4
Va r(A
)
1000 2000 epochs
1.5 × 100 1.6 × 100 1.7 × 100 1.8 × 100 1.9 × 100
2 × 100 2.1 × 100 2.2 × 100
Va r(A
)
1000 2000 epochs
0.00
0.25
0.50
0.75
1.00
|A |
1000 2000 epochs
0.00
0.25
0.50
0.75
1.00
|A |
1000 2000 epochs
0.60
0.65
0.70
0.75
TV D
(p ,
)
1000 2000 epochs
0.24
0.26
0.28
TV D
(p ,
)
Figure 5: Comparison between GDC and GDC++ using a set of Variance diagnosis metrics on two experiments for pointwise and distributional constraints.
Variance of the advantage This is defined by:
Var (A) . = Ex∼q ∥∥A(x)− µA∥∥2 2
(20)
where, µA ≡ Ex∼q A(x) is the mean of the advantage, which we showed above to be null after the addition of the baseline.
Expected absolute value of the advantage This metric is defined as:
µ|A| . = Ex∼q |A(x)|. (21)
It directly provides a standard measure of distributional discrepancy between p and πθ, in terms of TVD (Total Variation Distance). Indeed we have:
Ex∼q ∣∣∣∣p(x)q(x) − πθ(x)q(x) ∣∣∣∣ = 2 TVD(p, πθ). (22) Results Figure 5 shows that GDC++ obtains lower variance in the gradient estimates Var(Gθ) and the variance of the advantage Var (A) in both pointwise and distributional experiments compared to its non-baseline counterpart GDC.
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. Overall, these results support our hypothesis that adding a baseline to DPG reduces the variance during training and yields better convergence towards the optimal distribution p.

heading: 6 CONCLUSION
text: In this paper, we analyzed the nuanced relation between RM and DM approaches to fine-tuning language models: 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. We used these insights to inform an extension of DPG, consisting in adding a baseline to reduce the variance of gradient estimates.
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. In this paper, we have focused on variance reduction using baselines, but the space of possible enhancements is vast. 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 trust-region policy optimisation (Schulman et al., 2015) and proximal policy optimisation (Schulman et al., 2017b). 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.

heading: REPRODUCIBILITY STATEMENT
text: The source code for our experiments was based on the repository that Khalifa et al. (2021) published on GitHub.9 It is available for the reviewers and area chairs and will be made publicly available alongside the camera ready version of the paper. The two pretrained models used in our experiments are available on Hugginface Model Hub: gpt10 and mkhalifa/gpt2-biographies.11 In addition to that, in Appendix E we provide the hyperparameters used throughout our experiments and report our hardware configuration. In Appendix D, we describe in detail how DKL(p, πθ) and TVD(p, πθ) were estimated and provide an extended pseudocode for our training loop in Algorithm 2. Finally, in Appendix B we present proofs of all mathematical facts referred to in the paper.

heading: ETHICS STATEMENT
text: The focus area of this paper — fine-tuning large language models — is aligned with an important line of work on addressing the problem of social bias in large language models (Sheng et al., 2019; Liang et al., 2021). As the training data for large language models consists mainly of crawled user-generated content, a number of factors (from crawling methodology to Internet participation inequalities and moderation practices) leads to an over-representation of certain viewpoints and voices exceeding their prevalence in the general population. This poses a risk of amplifying biases and harms through a language model perpetuating these voices (Bender et al., 2021; Blodgett et al., 2020; Sheng et al., 2019). Numerous problems related to addressing data bias in language generation (e.g. controlling for gender distribution in generated texts) can be naturally posed as generative distributional control (GDC), the framework we focus our experiments on. The distributional character of these data bias problems lies in the fact that desirable properties of generated texts are defined for a collection of samples, not only for individual samples. Our theoretical analyses of reward maximization and distribution matching approaches as well as our algorithmic improvements to the GDC framework — termed GDC++ — are therefore also a contribution to the problem of bias in language models. However, we need to be aware that GDC++ , KL-control as well as controllable language generation techniques in general, can also be diverted to malicious uses such as spreading misinformation or generating harmful content.

heading: A EXTENDED RELATED WORK
text: Reinforcement learning for language generation Most previous attempts at steering language models to conform to global constraints defined over entire sequences have employed reinforcement learning. This includes using Reinforce (Williams, 1992a) for machine translation Ranzato et al. (2016), actor critic (Konda & Tsitsiklis, 2000) for abstractive summarization (Paulus et al., 2018), caption generation (Liu et al., 2016b), dialogue (Li et al., 2016b), and video captioning (Pasunuru & Bansal, 2017). Some approaches (for instance, in machine translation and summarization (Ranzato et al., 2016; Bahdanau et al., 2017)) directly optimize performance metrics such as BLEU and ROUGE at training time. Others use heuristic rewards (for instance Li et al. (2016b) for dialogue generation and Tambwekar et al. (2019) for story generation) in order to obtain certain a priori desirable features of generated sequences that then incentivize good performance on target metrics. Catastrophic forgetting is a frequent problem of these fine-tuning approaches: reward maximization happens at the expense of large deviations from the original model. This problem is sometimes addressed by imposing a penalty term to the rewards, such as the KL divergence between the trained policy and the auto-regressive model. This approach, termed “conservative fine-tuning", was applied to generating melodies with music theory rewards and organic molecules with synthesizability rewards by Jaques et al. (2017a) as well fine-tuning language models for controllable language generation by Ziegler et al. (2019). This solution often has hard time balancing between the reward term and the KL penalty term, leading to instability in training (Khalifa et al., 2021). Unlike this approach, KL-DPG determines an optimal distribution that satisfies both requirements.
RM and DM objectives in control problems While RM is the dominant approach to tackling control problems (Sutton & Barto, 2018) and is sometimes argued to be sufficient for any intelligent behavior (Silver et al., 2021), prior work explored the benefits of alternative objectives formulated as DM: minimizing divergence from some target distribution p. Prominent examples of (families of) DM objectives for control include active inference (Friston et al., 2010; Buckley et al., 2017) and control-as-inference (Kappen et al., 2012; Todorov, 2007; Levine, 2018). Hafner et al. (2020) propose a reverse KL from a joint distribution over observations and latent variables as a universal objective for action and perception that — depending on a choice of the target p — gives rise to many familiar objectives, including empowerment (Klyubin et al., 2005), maximum entropy RL (Haarnoja et al., 2017) or KL-control (Todorov, 2007). In a similar vein, Millidge et al. (2021) compare RM and DM objectives (or, evidence and divergence objectives, according to their terminology) in the context of exploration. They conclude that information-seeking exploration arises naturally in DM but not in RM. This is because, when the target distribution p involves latent variables, a DM objective decomposes into an information gain term that pushes the agent to seek observations that are most informative of latent variables. In contrast, RM objectives entail minimizing information gain between latent variables and observations.
Baselines in Reinforcement Learning In the context of reinforcement learning, baselines were introduced by Sutton (1984). Williams (1987; 1992a) has shown them to reduce variance in a number of use cases and also proved that they do not introduce bias. Dayan (1990) was the first to observe and confirm experimentally that the optimal constant baseline is not equal to expected reward in a simple two-arm bandit setting. This result was generalized to POMDPs (Partially Observable Markov Decision Processes) by Weaver & Tao (2001, section 3.1.3, p. 540) and variable baselines by Greensmith et al. (2004, theorem 13, p. 1489) who also proved bounds on the variance of gradient estimates. The optimal baseline, however, is rarely used in practice (Sutton & Barto (2018); for an exception, see (Peters & Schaal, 2008)). Outside RL, baselines were also used in the context of learning inference networks for amortized variational inference by Mnih & Gregor (2014) and found to yield similar variance reduction.
Energy-based models for language Energy-based models (EBMs) (Hinton, 2002; LeCun et al., 2006; Ranzato et al., 2007) are a family of models in which learning and inference are done by associating an unnormalized probability with each configuration of observed and latent variables. Early examples of EBMs applied to natural language processing include sequence labeling problems (e.g. tagging) exploiting global properties of a sequence (Andor et al., 2016; Belanger & McCallum, 2016). The recent surge of interest in EBMs has not left natural language processing unaffected (see Bakhtin et al. (2020) for a survey). Tu et al. (2020) proposed an energy-based inference networks
for non-autoregressive machine translation while Naskar et al. (2020) use an EBM for reranking candidate translations according to their predicted BLEU scores. Parshakova et al. (2019a) and Deng et al. (2020) augment an autoregressive language models with an additional global factor to obtain a lower perplexity on the training data. Clark et al. (2020) poses non-autoregressive language modeling as training an energy-based cloze task scorer using noise-contrastive estimation (Gutmann & Hyvärinen, 2010). He et al. (2021) obtain better calibration on natural language inference tasks by augmenting and training the classifier jointly with an energy-based model modeling the marginal distribution over samples, again using noise-contrastive estimation. In consequence, the classifier tends to assign more conservative (high-entropy) predictions to high-energy (less likely, possibly out of distribution) samples.

heading: B ADDITIONAL PROOFS

heading: B.1 OPTIMAL BASELINES IN RL
text: Despite its widespread use, the baseline as mean of reward
BRL = Ex∼πθ(x)R(x) (23)
is not the optimal constant baseline for reward maximization objectives in RL. The optimal constant baseline, i.e. one yielding the minimal variance of the gradient, is given by:
B∗ = Ex∼πθ [R(x) (∇θ log πθ(x)) 2 ]
Ex∼πθ [(∇θ log πθ(x)) 2 ]
. (24)
In order to maintain accessibility, in this section, we provide a self-contained derivation of this optimal form of baselines (24) and and connect it to the commonly used form (23).12
First, recall that R(x) is a reward associated with an input x. B is a baseline value subtracted from the reward that does not introduce bias in gradient estimation. Now let’s denote the gradient wrt an individual sample x as Gθ(x) where
Gθ(x) = [R(x)−B]∇θ log πθ(x), (25)
and the estimate of the gradient as
G(θ) = Ex∼πθGθ(x). (26)
Using the general identity var(z) = E[z2]− [Ez]2, the variance of the gradient takes the form:
Var(Gθ) = Ex∼πθ [Gθ(x)2]−G(θ)2 (27)
Now let’s take the gradient of this variance with respect to B and solve to find the baseline form with minimal variance:
dVar(Gθ)
dB =
d
dB Ex∼πθ [(Gθ(x))2]−
d
dB (Ex∼πθ [Gθ(x)])2. (28)
The second term of the right hand side of (28) is equal to zero, since B does not introduce bias into G(θ):
d
dB (Ex∼πθ [Gθ(x)])
2 =
d
dB (Ex∼πθ [(R(x)−B)∇θ log πθ(x)])
2
= d
dB (Ex∼πθ [R(x)∇ log πθ(x)]) 2 = 0.
12The formula for the optimal baseline in (24) was originally proved by Weaver & Tao (2001) but here we provide a simpler proof sketched by Sergey Levine in his slides: http://rail.eecs.berkeley.edu/ deeprlcourse-fa17/f17docs/lecture_4_policy_gradient.pdf
Plugging this back into (28), we obtain:
dVar(Gθ)
dB =
d
dB Ex∼πθ [(Gθ(x))2] = Ex∼πθ [ d
dB
[( R(x)2 +B2 − 2R(x)B ) (∇θ log πθ(x))2 ]] = Ex∼πθ (2B − 2R(x))(∇θ log πθ(x))2
= 2B Ex∼πθ (∇θ log πθ(x))2 − 2 Ex∼πθR(x) (∇θ log πθ(x)) 2 .
Then, solving dVar(Gθ)dB = 0 for B, we obtain the optimal form of the baseline B ∗ as required:
B∗ = Ex∼πθ [R(x) (∇θ log πθ(x)) 2 ]
Ex∼πθ [(∇θ log πθ(x)) 2 ]
. (29)
This can be interpreted as average reward (as in BRL) but weighted by gradient magnitudes (∇θ log πθ(x))2. Moreover, B∗ = BRL is recovered under the condition that the reward R(x) is uncorrelated (a fortiori independent) from (∇θ log πθ(x))2. If that were the case, we would have:
B∗ = Ex∼πθ [R(x) (∇θ log πθ(x)) 2 ]
Ex∼πθ [(∇θ log πθ(x)) 2 ]
(30)
= Ex∼πθ [R(x)] Ex∼πθ [(∇θ log πθ(x)) 2 ]
Ex∼πθ [(∇θ log πθ(x)) 2 ]
(31)
= Ex∼πθ [R(x)] = BRL. (32)

heading: B.2 UNBIASEDNESS OF PG BASELINE
text: Baselines are a standard variance reduction technique in the context of Policy Gradients (Sutton & Barto, 2018). 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. Equation (1) then takes the following form:
∇θEπθR(x) = Eπθ (R(x)−B)∇θ log πθ(x). (33)
To see that B does not introduce bias, we can rewrite (12) as:
Ex∼πθR(x)∇θ log πθ(x)−B Eπθ∇θ log πθ(x) (34) and note that the second term is null because ∑ x πθ(x)∇θ log πθ(x) = ∇θ ∑ x πθ(x) = 0.

heading: B.3 UNBIASEDNESS OF DPG BASELINE
text: Recall that the gradient estimate for DPG (Parshakova et al., 2019a) has the following form:
Ex∼πθ P (x)
πθ(x) ∇θ log πθ(x) (35)
After subtracting a baseline B = Z, it becomes
Ex∼πθ [ P (x) πθ(x) − Z ] ∇θ log πθ(x) = Ex∼πθ P (x) πθ(x) ∇θ log πθ(x)− Z [ Ex∼πθ∇θ log πθ(x) ] (36)
= Ex∼πθ P (x)
πθ(x) ∇θ log πθ(x)− Z [∑ x ∇θπθ(x) ]
(37)
Here, the second term does not introduce bias because Z [∑ x∇θπθ(x) ]
= 0, leaving us with the same exact form of gradient as in the DPG algorithm.

heading: B.4 UNBIASEDNESS OF DPGOFF BASELINE
text: Offline DPG, the off policy variant of DPG proposed in Parshakova et al. (2019b); Khalifa et al. (2021) has the following gradient estimate:
Ex∼q P (x)
q(x) ∇θ log πθ(x) (38)
Where q is a proposal distribution (another auto-regressive model) used to detach the training of πθ from the sampling process and allow more stable training.
Recall that the Baseline of DPGoff is of the form:
Boff(x) = Z πθ(x)
q(x) , (39)
The πθ(x)q(x) term is an importance weight correcting for the bias introduced by sampling from q.
Unbiasedness To show that subtracting a baseline Boff(x) = Z πθ(x)q(x) doesn’t introduce bias, let’s rewrite the gradient estimate with added baseline as a sum of two terms:
Ex∼q [P (x) q(x) − Z πθ(x) q(x) ] ∇θ log πθ(x) = [ Ex∼q P (x) q(x) ∇θ log πθ ] − [ Ex∼qZ πθ(x) q(x) ∇θ log πθ ] (40)
= [ Ex∼q P (x)
q(x) ∇θ log πθ
] − Z [∑ x ∇θπθ(x) ]
(41)
Here again the second term does not introduce bias because Z [∑ x∇θπθ(x) ] = 0.
Null Advantage on Average In the case of sampling with πθ in the online DPG choosing B = Z had the benefit that the advantageRθ(x)−B was centered around 0, namely: Ex∼πθ [Rθ(x)−Z] = 0. With the Boff(x) baseline for the DPGoff this important property is also maintained. The advantage now takes the form P (x)q(x) − Z πθ(x) q(x) and then:
Ex∼q [P (x) q(x) − Z πθ(x) q(x) ] = ∑ x P (x)− Zπθ(x) (42)
= Z − Z ∑ x πθ(x) = 0. (43)
To visualize things better, we elaborate the difference in forms of rewards, baseline and gradients before and after addition of the baseline between DPG (on policy) and DPGoff (off policy) in Table 2.

heading: C CODE GENERATION WITH COMPILABILITY CONSTRAINTS EXPERIMENTS

heading: C.1 EXPERIMENTAL SETUP
text: Energy-based model We represent a language model producing only compilable sequences as the following product-of-experts (Hinton, 2002) EBM:
P (x) = a(x)b(x), (44)
where a is the original language model pre-trained using a standard autoregressive language modeling objective and b(x) = 1 iff x is a syntactically correct Python program and b(x) = 0 otherwise.
Dataset In contrast with experiments with GPT-2, we trained a custom language model to obtain a. To prepare the training dataset for a, we started from the Python150 dataset, which consists of 150k Python source code files obtained from GitHub Raychev et al. (2016). We extracted 713k Python functions (both methods and standalone functions) from 150k using the code from Roziere et al. (2020) while filtering out functions that didn’t compile (b(x) = 0) or were less than 128 BPE tokens long. We then split the dataset into a training subset Dtrain and test subset Dtest.
Initial language model a: We implemented a using the GPT-2 Radford et al. (2019) architecture with 117m parameters (gpt2-small). First, we used Dtrain to train a byte-level BPE tokenizer. We included two special tokens, BOS and EOS, and obtained a vocabulary of 50k tokens. Then, we trained a on Dtrain for one epoch.
Compilability Scorer b We evaluate whether a sample x is compilable by first removing BOS and EOS tokens and then calling the compile_command function from codeop module of Python Standard Library13 with x as the argument. compile_command tries to compile a string of Python code and raises and exception if there is it fails (e.g. raises SyntaxError for invalid Python syntax and ValueError or OverflowError if there is an invalid literal in x). If compile_command returns a code object, b(x) = 1. Otherwise (if an exception is raised or None is returned), b(x) = 0. Note that our notion of compilability is concerned only with syntactic correctness and does not execute the body of a function.

heading: C.2 METRICS
text: In addition to Ex∼πθb(x), DKL(p, πθ), DKL(πθ, a), Distinct-1 (Li et al., 2016a) and Self-BLEU-5 (Zhu et al., 2018), we report the following metrics:
1. Perplexity measured on Dtest, a held-out subset of the data used for training a, calculated as exp [ − 1 N ∑ x∈Dtest log πθ(x) ] ,
13https://docs.python.org/3/library/codeop.html
where N is the total number of tokens in Dtest. 2. Sequence length, the average number of characters in generated sequence x after detokeniza-
tion, 3. AST node count, the average number of nodes in an abstract syntax tree (AST) of sequences
that compile. Samples are parsed to their corresponding ASTs using the ast module from Python Standard Library.14 Intuitively, this metric indicates the logical (as opposed to surface) complexity of generated programs.

heading: C.3 RESULTS
text: We report the performance of GDC and GDC++ as well as Reinforce on Table 3.
Reinforce with R(x) = b(x) improves compilability but that comes at a cost of large divergence from p and a. This divergence translates into a decrease in sequence length and logical complexity (in terms of the number of nodes in ASTs of generated sequences). Heavily decreased sequence length (most of the generated functions are one-liners) accounts for an artificial increase in diversity metrics (Self-BLEU-5 and Distinct-1).
GDC and GDC++ are the only method that consistently improve compilability rate while decreasing divergence from p, maintaining the diversity of a and only slightly decreasing sequence length and the number of nodes in ASTs. Moreover, as a by-product of improving compilability, GDC and GDC++ are also able to slightly decrease the perplexity and the frequency of PEP8 violations per character. The addition of baseline in GDC++ improves its performance in terms of constraint satisfaction, KL divergences and downstream metrics (e.g. lower Self-BLEU-5, higher Distinct-1).
14https://docs.python.org/3/library/ast.html

heading: D EXTRA DETAILS ON METRICS AND ALGORITHMS
text: Calculation of metrics relative to p, such as DKL(p, πθ), is not straightforward since the distribution p ∝ P is only implicitly represented by the unnormalized EBM P , and one cannot easily obtain direct samples from p. Instead, we apply the following workarounds. Given P and a proposal distribution q that we can sample from, using importance sampling (Owen, 2013), we calculate the partition function Z as follows:
Z = ∑ x P (x) = ∑ x q(x) P (x)/q(x) (45)
= Ex∼q P (x)/q(x). (46)
The precision of this estimate depends on the sample size and the quality of the proposal distribution q. We calculate a moving average estimate ZMA of Z which is then used inside the estimations of DKL(p, πθ) and DKL(p, q) (see below Algorithm 2, lines 7 and 8). ZMA is updated at each training iteration. ZMA is an unbiased estimate of Z because each Ẑi is an unbiased estimate of Z based on K samples. Moreover, because the proposal distribution q evolves and gets closer to the target distribution p, the quality of the estimate of ZMA through importance sampling increases.
With an estimate of Z, we can compute DKL(p, πθ) as
DKL(p, πθ) = ∑ x p(x) log p(x) πθ(x) (47)
= ∑ x p(x) log P (x) Zπθ(x) (48)
= − logZ + ∑ x p(x) log P (x) πθ(x) (49)
= − logZ + ∑ x q(x) p(x) q(x) log P (x) πθ(x) (50)
= − logZ + 1 Z Ex∼q P (x) q(x) log P (x) πθ(x) . (51)
Similarly, for TVD(p, πθ):
TVD(p, πθ) = 1
2 ∑ x |p(x)− πθ(x)| (52)
= 1
2 ∑ x q(x) ∣∣∣∣πθ(x)q(x) − p(x)q(x) ∣∣∣∣ (53)
= 1
2 ∑ x q(x) ∣∣∣∣πθ(x)q(x) − P (x)Z q(x) ∣∣∣∣ (54)
= 1
2 Ex∼q ∣∣∣∣πθ(x)q(x) − P (x)Z q(x) ∣∣∣∣ . (55)
See Algorithm 2 for a detailed pseudocode describing how metric computation is integrated in the training loop of KL-DPG.
Algorithm 2 KL-DPG with baseline (detailed) Require: P , initial policy q
1: πθ ← q 2: ZMA ← 0 3: for each iteration i do 4: for each step k ∈ [1,K] do 5: sample xk from q(·) 6: θ ← θ + α(θ) [ P (xk) q(xk) − Z πθ(xk)q(xk) ] ∇θ log πθ(xk)
7: Ẑi ← 1K ∑ k P (xk)/q(xk) 8: ZMA ← i∗ZMA+Ẑii+1 9: D̂KL(p, πθ)← − logZMA + 1/(KZMA) ∑ k P (xk) q(xk) log P (xk)πθ(xk)
10: D̂KL(p, q)← − logZMA + 1/(KZMA) ∑ k P (xk) q(xk)
log P (xk)q(xk) 11: if D̂KL(p, πθ) < D̂KL(p, q) then 12: q ← πθ

heading: Ensure: πθ

heading: E HYPERPARAMETERS AND TRAINING DETAILS
text: We implemented all models using PyTorch (Paszke et al., 2019) and HuggingFace (Wolf et al., 2019). Based on Khalifa et al. (2021) published source code: https://github.com/naver/gdc. Each training run took approximately 5 days on 2 Nvidia V100 GPUs. For a detailed list of hyperparameter values, see Table 4; for a description of hyperparameters specific to Ziegler and GDC, see (Ziegler et al., 2019) and (Khalifa et al., 2021).

heading: F EXTENDED EVALUATION
text: between GDC and GDC++ is underlined. Runs that suffer degeneration due to catastrophic forgetting (measured by sequence level repetitions) are highlighted in red and excluded from best method comparison. Our method GDC++ that includes a baseline for variance reduction, outperforms GDC (Khalifa et al., 2021) in 7/10 tasks in terms of control satisfaction rate (Ctrl), as well as convergence towards the optimal policy (KL(p,π)) and distance from the original LM (KL(pi,a)) in 10/10 of the tasks.
“amazing" (with overall frequency 10−3) highlighted in green. Tokens are highlighted with yellow with different intensities to indicate their overall frequencies in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity). φ(x) = 1 indicates the satisfaction of the constraint in the sample and reps the number of repetitions of the very same sample in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity).
“WikiLeaks" (with overall frequency 10−4) highlighted in green. Tokens are highlighted with yellow with different intensities to indicate their overall frequencies in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity). φ(x) = 1 indicates the satisfaction of the constraint in the sample and reps the number of repetitions of the very same sample in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity).
“politics". Tokens are highlighted with yellow with different intensities to indicate their overall frequencies in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity). φ(x) = 1 indicates the satisfaction of the constraint in the sample and reps the number of repetitions of the very same sample in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity).
“science". Tokens are highlighted with yellow with different intensities to indicate their overall frequencies in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity). φ(x) = 1 indicates the satisfaction of the constraint in the sample and reps the number of repetitions of the very same sample in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity).
sentiments. Tokens are highlighted with yellow with different intensities to indicate their overall frequencies in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity). φ(x) = 1 indicates the satisfaction of the constraint in the sample and reps the number of repetitions of the very same sample in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity).
sentiments. Tokens are highlighted with yellow with different intensities to indicate their overall frequencies in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity). φ(x) = 1 indicates the satisfaction of the constraint in the sample and reps the number of repetitions of the very same sample in the generated corpus of 10k samples (lower intensity indicates higher corpus level diversity).

Review summaries:
The authors explore the connections between reward maximimization (RM) with REINFORCE and distribution matching (DM) using distributional policy gradients (DPG), which uses importance sampling from a proposal distribution $q$ to minimize  $L = E_p[log\pi]$, where $p(x) = P(x)/Z$ is a distribution (energy-based model) that is difficult to sample from (but incorporates generation constraints, such as gender percentage, etc.). The corresponding DM sequence-level reward in REINFORCE is p/q. The authors derive and utilize $E_q[p/q] = \pi/q*Z$ as a baseline, which lowers variance, and leads to improvements in loss L and sampling efficiency. 

This paper shows a connection between the reward maximization (RM) approach and the distribution matching (DM) approach for fine-tuning language models. The authors further suggest borrowing the baseline idea from reinforcement learning and applying it to DM to reduce variance. Experiments show that adding a baseline not only makes training more stable, but also converges to a better solution.

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