Attribute Controlled Fine-tuning for Large Language Models: A Case Study on Detoxification (2024)

Abstract: We propose a constraint learning schema forfine-tuning Large Language Models (LLMs)with attribute control. Given a training corpusand control criteria formulated as a sequence-level constraint on model outputs, our methodfine-tunes the LLM on the training corpus whileenhancing constraint satisfaction with minimalimpact on its utility and generation quality.Specifically, our approach regularizes the LLMtraining by penalizing the KL divergence be-tween the desired output distribution, which sat-isfies the constraints, and the LLM’s posterior.This regularization term can be approximatedby an auxiliary model trained to decomposethe sequence-level constraints into token-levelguidance, allowing the term to be measuredby a closed-form formulation. To further im-prove efficiency, we design a parallel schemefor concurrently updating both the LLM andthe auxiliary model. We evaluate the empiricalperformance of our approach by controlling thetoxicity when training an LLM. We show thatour approach leads to an LLM that producesfewer inappropriate responses while achievingcompetitive performance on benchmarks and atoxicity detection task

## introduction large language models (llms) have demonstrated impressive performance across a variety of tasks which has led to their widespread adoption for a multitude of ai applications. however, they carry the risk of producing inappropriate, unsafe, unfair outputs @xcite @xcite ideally, llms should learn to comply with constraints and policies specified by users. for example, in a userfacing application like a chatbot, llms should never generate toxic or offensive responses, nor to divulge sensitive information. while there are several post hoc methods to moderate llm outputs @xcite @xcite , they lack an efficient and principled approach to training llms to adhere to constraints. we start by defining a sequence-level oracle as a function that takes an llm's output and adjudicates whether it satisfies a predefined set of attribute constraints. in practice, the oracle can be a rule-based, model-based, or mixed system (e.g., a classifier that decides whether a sentence is toxic). given a pre-trained llm and the oracle, we aim to fine-tune an llm to achieve the following: 1) attribute control: the llm output passes the oracle with a high probability. 2) utility preservation: the llm maintains performance comparable to the original llm on utility benchmarks. 3) training efficiency: the cost of fine-tuning with attribute control is similar to that of the typical fine-tuning. while existing approaches can meet some of these criteria, achieving all of them is challenging. for example, filtering training data with the oracle function before fine-tuning @xcite is a simple and efficient method. however, this approach could be less effective. taking toxicity control as an example, if we filter out the toxic data from a fine-tuning corpus, in a regular context the model will learn not to generate toxic contents. nevertheless, it might still be possible to trigger the generation of offensive responses given a toxic prompts, due to the fact that toxic prompts are out-of-distribution in relation to the fine-tuning corpus. another promising approach is reinforcement learning (rl) considering controlling criteria in the reward function @xcite . however, rl setups tend to be inefficient and require preference data generation which adds significant overhead in comparison to generic fine-tuning. in this work, we propose a novel solution to training an llm with a set of attribute constraints. inspired by the classic idea of constraint-driven learning @xcite and posterior regularization @xcite , we incorporate constraints as a regularizer in fine-tuning. specifically, we estimate the closest distribution from the current model that satisfies the constraints and penalize the gap from the current model distribution to this estimated distribution to regularize the llm during fine-tuning. we iterate through this process to push the llm closer to the feasible region of generations, making the estimation progressively more accurate. this iterative fine-tuning process updates the base llm and regularizer sequentially, causing run time to be significantly longer than the typical fine-tuning. thus, we parallelize our algorithm by updating the base llm and regularizer simultaneously based on their status in the last iteration. empirically, the parallelization achieves the same level of performance compared to sequential finetuning, and the time complexity is the same as a typical fine-tuning approach. to validate the effectiveness of our proposed method, we conduct a case study in detoxification, considering three scenarios involving different datasets. in the first scenario, we fine-tune llms on datasets rich in toxic language with an attribute control that prevents the generation of toxic outputs. our approach successfully passes stress tests and produces responses with lower toxicity compared to all baseline models. in the second scenario, we explore whether the attribute control can retain the utility of the llm while reducing the toxicity of its responses. training only on a small dataset will lead to catastrophic forgetting. therefore, we fine-tune the llm on a mix of data comprising toxigen @xcite and wikitext @xcite datasets with attribute control. our method demonstrates the best balance between model utility and toxicity management compared to similar techniques. finally, we assess whether the llm can effectively identify toxic content without generating it, a critical skill since the model must recognize toxic elements to avoid producing them. in standard fine-tuning, these goals often conflict: the model learns to identify toxicity through training on a toxic corpus, which paradoxically increases the generation of toxic content. however, our method successfully mitigates the generation of toxic content while maintaining classification performance on par with traditional fine-tuning techniques. we summarize our contributions as follows: ## related work prior work exists on the controlled generation problem and it can be divided into two fronts. solutions that apply at inference time during decoding, and solutions that apply at fine-tuning. attribute controlled decoding for llms several methods have been explored to control llm generation during decoding. some prominent methods include activation editing @xcite which adjusts the activation vectors in the llm, and weight editing @xcite which adjusts the weights in llm. dathathri et al. (2020) (pplm) leverages an auxiliary model to steer the base llm distribution. following this line of work, krause et al. (2021) (gedi) and liu et al. (2021) (dexpert) used contrastive learning as an objective to achieve attribute control during decoding. yang and klein (2021) (fudge) leveraged an external token-level auxiliary model for their work. this was followed by (meng et al., 2022b) (nado) who trained a token-level auxiliary model by decomposing the controlling criteria via optimization and approximation. zhang et al. (2023) (gelato) leverage a probabilistic circuit to tractably incorporate symbolic constraints. our work is inspired by nado @xcite , and we take it as a sub-component in our fine-tuning approach. however, our paper significantly differs from nado. firstly, nado presents an inference method that guides generation by reweighting output distributions without updating the model weights of the base model. in contrast, our approach involves fine-tuning the model. we use nado to estimate the optimal distribution that satisfies the constraints. however, our design of the posterior regularizer, the iterative-updated scheme, and the parallel computing algorithm is novel. moreover, the effectiveness of regularizing the polygon is representing the feasible region q where the constraints are satisfied. on (a) it shows the regularizer term is defined as the closest distance from p θ to q. regularized by kl-divergence from q, on (b) we show the llm distribution p θ is gradually pushed towards the feasible region. attribute controlled fine-tuning for llms when controlling attributes during fine-tuning, the most straightforward way is to filter out training data that contain the attribute @xcite . however, this approach tends to cause performance degradation as it can filter out large portions of the training set and does not actively leverage negative examples. neuro-symbolic approach @xcite incorporates symbolic constraints as loss added to the training objective; however, it cannot handle complex or implicit constraints. rl @xcite @xcite can be utilized to control an attribute via the use of an attribute-related reward. rlhf @xcite @xcite and rlaif @xcite leverage feedback from humans and llms, respectively, to control the required attributes. however, they are more focused on being aligned with humans (or llms) rather than specific attribute control. rl-based methods are effective but often inefficient due to the large variance in feedback provided by the reward models. ## case study on detoxification to test the effectiveness of the proposed approach, we apply it to detoxify an llm. toxicity, as discussed in section 1, is of significant importance as a metric for the evaluation of llm @xcite @xcite . in all experiments, we set δ = 1 to set toxicity as a hard constraint. detailed notes on data pre-processing, hyper-parameter choice for model training, and the architecture of auxiliary models can be found in the appendix. ## conclusion we propose a novel fine-tuning approach for attribute control in llm generations and we demonstrate its effectiveness using toxicity as our chosen attribute. while this work focuses on toxicity, our approach is general enough to accommodate other types of attributes as well. with adaptive regularizers, our method can further extend to control multiple attribute across various domains.