CE-U: Cross Entropy Unlearning

Large language models memorize sensitive data from their pretraining corpora Jang et al. (2023). In this work, we propose CE-U (Cross Entropy Unlearning), a loss function for unlearning. CE-U addresses fundamental limitations of gradient ascent approaches that suffer from vanishing gradients when model confidence is high and exploding gradients when confidence is low. We also unify standard cross entropy learning and unlearning into a single framework. On the TOFU benchmark for unlearning Maini et al. (2024), CE-U achieves state-of-the-art results on LLaMA2-7B models without using an extra oracle model or additional positive samples. Our analysis reveals that the problematic gradient ascent component also exists in reinforcement learning algorithms like DPO Rafailov et al. (2023) and GRPO Shao et al. (2024). This suggests that applying CE-U approach to reinforcement learning could be promising to improve stability and convergence.