SOUL : Unlocking the Power of Second-Order Optimization for LLM Unlearning (2024)

Abstract: Large Language Models (LLMs) have highlighted the necessity of effective unlearning mechanisms to comply with data regulations and ethical AI practices. LLM unlearning aims at removing undesired data influences and associated model capabilities without compromising utility beyond the scope of unlearning. While interest in studying LLM unlearning is growing, the impact of the optimizer choice for LLM unlearning remains unexplored. In this work, we shed light on the significance of optimizer selection in LLM unlearning for the first time, establishing a clear connection between second-order optimization and influence unlearning (a classical approach using influence functions to update the model for data influence removal). This insight propels us to develop a second-order optimization-based LLM unlearning framework, termed Second-Order UnLearning (SOUL), which extends the static, one-shot model update using influence unlearning to a dynamic, iterative unlearning process. Our extensive experiments show that SOUL consistently outperforms conventional first-order methods across various unlearning tasks, models, and metrics, indicating that second-order optimization offers an effective and broadly applicable solution for LLM unlearning.

## introduction llms have emerged as transformative technology, greatly enhancing natural language processing capabilities from text generation to simulating humanlike interactions @xcite . while offering substantial benefits, llms also present challenges, such as the risk of misuse in generating private, toxic, or illegal content @xcite @xcite , perpetuation of biases @xcite , and the potential for aiding in developing cyberattacks or bioweapons @xcite . to address the aforementioned risks, the problem of llm unlearning arises, aimed at eliminating specific undesirable data influences and their corresponding model generation capabilities while ensuring that model utility is not compromised out of the unlearning scope @xcite @xcite @xcite @xcite @xcite . while the concept is appealing, the development of effective unlearning algorithms remains challenging. a straightforward approach involves retraining the model from scratch after removing the undesired training data, driven by data privacy concerns @xcite . however, this method is impractical due to the extremely high cost associated with retraining llms from scratch. therefore, model fine-tuning under a predefined unlearning objective has become the primary approach to solve most llm unlearning problems @xcite @xcite . unfortunately, there is a lack of effective fine-tuning techniques for llm unlearning. for example, classical gradient ascent-based fine-tuning techniques are susceptible to over-forgetting, which can hamper the original model utility @xcite @xcite . conversely, less aggressive fine-tuning techniques, such as fine-tuning solely on the retain set (i.e., the data set irrelevant to the forgetting data points) @xcite , could result in under-forgetting, failing to completely erase the influence of forgotten data. as a result, it is hard to strike the optimal balance between unlearning effectiveness and model utility preservation. several recent efforts have been made to develop improved model fine-tuning techniques for llm unlearning. for example, studies have delved into designing fine-tuning loss functions tailored for llm unlearning @xcite @xcite . a currently popular choice is the regularized optimization objective that integrates unlearning efficacy loss with model utility loss, as seen in approaches such as the gradient difference (graddiff) @xcite @xcite , preference optimization (po) @xcite and negative preference optimization (npo) @xcite . additionally, other llm unlearning techniques incorporate the model's prior into fine-tuning. for instance, finetuning is selectively applied to a subset of model units deemed essential for the unlearning task @xcite . this approach has led to the emergence of localization-informed llm unlearning @xcite . furthermore, input prompt strategies have been employed, enabling unlearning through model queries and/or adjusting only a small fraction of parameters @xcite @xcite . despite the recent progress of llm unlearning, the majority of existing fine-tuning-based approaches have relied on first-order (fo) optimization to conduct unlearning. to our knowledge, there have been no prior studies that specifically investigate llm unlearning from the perspective of optimizer design. in this work, we unveil the power of second-order (so) optimizer in llm unlearning and demonstrate its superiority over fo optimizer in various fine-tuning scenarios. we term the second-order optimization-based unlearning framework as soul (second-order unlearning). we will show that soul not only offers a viable approach for enhancing unlearning efficacy but also stays effective in preserving model utility. such an optimizer-induced advantage holds consistently across various llm unlearning objectives and formulations, providing a generic improvement. we summarize our contributions below. machine unlearning for non-llms. the concept of machine unlearning has emerged from data protection regulations, such as the 'right to be forgotten' @xcite , which were initially not specifically targeted at llms @xcite @xcite . as the field has progressed, the applications of machine unlearning have rapidly expanded into diverse areas such as image classification @xcite @xcite ), text-to-image and image-to-image generation @xcite @xcite @xcite , and federated learning @xcite . in the literature, retraining a model from scratch by excluding forgotten data points has been considered as 'exact' unlearning @xcite @xcite . however, the significant computational costs associated with retraining from scratch and the need for access to full training data have spurred the development of scalable and efficient 'approximate' unlearning techniques @xcite @xcite @xcite . additionally, some methods provide provable and certified data removal, often employing differential privacy to ensure compliance and verifiability @xcite @xcite . llm unlearning. the exploration of machine unlearning in the context of llms has garnered increasing interest @xcite @xcite @xcite @xcite @xcite . seminal works by @xcite and @xcite have elucidated the need for machine unlearning within llms, delineating clear motivations from both application-centric and regulatory standpoints. some research efforts @xcite @xcite @xcite have concentrated on employing gradient ascent to facilitate forgetting in targeted datasets. other studies such as those by @xcite ; @xcite have examined preference optimization, crafting alternative responses (e.g., reject) to realize unlearning. in addition, some unlearning methods have explored and exploited the data-model interactions that could affect llm unlearning @xcite @xcite , such as weight localization-informed unlearning @xcite , and altering the hidden representations of llms to achieve unlearning @xcite . furthermore, input-based unlearning methods have leveraged the inherent in-context learning capabilities of llms to promote knowledge decay. for instance, @xcite developed system prompts that instruct models to avoid generating unwanted knowledge, while @xcite applied in-context learning strategies to address unlearning. last but not least, some recent benchmarks have been developed for the evaluation of llm unlearning, such as tofu for fictitious unlearning @xcite and wmdp for unlearning hazardous knowledge in llms @xcite . despite the proliferation of existing research, the influence of optimizer selection in llm unlearning remains unexplored. ## primer on llm unlearning problem setup. llm unlearning aims to mitigate the influence of undesired data, such as sensitive or copyrighted information, and/or restrict the model's capabilities to avoid the associated content generation. this process also requires preserving the llm's utility for unrelated tasks and avoiding full retraining to maintain computational efficiency. following the generic formulation of llm unlearning in @xcite , the unlearning problem can be conceptualized as removing the influence of a designated 'unlearning target'-whether it pertains to data, knowledge, or model capabilitiesfrom a pre-trained llm (denoted as θ o ). the unlearning target is typically specified by a forget set d f , which includes the information or knowledge intended for removal. to preserve the llm's generation capability (i.e., utility) after unlearning, a retain set d r is also introduced. this set comprises data that is irrelevant to the unlearning target. given the aforementioned setup, the problem of llm unlearning is often formulated as a regularized optimization problem, fine-tuned from θ o over the forget set d f and the retain set d r : here ℓ f and ℓ r represent the forget loss and the retrain loss respectively, and λ ≥ 0 is a regularization parameter to strike a balance between unlearning and utility preservation. note that problem (1) is not the only formulation of llm unlearning. yet, it remains the prevailing mainstream formulation in the field, although there have been research efforts to explore the optimization-free based methods, such as in-context learning or input-level prompting @xcite . some specifics of llm unlearning @xcite . while problem (1) may appear as a straightforward optimization task initially, complexities arise in determining the effective forget loss ℓ f and achieving the optimal balance between unlearning and utility. these questions remain challenging in the literature. we present three representative llm unlearning approaches and illustrate how they relate to the specifics of problem @xcite . at λ = 0, problem (2) simplifies to maximizing the training loss on forget set. this method is known as gradient ascent (ga) @xcite . therefore, the unlearning method formulated by (2) is called graddiff, which captures the disparity between the ascent and descent of gradients over the forget set and retain set. ## second-order optimization to enhance llm unlearning: why & how in this section, we shed light on a missing factor of llm unlearning: the choice of optimizer, which has been overlooked in the literature yet crucial for the effectiveness of unlearning. gaining insights from influence unlearning. influence unlearning is a one-shot machine unlearning technique that utilizes the influence function approach @xcite to assess and quantify the impact of the forget set d f on the pre-trained model θ o . diverging from iterative optimization approaches like graddiff (2) and po (3), influence unlearning involves a single weight modification step, updating θ o based on the influence exerted by the forget set on the weight space. while influence unlearning is a classic technique, its usage has been limited to vision tasks and small models @xcite . even within the realm of vision tasks, it is not deemed a state-of-the-art (sota) approach to unlearning @xcite . this is because influence unlearning relies on several strong approximations in its derivation and computation, as elaborated on below. let θ mu denote a retrained model from scratch on the retain set d r , i.e., the solution to the optimization problem min θ e (x,y)∈dr [ℓ(y|x; θ)] with random initialization, where ℓ is the training loss introduced in (2). the objective of influence unlearning is to derive the weight modification from the pre-trained model θ o to the retrained model θ mu , i.e., θ muθ o . to this end, a weighted training problem is introduced: here 1 dr denotes an element-wise indicator function that takes the value 1 if the data point belongs to the retain set d r and 0 otherwise. based on (5), influence unlearning then aims to derive: the derivation of ( 6 ) is highly non-trivial as the retrained model θ(w mu ) cannot be directly obtained and is implicitly defined through the optimization problem min θ ℓ(θ, w mu ). to proceed, the influence function approach @xcite @xcite simplifies (6) by applying a first-order taylor expansion to θ(w mu ) at @xmath0: as evident from the above derivations, influence unlearning encounters two primary limitations that hinder its application to llm unlearning: the computational complexity associated with inverting the hessian matrix, and the diminished accuracy stemming from approximations utilized in taylor expansion and second-order information acquisition. an intriguing observation from ( 8 ) is that influence unlearning conforms to the generic form of so optimization @xcite . as in newton's method, one uses a so approximation of a loss function ℓ to locate its minima. this yields a descent algorithm based on a newton step @xcite : the consistency observed in the formats of influence unlearning (8) and second-order optimization (9) prompts us to consider whether we can integrate second-order optimization into influence unlearning, thereby transforming the latter into an effective iterative unlearning approach. soul: second-order unlearning for llms. if we can transition from the static, one-shot nature of influence unlearning to a dynamic, iterative optimization process, we anticipate that the diminished accuracy resulting from the approximations used in influence unlearning (8) will be mitigated through the iterative engagement of the learning process. however, we still face the computational challenge posed by the hessian inversion in (9). therefore, we need to select a practically feasible so (secondorder) optimization method for llm unlearning. sophia (second-order clipped stochastic optimization) @xcite , a simple scalable so optimizer, is well-suited since it utilizes a simple diagonal matrix estimate of the hessian and has shown its effectiveness in llm pre-training. sophia modifies the vanilla newton's method to next, we link influence unlearning (8) with the so optimizer and propose the so unlearning approach. recall from ( 8 ) and (4) that the change in data weights (1w mu ) encodes the influence of the forget set d f in model training. therefore, we can interpret the term h -1 ∇ θ ℓ(θ 0 , 1w mu ) in (8) as a second-order optimization-based ascent step over the forget set. this contrasts with the original sophia update (10), which executes the descent using the clipped newton step. let us take graddiff (2) as an example. in the context of llm unlearning, so optimization will be conducted in two modes: the descent step over the retain set and the ascent step over the forget set. we outline the proposed so optimization-based llm unlearning approach soul in algorithm 1. algorithm 1 soul to solve problem (2) 1: initialize: θ0 = θo, @xmath1, @xmath2, @xmath3, learning rates {ηt}, and ema parameters β1 and β2 2: for @xmath4: estimate hessian diagonal ĥt-1 as sophia at θt-1, 6: based on mt and ht, update θ based on (10): when considering po-type problems like (3), the proposed algorithm can only operate in the descent mode. this is because the preference (i.e., the unlearning response y f ) has already been defined, and the corresponding forget loss is minimized rather than maximized in (2). in this scenario, soul optimizes both forget loss and retain loss through descent mode unification. this makes soul's implementation same as sophia. ## conclusions in this paper, we investigate the role of optimizer choice in llm unlearning, linking second-order optimization to influence unlearning. building on this, we propose a second-order llm unlearning framework, agnostic to loss function, to augment existing approaches. extensive experiments across various unlearning tasks, models, and metrics consistently show the superiority of second-order unlearning. these results advocate for the development and adoption of optimizers tailored for effective llm unlearning.