BLUR: A Bi-Level Optimization Approach for LLM Unlearning
Hadi Reisizadeh, Jinghan Jia, Zhiqi Bu, Bhanukiran Vinzamuri, Anil Ramakrishna, Kai-Wei Chang, Volkan Cevher, Sijia Liu, Mingyi Hong
TL;DR
This paper tackles the problem of unlearning in large language models by introducing BLUR, a bi-level optimization framework that prioritizes forgetting over retaining utility. By modeling unlearning as a hierarchical problem where the lower-level forget objective constrains the solution space and the upper-level retain objective selects the best utility-preserving solution, BLUR provides a gradient-based update that orthogonally projects retain information away from forgetting directions. The authors prove convergence properties for nonconvex settings and demonstrate through extensive experiments on TOFU, MUSE, and WMDP benchmarks that BLUR outperforms state-of-the-art baselines in unlearning efficiency while preserving model utility, with ablations highlighting the impact of hyperparameters. The work offers a principled path toward compliant and ethical LLM unlearning and introduces a flexible meta-algorithm that can accommodate various forget and retain losses, potentially influencing future standardizations in data-removal and privacy-preserving AI.
Abstract
Enabling large language models (LLMs) to unlearn knowledge and capabilities acquired during training has proven vital for ensuring compliance with data regulations and promoting ethical practices in generative AI. Although there are growing interests in developing various unlearning algorithms, it remains unclear how to best formulate the unlearning problem. The most popular formulation uses a weighted sum of forget and retain loss, but it often leads to performance degradation due to the inherent trade-off between forget and retain losses. In this work, we argue that it is important to model the hierarchical structure of the unlearning problem, where the forget problem (which \textit{unlearns} certain knowledge and/or capabilities) takes priority over the retain problem (which preserves model utility). This hierarchical structure naturally leads to a bi-level optimization formulation where the lower-level objective focuses on minimizing the forget loss, while the upper-level objective aims to maintain the model's utility. Based on this new formulation, we propose a novel algorithm, termed Bi-Level UnleaRning (\texttt{BLUR}), which not only possesses strong theoretical guarantees but more importantly, delivers superior performance. In particular, our extensive experiments demonstrate that \texttt{BLUR} consistently outperforms all the state-of-the-art algorithms across various unlearning tasks, models, and metrics. Codes are available at https://github.com/OptimAI-Lab/BLURLLMUnlearning.
