Controlled Text Generation as Continuous Optimization with Multiple Constraints

1 Introduction . Recent advances in language models [ 11 , 51 , 52 ] trained on large-scale web text corpora have led to great improvements in state-of-the-art on many natural language processing ( NLP ) tasks including the ability to generate increasingly coherent text [ 3 ] . However , once such models are trained , they are prone to degeneration [ 66 ] and biased , non-factual outputs [ 16 , 45 ] as it is difficult to control the characteristics or attributes of the generated text without architectural modifications [ 26 , 28 , 35 ] and fine-tuning the models on attribute-specific corpora [ 29 , 7 ] . This can be even more challenging if multiple attributes are involved as labeled data for each combination of attributes can be difficult to obtain . We focus on controlled text generation where the goal is to decode from a text generation model such that the outputs satisfy certain constraints , which the model was not necessarily trained on . For example , given a dialogue generation model , additionally constraining the generated responses to be polite , although the model was not optimized for politeness during training . Recent works address this problem with left-to-right autoregressive decoding , and modify the vocabulary distribution at every step directly using classifiers or language models trained on attribute specific corpora [ 72 , 40 , 39 ] , or indirectly via backpropagating gradients through model activations [ 8 ] . While exhibiting high level of attribute control , by design these methods can only work with categorical attributes ( typically only one attribute ) and condition only on the left context while decoding . Additionally , they often require several heuristics to work and are prone to adversarial outputs [ 64 ] . 1The code is available at https : //github.com/Sachin19/mucoco 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To address these concerns , we propose the following decoding algorithm . Given a pretrained language model , we posit decoding from it as an optimization problem . First , we relax this discrete optimization problem to a continuous one by representing each token as a simplex on the target vocabulary [ 19 ] . This allows us to use continuous optimization techniques like gradient-descent considering each token distribution as parameters ; while keeping the language model ’ s parameters fixed ( §2 ) . Second , we represent each target attribute to control as a differentiable function . We formulate controllable decoding as a multi-objective optimization problem , with maximizing the log-probability of the language model as well as target attributes as objectives . To make this optimization feasible via gradient-descent , we repurpose it to a constraint optimization problem and solve the dual using the modified differential method of multipliers [ 48 ] . We call the algorithm MUCOCO , for incorporating multiple constraints through continuous optimization . We validate MUCOCO on three conditional text generation tasks with different types of sentence level constraints : ( 1 ) Adding formality and cross-lingual similarity in a machine translation model ; ( 2 ) Ensuring transfer and content-preservation in a style-transfer model , and finally ( 3 ) Incorporating multiple styles and attributes ( e.g. , formality , sentiment magnitude , writer ’ s age group ) in a paraphrasing model . With automatic as well as human evaluations we find that our proposed method outperforms strong baselines . 2 MUCOCO : Constrained Decoding as Multi-Objective Optimization . For a given language generation task , let G model the conditional probability p ( y|x ) of the output sequence y = y1 , . . . , yn , given the input sequence x = x1 , . . . , xn . This model can be parameterized using any differentiable architecture like Transformers [ 62 ] or LSTMs [ 20 ] and trained with any loss function [ 12 , 30 ] , either autoregressively or non-autoregressively [ 17 ] . Traditionally , given an input x , decoding from such a model requires finding the output sequence with the highest probability or the lowest negative log-probability , y∗ = arg miny∈Y − logP ( y|x ) . Here Y is the set of all possible output sequences . In practice , searching Y to find the highest probability generation is intractable as the space of possible sequences grows exponentially with sequence length and has also been shown to produce undesirable solutions [ 60 ] . In most prior work , simple heuristics like beam search , or sampling have been adopted to find approximate solutions , where the text is generated one token at a time ( usually left to right ) with the output of step t being fed to the input at step t+ 1 . In this work , however , given G and an input sequence x , we are interested in finding an output sequence y that not only maximizes the output probability but also optimizes multiple objectives defined over x and y . More formally , we seek to find a y that minimizes all of the following objectives y∗ = arg min y∈Y ( − log p ( y|x ) , f1 ( y ) , . . . , fu ( y ) , g1 ( x , y ) , . . . , gv ( x , y ) ) ( 1 ) Here each fi is a function defined over the output sequence y , for example , the negative log-probability of an attribute ( e.g. , formality ) classifier we want the output sequence to satisfy . And each gj is a function defined over both the input and output sequence , for example , semantic similarity between x and y [ 55 ] . We assume all fi and gj are differentiable . This is a multi-objective optimization with several possible solutions . Since there are many objectives to minimize , a left-to-right decoding strategy like beam search or sampling will simply not work due to several reasons . First , the objectives fi and gj are sentence-level and hard to define accurately only on generated left-context [ 72 , 40 ] . Even if we are able to define them , as we add more objectives this process becomes very computationally expensive . Following prior work [ 19 , 50 ] , we formulate this as a continuous optimization process instead of a standard discrete one , and then use standard algorithms for continuous optimization ( like gradient descent ) for decoding . We maintain a soft-representation of the sequence y , ỹ = ( ỹ1 , . . . , ỹn ) , where each ỹk ∈ ∆V is a simplex over the target vocabulary of size V , representing the probability of the k-th token . To decode a sentence , we initialize each ỹi uniformly over V , and treat the entire output sentence as the parameters for gradient descent keeping the parameters of G , fi , gj fixed . After gradient descent has converged , we generate discrete text by selecting the token with the highest probability in ỹk . We provide more details on the optimization procedure in §2.2 . To make optimization feasible , a multi-objective problem generally yields itself to the following formulation : arg min y −α log p ( y|x ) + u∑ i=1 λifi ( y ) + v∑ j=1 µjgj ( x , y ) , ( 2 ) for some statically or dynamically computed weights λi and µj for each i and j , where α+ ∑ i λi +∑ j µj = 1 . Although this weighted summation formulation is intuitively appealing , it typically requires an expensive grid-search over the various scalings or use of a heuristic [ 25 , 6 , 18 ] . Furthermore , this formulation by definition assumes a trade-off between the different objectives by essentially assigning an importance weight to each of them . This problem is further exacerbated when different objectives have widely varying scales2 with smaller scale objectives just getting ignored . More concretely , a multi-objective formulation as we define in ( 1 ) admits several possible “ optimal ” solutions also known as the Pareto set [ 9 ] . The image of the Pareto set is called the Pareto front . Since we define all objectives using neural networks , the Pareto front in our case is non-convex , where linear combinations of objectives are shown to be unsuccessful in finding good solutions [ 36 , 37 , 10 ] ( see figure 1 for an example ) . Ideally , our goal is a tunable optimization algorithm that finds solutions on the Pareto front , i.e. , every solution on the Pareto front should have a hyperparameter value for which the optimization algorithm finds that solution . In order to achieve this , we reframe our optimization problem as a Lagrangian optimization problem instead . We choose one of the losses as the primary objective and consider other losses as constraints . The goal is to minimize the primary loss subject to the secondary losses , each below a threshold value . More formally , arg min y − logP ( y|x ) subject to fi ( y ) ≤ i , i ∈ { 1 , · · · , u } gj ( x , y ) ≤ ξj , j ∈ { 1 , · · · , v } . Here i and ξj are tunable hyperparameters whose values ’ change can result in different solutions on the Pareto front . This formulation leads to an intuitive interpretation of the decoding process that the generated text from the model G should satisfy the constraints while being as faithful to the primary objective as much as possible.3 Consequently , the Lagrangian we end up with looks similar to our original total loss linearly combined as in ( 2 ) given by L ( y , λ1 , . . . , λu , µ1 , ... µv ) = − log p ( y|x ) − u∑ i=1 λi ( i − fi ( y ) ) − v∑ j=1 µj ( ξj − gj ( x , y ) ) ( 3 ) 2For example , classifier log-probabilities are in ( 0 , inf ) while sentence similarities usually lie in [ 0,1 ] . 3For example , defining fi ( y ) = p ( a|y ) as the probability of a desired attribute a in y leads to a natural threshold of fi ( y ) > 0.5 . For a well-calibrated fi , an even higher threshold could be used for inducing highly indicative features of a in y. where λi , µj are Lagrange multipliers , and an optimal output y∗ can be obtained as y∗ = arg miny maxλi≥0 , µi≥0 L ( y , λi , µi ) . However , the traditional method of solving the dual function to find λi , µj that matches i , ξj , respectively , again leads to a linear trade-off between the various objectives . When the Pareto front is non-convex as in our case , with gradient-descent , the constraints can be ignored and we still can not always find optimal solutions by tuning i , ξj [ 48 ] . 2.1 Modified Differential Method of Multipliers . The fundamental issue in both linear combination of objectives and solving the dual is that fixed scalings λi and µi ( manually pre-determined or obtained by solving the dual ) do not work well with gradient descent to minimize for y . Following prior work on differential method of multipliers [ 48 ] , we propose to use a single gradient descent to optimize for both Lagrangian multipliers and y simultaneously as follows : y ( t ) = y ( t−1 ) − η1∇yL , λti = λt−1i + η2∇λiL , µ t i = µ t−1 i + η2∇µiL . ( 4 ) We follow the gradient of L downwards for the y ( descent ) and upwards for the multipliers ( ascent ) while making sure that the multipliers remain positive ( by setting the multipliers to 0 whenever they become negative ) . Intuitively , this algorithm works by increasing the value of the multiplier with each gradient step as long as the constraint is violated . But when the constraint is suddenly satisfied and the multiplier is still large , it might take a number of gradient steps before the gradient descent pushes it to 0 , thus causing the solution to be pushed further away from the constraint . As soon as the multipliers become 0 ( or negative ) , the constraint is ignored and the process continues . However when the optimization hits the constraint again , this whole cycle repeats , resulting in “ oscillations ” . We introduce a dampening parameter to each of the multipliers to reduce these oscillations ( again following Platt and Barr [ 48 ] ) and update the Lagrangian as follows : L ( y , λi , µj ) = − log p ( y|x ) − u∑ i=1 ( λi − ζi ) ( i − fi ( y ) ) − v∑ j=1 ( µj − νj ) ( ξj − gj ( x , y ) ) , ( 5 ) where ζi = d ∗ stop-gradient ( i − fi ( y ) ) , νj = d ∗ stop-gradient ( ξj − gj ( x , y ) ) and d is a hyperparameter . d does not affect the final y , just how quickly the algorithm converges to it ( We use d = 1 in all experiments ) . stop-gradient ( · ) indicates that the argument is detached from the computational graph and does not contribute to the gradient computation . When a constraint is not satisfied ( i − fi ( y ) < 0 , hence ζi < 0 ) , the dampening parameter ζi being negative incurs higher penalty on the violation than when not using any dampening , without actually increasing the value of λi too much . But when the constraint is satisfied , it helps quickly reduce the value of penalty being incurred on the constraint while the multiplier converges to 0 .

Summary: This paper presents a decoding algorithm for controlled text generation. The problem is formulated as finding the most likely output $y$ under given constraints. The authors adopt the method of Lagrange multipliers, and introduce dampening in gradient descent to reduce oscillations. To deal with discrete $y$, the authors choose to optimize over probability simplex $\tilde y$ using exponentiated gradient descent. The authors also report that converting $\tilde y$ to a one-hot vector via a straight through estimator helps to converge faster.  Experiments include increasing formality in style transfer and machine translation. Compared with greedy decoding/beam search, the proposed method with formality and semantic similarity constraints achieves better content preservation/BLEU score and better transfer accuracy. It also compares favorably to the recent FUDGE method for controlled generation.