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: In this paper, the authors propose a controllable decoding mechanism for language generation models (e.g., GPT2) using constrained optimization. The proposed method explicitly formulate the sentence decoding as a constraint optimization problem where the free parameters are the distribution over the vocabulary at each step and two types of constraints are defined one over the generate text itself (f(Y)) and one over the similarity between the input and the outputs (g(x,y)). Please refer to Eq. 1 where the problem is clearly formalized. Since directly solving the problem in Eq 1. is not computationally feasible the author reformulates the problem using the Lagrangian formulation for container optimization problems (Eq. 3). However, this min-max formulation introduces several problems (e.g., the tradeoff between different objectives, hard to find opt for non-convex functions). To cope with these issues, the authors proposed a gradient descent/ascent solution (Section 2.1), a soft relaxation for each token using the probability simplex (exponential gradient descent + normalization) (Section 2.2), and an annealing threshold (Section 2.3).   The authors benchmarked the proposed decoding schema in both style transfer and style-controlled machine translation. In the style transfer, they follow the methodology in Krishna et al. 2020 and they study two contains such as formality and the semantic similarity between the input and the output. Similarly in the style-controlled machine translation, they study both the formality style and the cross-lingual similarity between input-output. In both settings, the author compared the proposed method with other decoding techniques (greedy and beam-search) and other controllable language generation techniques (e.g., FUDGE) using both automatic and human evaluation.  Finally, the authors show some preliminary results in controlling multiple attributes (Age, style pos/neg), and discuss the speed and memory requirements of the model.