Controllable Text Generation with Neurally-Decomposed Oracle
Anonymous Author(s)
AffiliationAddress email
Abstract
We propose a general and efficient framework to control auto-regressive generationmodels with NeurAlly-Decomposed Oracle (NADO).
Given a pre-trained baselanguage model and a sequence-level boolean oracle function, we propose todecompose the oracle function into token-level guidance to steer the base modelin text generation.
Specifically, the token-level guidance is approximated by aneural model trained with examples sampled from the base model, demanding noadditional auxiliary labeled data.
We present the closed-form optimal solution toincorporate the token-level guidance into the base model for controllable generation.
We further provide a theoretical analysis of how the approximation quality of
NADO affects the controllable generation results.
Experiments conducted on twoapplications: (1) text generation with lexical constraints and (2) machine translationwith formality control demonstrate that our framework efficiently guides the basemodel towards the given oracle while maintaining high generation quality.
Introduction
Auto-regressive language models have been widely used for text generation.
With the recent development of large-scale pre-trained language models Radford et al.
; Brown et al.
; Raffelet al.
; Lewis et al.
, they have achieved state-of-the-art performances in applicationssuch as machine translation (Bahdanau et al., 2015;
Luong et al., 2015), image captioning (Andersonet al., 2018; You et al., 2016) and open domain text generation (Zhang and Lapata, 2014;
Yao et al.
Vinyals and Le, 2015; Shang et al., 2015; Lu et al., 2018).
However, many applications requireto control model output with specific sequence-level attributes.
The attributes can be specified bya set of rules 1 or by an abstract concept ( e.g. , the generated text follows a particular writing style).
How to control auto-regressive language models to satisfy these attributes is an open challenge.
In this paper, we propose a general and flexible framework for controllable generation.
Given a basepre-trained language model and a sequence-level oracle function indicating whether an attribute issatisfied, our goal is to guide the text generation by the oracle to satisfy certain attributes.
To this end,we propose decomposing the sequence-level attributes into token-level guidance.
When generatingi − th token in the output sequence, instead of sampling from the base model given the prefix, wemodify the probability distribution of the output token based on the token-level guidance.
Specifically,we formulate the control as an optimization problem based on posterior regularization (Ganchevet al., 2010) and solve the close-form optimal solution to incorporate the token-level guidance for x +
Sampling
Oracle ? ? ??
Training generated text
A girl with hair sits at a table looking around.
A woman is sitting and looking down.
hair
[Figure]
a) We take lexically constrained generation as an example, where the oracle is checking whether all
keywords in the input
x are incorporated in gener- ated text y produced by base model p .
With proper training with the
examples sampled from the base model p (dashed arrow and
labeled with the oracle , the oracle is decomposed
into
token
-level
NADO
and realized
by an auxiliary model R θ .
The
R
θ
provides guidance
when
generating
text with the base model (
see
details in Fig. 1(b)).
[Figure] [Figure]
The illustration of the controlled generation pro- cess.
Both the base model and the auxiliary model (NADO) take input x and the generated sequence (prefix) y <
L as input.
The base model, in each step, outputs a token distribution p ( y i | x , y <i ) .
Guided by NADO R θ , we obtain the distribution q
(See Sec. 3.2), based on which we generate the current token.
text generation.
The decomposition is approximated by an auxiliary neural network model, called
NeurAlly-Decomposed Oracle (NADO), which is trained on data sampled from the base model andsupervised by the sequence-level oracle (see the illustration Fig.
1a).
We further provide theoreticalanalysis on how NADO’s approximation quality affects the controllable generation results.
Note thatin the entire process, we treat the base model and the sequence-level oracle as black-box functions,without the need for any refacting or fine-tuning.
A few existing controllable generation works ( e.g. , Lu et al.
(2021)) design decoding algorithms forlexical constraints .
However, their approach is purely algorithmic and cannot generally be applied toother types of oracles.
Another line of work such as PPLM (Dathathri et al., 2020), GeDI (Krauseet al., 2021), and FUDGE (Yang and Klein, 2021) also aims to guide the base model with an auxiliarymodel.
However, they either shift the base model distribution in a post-hoc manner without theoreticalguarantee, or/and require external labeled data to train the auxiliary model.
Instead, we derive aclosed-form solution for the optimal way to incorporate the oracle, without requiring external labeleddata or token-level guidance.
Furthermore, since NADO is trained on the data sampled from the basemodels, it aligns better with the base model and thus can achieve better control.
We conduct experiments on lexically constrained generation (LCG) tasks and a machine translation(MT) formality change task.
In LCG tasks, the oracle is a rule-based keyword checker.
We achievealmost perfect keyword incorporation with significantly boosted BLEU scores compared to previousapproaches that design specific decoding algorithms
Lu et al.
In the formalitycontrolled MT task, we are provided with a formality oracle predicting whether a sentence is formalor not, and the goal is to guide the model to generate formal translations.
Compared with recentwork Yang and Klein (2021), we improve the BLEU score by 3 points as well as improve theformality rate, demonstrating NADO’s superior ability in incorporating external oracle supervision.
Both experiments demonstrate the effectiveness of our framework in dealing with various types ofcontrol and generating high-quality texts.
Related Work
Controllable Text Generation with Auto-regressive Models Most previous work on controllabletext generation are based on the auto-regressive framework.
Zhang et al.
(2022) summarize thesemethods into three categories: fine-tuning , refact/retraining and post-processing .
The first two categories, e.g., CTRL (Keskar et al., 2019), prompt methods (Shin et al., 2020; Lester et al., 2021;
Li and Liang, 2021), are usually weaker in controlling and inefficient in training considering the sizeof language models are dramatically increasing nowadays.
Generally, the post-processing methodsare considered expensive in inference and low quality in generated texts.
However, our framework,as a kind of post-processing method, is able to achieve high generation quality demonstrated inexperiments and is efficient in the inference time.
Controllable Text Generation via Post-processing There are two major lines in post-processing: (1)adapting the decoding algorithm and (2) guiding generation with an auxiliary model.
For some tokenlevel controlled generation tasks like lexically constrained generation, we can inject the constraintsinto the decoding algorithm (e.g., constrained beam search (Anderson et al., 2017) and
NeuroLogicdecoding (Lu et al., 2021)).
Though shown effectiveness in lexically constrained generation, thesealgorithmic methods fail to fundamentally touch the token distribution, and are hard to handle otherabstract attributes.
In the second line, PPLM (Dathathri et al., 2020) proposes an auxiliary discriminator for the expectedattribute to guide the model; GeDi (Krause et al., 2021) and DEXPERTS (Liu et al., 2021) applycontrastive learning and train an auxiliary language model to reweight the token distribution ineach step; Plug-and-Blend (Lin and Riedl, 2021) further extends the GeDi framework by adding aplanner architecture.
FUDGE (Yang and Klein, 2021) leverages external token-level oracle to train adiscriminator for guiding the base model.
These methods either require external token-level oracleguidance or auxiliary labeled datasets to train the auxiliary models.
However, the distribution of thedata used to train the auxiliary model is different from what the based model is trained on.
Thisdistributional discrepancy causes the drop of generating quality as we will show in the experiments.
For example, given a controlling attribute a , Fudge generates next token
y i based on Bayesian rule
P ( y i | y <i , a ) ∝ P ( a | y ≤ i ) P
y i |
y <i ) .
However, their P ( a | y ≤ i ) and P ( y i | y <i ) are not estimatedbased on the same distribution.
In contrast, NADO is trained with data sampled from the base model.
Therefore, it learns to incorporate with the base model, which avoids the distributional discrepancy.
We also provide a principle, theoretical framework to discuss the optimal solution of incorporatingthe sequence-level oracle.
Methodology
We approach the sequence-level controllable generation problem by decomposing the sentence-leveloracle into token-level guidance.
We formulate this as an optimization problem.
Since the token-levelguidance is intractable, we propose to train an auxiliary model, called NeurAlly-Decomposed Oracle(NADO), to approximate it.
During the inference time, NADO guides the base model to generatesequences that satisfy the oracle constraints.
In the rest of this section, we discuss 1) the formulation to decompose the sequence-level oraclefunction into token-level guidance; 2) the formulation to incorporate the token-level guidance intothe base model to achieve control; 3) the approximation of the token-level guidance using NADO; 4)a theoretical analysis of the impact of NADO approximation to the controllable generation results;and 5) the training of NADO.
3.1 Notations and Problem Formulation
We use x ∈ X to denote the input and y ∈ Y to denote the generated sequence.
y i is the i − thtoken in y and y <
i is the sequence prefix from the beginning to the ( i − 1) − th token.
We denotethe base auto-regressive generation model as p ( y i | x , y <i ) , hence the sequence-level distribution isgiven by p ( y | x ) =
(cid:81)
i p ( y i | x , y <i ) .
A sequence-level oracle is defined as a boolean function C :
X ×Y → { 0 , 1 } .
We formalize the optimization objective based on posterior regularization (Ganchevet al., 2010).
Basically, we explore a token-level distribution q ∗ ( y i | x , y <i ) and its correspondingsequence-level distribution q ∗ ( y | x ) , satisfying i.e. can be treated an auto regressive model.
q ∗ ( y | x ) =
if C ( x , y ) = 0 , i.e. , q ∗ only generates sequences satisfying the oracle C.3.
Given an input x , KL ( p ( y | x ) ∥ q ∗ ( y | x )) is minimized, i.e. , q ∗ should be as similar to thebase model as possible.
3.2 Token-level Guidance and Closed-Form Solution For q ∗
Before we compute the solution for q ∗ , given the base model p and oracle C , we first define thetoken-level guidance as a success rate prediction function R Cp ( x ) , which defines the probability ofthe sequence generated by p satisfies the oracle given input x .
We further define R Cp ( x , y <i ) as theprobability given input x and prefix y <i .
By definition, we have
[Equation]
With the function R Cp , we are now able to derive the closed-form solution of q ∗ considering conditions(2) and (3) defined in Sec. 3.1.
Given input x , we define the feasible sequence-level distribution set
Q as
[Equation]
then the sequence-level closed-form solution for q ∗ is given by
[Equation]
Considering condition (1) and making q ∗ tractable, we decompose q ∗ ( y | x ) into token-level.
Theclosed-form solution is given by
[Equation]
The decomposition is unique.
The proof and detailed derivation can be found in the appendix.
Control with Soft Constraints.
In Eq.
(2) we define the feasible distribution set as distribution thatthe possibility of a sequence violate the oracle function is 0 .
However, in some applications, we expectto control the generation with soft constraints.
For example, we want the model to generate sentenceabout sports with probability r = 0 . 8 .
Our framework also supports controlling the generation withsoft constraints.
To achieve this, with a pre-defined ratio r ∈ [0 , 1] , we alternatively define a generalfeasible set Q as
[Equation]
where Eq.
(2) is the special case when r = 0 .
The general token-level closed-form solution is
[Equation]
where α = r R Cp ( x ) , β = 1 − r 1 − R Cp ( x ) .
Similar to Eq.
(4), once we have access to R Cp , we can directly compute the closed-form solutioneven though the form is much more complicated.
In this paper we only focus on the hard constraint( r = 0 ), however, here we demonstrate that our framework is capable for the relaxed controlling.
3.3 Approximating R Cp by NADO and Theoretical Analysis
Unfortunately, function R Cp defined in Eq.
(1) is intractable.
We cannot enumerate all the possiblesequence y since the space is exponentially large.
Hence, we train a neural model NADO to approximate this well-defined function.
We use R θ to denote NADO parameterized by θ .
Inthis section, we derive bounds to have a theoretical analysis about the correlation between errorsin approximation and errors in corresponding sequence-level distribution.
Generally, when R θapproximates R Cp precisely enough, we have an upper bound for the sequence-level distributiondiscrepancy.
The following lemma provides the formal definition.
Lemma 1 We define distribution
[Equation]
If there exists δ > 1 such that given input x , ∀ y <i , 1 δ < R θ ( x , y <i ) R Cp
[Equation]
x , y <i ) <
δ, we have
[Equation]
where L is the length of the sequence y .
The proof is in the appendix.
We also notice that by definition, R Cp satisfies the following equation:
[Equation]
If R also satisfies Eq.
(6), we can tighten this bound.
Formally,
Lemma
Given the condition in Lemma 1, if q is naturally a valid distribution without normalization ( i.e. , (cid:80)
y i R
[Equation]
( x , y <i ) R θ
( x , y <i − 1 ) p ( y i | x , y <i )
= 1 ), we have
[Equation]
The proof is in the appendix.
This two bounds show that in training of model R , we should push R θto satisfy Eq.
(6) as well as approximating R Cp .
3.4 Training NADO
In Fig.
1b we show the architecture of NADO.
Generally, NADO can be any seq2seq model.
Intraining, it takes x , y as input and predicts from R θ
( x , y < 0 ) to R θ ( x , y <T ) .
In the inference, thereare two parallel forward pass 2 to compute the token distribution q .
Considering the decoder size ofthe NADO is usually far less than the base model, the whole forward pass takes no more than 2x basemodel forward pass.
Now we discuss about the training objective.
In training, with some predefined input distribution X ,we sample x ∼ X , y ∼ p ( y | x ) .
We take these sampled ( x , y ) pairs as training examples, and usethe boolean value C ( x , y ) as their labels for all steps.
We use cross entropy (denoted as CE ( · , · ) ) asthe loss function, formally, L CE ( x , y , R θ )
=
(cid:80) T i =
CE ( R θ ( x , y <i ) , C ( x , y )) .
Given a particularinput x , in expectation, we have
[Equation]
Therefore, L CE empirically estimates the cross entropy loss between R θ and the ground truth R C pwhich is intractable.
As we analyze above, we also regularize R θ for satisfying Eq.
(6) based on KL-divergence:
[Equation]
[Equation]
q is KL-divergence regarding p and q as two Bernoulli distributions.
We use a hyper-parameter λ > 0 to balance these losses.
The final training loss is
[Equation]
3.5 Sampling
In Sec.
3.4 we describe that we train NADO by sampled data from base model p .
One advantageis that we are able to leverage different sampling strategies to better adapt to different applicationscenarios.
It is also possible to leverage reinforcement learning to train R θ , and we discuss ourconnection to reinforcement learning in the appendix.
In this section, we introduce two samplingstrategies and their corresponding properties.
Sampling with Temperature Control.
In some task, the output sequences are not diverse much,in other words, the token distribution in each step is very peaky.
Since our NADO is trained on thesampled examples, we expect those examples to cover as much tokens combination as possible toavoid overfitting.
Therefore, we add temperature factor T to smooth the distribution (Ackley et al.,1985).
Specifically, we sample y from distribution p ( y | x ) 1 T , and add coefficient p ( y | x ) 1 − 1 T whencomputing the cross-entropy loss.
Formally, the expected loss is
[Equation]
which is same as the original expected loss in Eq.
Importance Sampling.
In practice, the training process of NADO can be extraordinarily difficultwhen samples generated by the base model p hardly satisfy C .
i.e. E y ∼ p
( y | x ) [ p ( C | x , y )] ≃ 0 .
Hence,we introduce the importance sampling (Hammersley and Morton, 1954) to tackle this issue.
Basically,we leverage existing partially trained ˆ R θ to form distribution ˆ q .
Although ˆ R θ is not well-trained, it isstill able to provide positive guidance to produce samples satisfying C .
Note that ˆ q does not have tobe updated in each training epoch.
With coefficient p ( y | x ) ˆ q ( y | x ) , the expected loss is same as the originalexpected loss:
[Equation]
Experiments
We conduct experiments on two tasks: lexically constrained generation (LCG) and machine translation(MT) with formality change.
For the former, we use GPT-
(Radford et al., 2019) as the base modeland for the latter, we use a sequence-to-sequence model, MarianMT (Junczys-Dowmunt et al., 2018).
We demonstrate our framework is generally effective in both scenarios.
The boolean oracle is arule-based function checking whether all lexical constraints are satisfied in LCG task, while in MT itis a classifier trained on an external dataset identifying the formality of the text.
We put all detailsabout hyper-parameter settings in the appendix.
4.1 Text Generation with Lexical Constraints
We evaluate our model on two general classes of LCG problems: 1) unsupervised LCG, i.e. annotationfor lexical constraints are not available during training, but are expected to be in their exact orderand lexical form during inference; 2) supervised LCG, i.e. annotation for lexical constraints areavailable, yet the words may appear in a different lexical form (e.g., “look” can appear in the pasttense “looked”) or a different order in the generated text.
In both cases, we define oracle C as a boolean function indicating whether the generated sequence satisfies all of the lexical constraints .
We do not naturally have negative samples ( i.e. the sequences that do not satisfy all constraints) totrain the auxiliary model in both settings, thus, it is non-trivial to compare against methods requiringboth positive and negative labeled data for training the auxiliary model like FUDGE and GeDi.
Data Setup For unsupervised LCG, we follow the settings in POINTER (Zhang et al., 2020) andconduct our experiments on Yelp!
Review and News dataset.
Each of the unsupervised LCG datasetcontains a great number of un-annotated, raw sequences for training (160K for Yelp!
Review and268,586 for News).
During inference, the model is expected to generate text lexically constrained inthe exact order and form by a specific number of keywords (7 for Yelp!
Review and 4 for News).
Forsupervised LCG, we evaluate the proposed method on CommonGen (Lin et al., 2020).
CommonGenis a supervised LCG task that aims to examine the commonsense of neural text generation models.
For training, it contains 32,651 unique key concepts ( i.e. the constraints) with 67,389 completedsequences in total.
It also contains a validation set with 993 concepts and 4018 reference sequences.
For a more robust evaluation, the dataset maintains an open leaderboard that benchmarks differentapproaches on a withheld test set.
We follow most of the data configurations specified in the originalpaper that first introduced the datasets.
General Model Setup We investigate the effectiveness of different factors in our framework byenumerating different combinations of them.
We implement two types of base model:
• (Seq2seq base model)
A sequence-to-sequence model p ( y | x ) that takes into account thelexical constraints
as condition sequence input
• (DA base model)
A language model that is only domain-adapted to p ( y ) but unconditioned
on anything.
This is a challenging setting, since we impose the lexical constraints only
with NADO.
This setting is to better verify the effectiveness and efficiency of the proposed
method and control
irrelevant factors.
Under both p ( y | x ) and p ( y ) settings, we fine-tune the base model from the pre-trained GPT2-Large.
During training, NADO is trained as a Seq2seq-like model, which takes in the keys (for unsupervised
LCGs, they are generated by randomly sampling a specific number of natural words in the originalsentence) and generates the token-level guidance R θ ( x , y ≤ i ) .
For each pseudo key, we sample 32target text with top-p ( p = 0 . 8 ) random sampling from base model p .
We conduct experiments to testdifferent training setups for NADO:
• (NADO training)
The proposed training process described in Sec. 3.4.
• (Warmup)
We warm up NADO by maximizing the likelihood of positive samples, but
only backpropagating the gradient to the parameters of R θ .
The warm-up R θ is used for
importance sampling described in Sec. 3.5.
With DA base models, however, the warmup
process is always incorporated for practical success of training (see the results for DA
pretrained w/o warmup)
We also consider the setting with warmup only, which can be treated as a stronger baseline to verifythat the major improvement of our framework is not coming from the extended capacity in NADO.
Results and Analysis
We compare the performance under different setups of our model to previousstate-of-the-art methods on LCG tasks, including insertion-based models (Levenshtein Transformer(Gu et al., 2019) with Lexical Constraints (Susanto et al., 2020), InsNet (Lu and Peng, 2021), etc.)and decoding-based algorithms.
We also compare the results with a simple baseline which addressthe problems with a standard Seq2seq pipeline.
The results are as shown in Table 1.
NADO consistently improves the BLEU score and coverage in different setups.
Furthermore, underthe best setting of each task (see bolded items in the table), NADO performs significantly better thanmost baselines in generation quality and can achieve very good lexical constraints coverage rate.
Compared to InsNet, it is much easier for an autoregressive model with NADO to handle flexiblereordering/transformation of lexical constraints.
This is reflected in the performance comparison of
InsNet and NADO on CommonGen dataset.
Under most settings, a Seq2seq base model makes iteasier for the framework to perform well, as it guarantees a reasonable level of lexical constraint
w/ Regularization w/o Regularization
Epochs
BLEU-
[Figure]
w/ Regularization w/o Regularization
Epochs
Coverage %
[Figure]
[Table]
[Table]
coverage in even the initial state of the model.
However, results from using a DA pretrained basemodel with R θ warmup and NADO under importance sampling show that it is still possible to obtaina powerful model with the proposed method under a more challenging setting
[Table]
To study the importance of the regularization term
[Table]
we conduct an ablative study under the optimalsetting on the CommonGen dataset (Seq2seq base model with NADO only).
The results are shownin Figure 2.
While the success of achieving lexical control does not degenerate when NADO w/oregularization overfits, adding regularization can significantly improve the robustness of NADOgeneration quality when training NADO for more epochs.
4.2 Machine Translation with Formality Change
Datasets and Setup We follow the experimental setting in FUDGE (Yang and Klein, 2021) toformalize the results of machine translation.
Given an informal source sentence, our goal is totranslate it into formal sentence written in the target language.
We conduct our experiments on Fisherand CALLHOME Spanish-English Speech Translation Corpus (Post et al., 2013), where both of the
Spanish source and English reference are informal and casual.
Instead of evaluating the translation [Table] on original references, we use the formal and fluent rewritten version of references (Salesky et al.
to evaluate the translation quality by BLEU scores.
In the training process, the formal versionreference is unseen to the models.
We also evaluate the formality scores by a discriminator trained on
GYAFC formality dataset (Rao and Tetreault, 2018) as what FUDGE paper does.
In this experiment,pre-trained Marian MT model (Junczys-Dowmunt et al., 2018) is used as the base model.
In FUDGE, the authors train an auxiliary model also on GYAFC modeling token-level guidance
P ( formal | y <i ) , and leverage it to guide the base model by Bayesian rule P [Equation] y
[Equation]
i , formal
[Equation]
For the formality supervision, FUDGE leverages an external token-level oracle.
In NADO, we loadthe same oracle but exclusively leverage sequence-level binary supervision as oracle C .
We randomlychoose 10,000 (7.2%) source texts from the training set as input examples, and sample 8 target textsby sampling with temperature T from base model p for each source text.
We use those sampledexamples to train NADO.
In total, we have 80 , 000 training samples, which is similar to the number oftraining data (105k) for the token-level oracle in FUDGE.
All the methods are using greedy decoding.
Results and Discussion
The experimental results are shown in Table 2.
Compared to FUDGE,although only the sequence-level supervision is leveraged, we are consistently better in both metrics,especially in BLEU score we boost about 3 points.
We conjecture that the improvement is because ourformulation is more principle and correct.
In methods using auxiliary model to guide the base model,including FUDGE, their formulation is based on Eq. 9.
However, the auxiliary model is trained ona distribution different from where the base model is pretrained on, which leads to a distributionaldiscrepancy issue.
In other words, directly multiplying these two terms is not rigorous, since they areestimated on two different distributions.
On the contrary, NADO is trained specifically to the basemodel.
This avoids the discrepancy issue and provides an accurate guidance.
Considering we areusing the same oracle function and similar number of training samples, the higher generation qualityreflected by BLEU scores supports our conjecture.
In sampling, for each input we sample 8 examples to train R θ , which are usually identical in this task.
Applying temperature in sampling allows NADO to be trained with more diverse data.
Results showthat with a properly set temperature, we can further improve the generation quality.
Conclusion
We purpose a general and efficient framework for controllable generation.
We leverage an auxiliaryneural model, NADO, to approximate the decomposed oracle guidance, and incorporate it with afixed base model.
By training with sampled data from the base model, NADO aligns better with thebase model, and our framework is more flexible dealing with various application scenarios providedby different sampling methods.
As NADO is a general framework, in the future, we plan to applyit in boarder application scenarios.
For example, reducing societal bias (Sheng et al., 2019)
(e.g.,gender or racial bias) in generation by providing corresponding oracle.
[-Lemma 
-]
[-With the help of KL-regularization,we set λ = 1 . 
-]