BOLT : Fast Energy-based Controlled Text Generation with Tunable Biases (2023)

Abstract: Energy-based models (EBMs) have gained popularity for controlled text generation due to their high applicability to a wide range of constraints. However, sampling from EBMs is non-trivial, as it often requires a large number of iterations to converge to plausible text, which slows down the decoding process and makes it less practical for real-world applications. In this work, we propose BOLT, which relies on tunable biases to directly adjust the language model’s output logits. Unlike prior work, BOLT maintains the generator’s autoregressive nature to assert a strong control on token-wise conditional dependencies and overall fluency, and thus converges faster. When compared with state-of-the-arts on controlled generation tasks using both soft constraints (e.g., sentiment control) and hard constraints (e.g., keyword-guided topic control), BOLT demonstrates significantly improved efficiency and fluency. On sentiment control, BOLT is 7x faster than competitive baselines, and more fluent in 74.4% of the evaluation samples according to human judges.

## introduction generating text using pre-trained language models (plms) to satisfy user-specified constraints is an important task to allow practical usage of plms. common controlled text generation methods include training conditional language models @xcite or attribute-based fine-tuning of plms @xcite . yet, these methods are often resource-intensive and infeasible for large models like @xcite . furthermore, these methods assume access to large amounts of attribute-specific data and are inflexible for new constraints. on the contrary, inference-time methods @xcite @xcite fine-tuning. in particular, energy-based models (ebms) @xcite have demonstrated greater flexibility, since they can accommodate arbitrary energy functions @xcite @xcite . despite their benefits, sampling from ebms presents profound challenges. notably, the sampling process, which is often done through langevin dynamics @xcite or gibbs sampling @xcite , requires a substantial number of iterations to converge to readable sequences of text. this can significantly slow down the decoding process, rendering the methods unusable in real-world applications. in this paper, we propose bolt 1 , that uses a sequence of tunable biases over logits of the plm's output layer, to steer the generation towards specified constraints. the biases are tuned through a gradient-based process, with the goal of minimizing the energy of the generated sequences. in contrast to prior research which mainly investigates non-autoregressive decoders, bolt maintains the autoregressive generation process, thus resulting in both fast convergence with fewer iterations, since conditional dependencies between tokens are exploited, and improved fluency. fig. 1 demonstrates that the sampling process of recent ebm-based methods-mucola @xcite , mix&match @xcite , and cold @xcite -is slower on a sentiment control task, e.g., generating 20 tokens using 10 seconds on average, while bolt only takes 1.4 seconds. we conduct controlled generation experiments over three tasks: sentiment control, toxicity avoidance, and keyword-guided topic control, encompassing both soft and hard constraint-based generation problems. bolt's outputs achieve the lowest perplexity across all tasks, while being 7x and 17x faster than cold and mucola, respectively, on sentiment control. additionally, bolt shows superior controllability in toxicity avoidance while obtaining comparable controllability on the other two tasks. lastly, according to human evaluation, 74.4% and 51.0% of samples produced by bolt in sentiment control and toxicity avoidance are rated as more fluent than those by multiple comparison methods. ## related work popular methods for controlled generation often rely on attribute-conditioned language modeling @xcite , model fine-tuning @xcite , or prompt tuning @xcite , all requiring intensive model training and attribute-specific data. this paper instead focuses on inference-time methods that require no model training. prior work under this paradigm mainly adjusts the output token probabilities toward constraint-satisfying sequences @xcite . for instance, @xcite leverage gradients from an attribute classifier to update the lm hidden state to guide the generation. however, one notable drawback of such techniques is the requirement of learning specialized models such as attribute classifiers @xcite and future-aware classifiers @xcite . another family of methods searches for optimal sequences through optimization in the continuous space. for instance, mucoco @xcite uses constrained continuous optimization, solved by lagrangian multipliers and gradient descent. @xcite further enhance the gradient-based energy functions one-hot vector tunable bias logits from plm + ste + + ste ste + + ste ste step 1 step 𝒕-1 step 𝒕 step 𝒕+1 step 𝑳 ## the bolt model energy-based controlled generation aims to produce a sequence of tokens that minimize an energy function, with lower energy indicating more constraints being satisfied @xcite . while sampling techniques such as rejection sampling can be used to sample lowenergy sequences @xcite , such sampling requires the usage of an appropriate proposal distribution and is typically slow in practice. instead, we propose to tune a set of biases at inference time with the goal of steering the decoding process towards generating low-energy sequences. the overview of our framework is displayed in fig. 2 . at each decoding step t, we add the tunable bias y b @xmath0: therefore, we model the weights using a decreasing linear function of t, i.e., w @xmath1, which is found to work best in practice. 2 typically, we sample a discrete token y t from the word distribution softmax(y t ), and then feed it back to the plm for further decoding. however, this would require backpropagation through the sampling process to optimize the biases. as a workaround, we use the straightthrough gradient estimator (ste) @xcite , which converts y t to a one-hot vector ȳt in the forward pass and bypasses ȳt in the backward pass to allow gradients to be applied to y t . foot_2 ȳt designates the argmax token, i.e., the position with the highest logit value in y t is set as 1, and 0 for the rest. the one-hot vector ȳt is fed to the plm for next-step decoding. after decoding for l steps, we obtain a sequence of one-hot vectors ȳ[1:l] =[ȳ 1 , ȳ2 , ..., ȳl-1 , ȳl ]. then, we update y b t with gradient descent to minimize the energy function e(ȳ [1:l] ). 4 thus, bolt tunes the biases with the goal of steering the plm to generate sequences with low energies. finally, the output sentence [y 1 , y 2 , ..., y l-1 , y l ] can be derived from ȳ[1:l] through multiple iterations of gradient descent until the constraints are satisfied (e.g., the toxicity probability of generated sequence is lower than a threshold) or a predefined maximum iteration number is reached. ## conclusion we introduce bolt, an energy-based model for controlled text generation. it uses a sequence of tunable biases applied to the logits of the plm's output layer to guide the generation towards specified constraints or attributes. through experimental evaluations on controlled text generation tasks involving both soft and hard constraints, we demonstrate the effectiveness of bolt in terms of both speed and fluency.