SHA 256 at S em E val-2025 Task 4: Selective Amnesia – Constrained Unlearning for Large Language Models via Knowledge Isolation (2025)

Abstract: Large language models (LLMs) frequently memorize sensitive information during training, posing risks when deploying publicly accessible models. Current machine unlearning methods struggle to selectively remove specific data associations without degrading overall model capabilities. This paper presents our solution to SemEval-2025 Task 4 on targeted unlearning, which introduces a two-stage methodology that combines causal mediation analysis with layer-specific optimization. Through systematic causal tracing experiments on OLMo architectures (1B and 7B parameters), we identify the critical role of the first few transformer layers (layers 0-5) in storing subject-attribute associations within MLP modules. Building on this insight, we develop a constrained optimization approach that freezes upper layers while applying a novel joint loss function to lower layers—simultaneously maximizing forget set loss via output token cross-entropy penalties and minimizing retain set deviation through adaptive regularization. Our method achieves 2nd place in the 1B model track, demonstrating strong task performance while maintaining 88% of baseline MMLU accuracy. These results establish causal-informed layer optimization as a promising paradigm for efficient, precise unlearning in LLMs, offering a significant step forward in addressing data privacy concerns in AI systems.

## introduction large language models (llms), pretrained on massive datasets via self-supervised learning, often inadvertently memorize sensitive information @xcite . this can include personally identifiable information such as email and home addresses, social security numbers linked to individual names, and even copyrighted creative content @xcite @xcite . the widespread opensourcing of these models raises concerns about the potential exposure and misuse of such data @xcite @xcite . while retraining with filtered datasets is a viable solution, the need for frequent updates to address newly discovered vulnerabilities or comply with evolving data privacy regulations (e.g., "right to be forgotten" requests @xcite ) makes this approach prohibitively expensive. each full retraining requires significant computational resources, time, and energy, leading to a substantial economic and environmental burden. unlearning offers a promising solution by allowing models to remove or modify specific information without full retraining @xcite . unlearning methods seek to efficiently update llms by altering the model in a way that eliminates unwanted information while minimizing the impact on the model's overall performance and capabilities @xcite . however, current unlearning solutions often struggle to balance effective unlearning with preserving the model's general usefulness @xcite . this is largely due to their broad, non-selective application of unlearning techniques, which can unintentionally erase useful information. further, these methods may be vulnerable to membership inference attacks (mia) @xcite , and exhibit difficulty in preserving knowledge within the retain set while effectively unlearning the forget set. to address these limitations and foster research into more effective and robust unlearning strategies, semeval 2025 task 4, unlearning sensitive content from large language models @xcite , challenges participants to develop methods that can selectively remove sensitive information from llms while preserving their core capabilities. in this work, we address the challenge of tar-geted unlearning by first performing knowledge isolation using causal mediation analysis @xcite . causal mediation analysis helps identify the specific layers within the llm responsible for storing the factual knowledge to be unlearned. through experiments with the provided fine-tuned olmo models @xcite (both 1b and 7b parameter versions, fine-tuned by the task organizers to memorize the forget and retain sets), we empirically determine that the initial layers (specifically layers 0-5) have a disproportionately high impact on factual recall. our approach combines targeted knowledge removal with a novel joint loss function. by focusing on causally identified lower layers (layers 0-5) and using cross-entropy loss on output tokens, we aim to disrupt specific subject-attribute associations while preserving overall model performance. this method seeks to achieve effective and efficient unlearning of sensitive content in llms by isolating knowledge, applying carefully designed loss functions, and implementing targeted parameter updates. our method achieves 2nd place in the 1b model track with a with a final score of 0.652, demonstrating a strong task aggregate performance (0.973) while maintaining 88% of baseline mmlu accuracy. the 7b variant shows comparable forget set eradication (0.964 task score) but highlights scalability challenges through a 46% mmlu decrease, underscoring the need for layer-specific capacity analysis in larger models. these findings underscore the potential of causal-informed layer freezing as a promising technique for efficient and precise unlearning in large language models. ## system overview machine unlearning in llms is a highly challenging task, and most techniques can not guarantee complete erasure of target knowledge @xcite . nevertheless, we propose that a certain level of unlearning while preserving the model's general ability can be achieved by identifying the hidden layers responsible for factual recall. we therefore divide our unlearning approach into two steps. first, we employ causal mediation analysis @xcite to identify layers, critical for storing factual information. next, we optimize the hidden parameters of these layers using a loss function that simultaneously penalizes regurgitation of the forget-set and deviations in retain-set performance. ## experiments this section details the analysis of our main evaluation results in the leaderboard, as well as our parameter selection process. ## conclusion our systematic investigation establishes that targeted unlearning in large language models can be significantly enhanced through causal-informed layer optimization. combining causal mediation analysis with constrained parameter updates to mlp modules in early transformer layers (0-5), we demonstrate precise eradication of sensitive subject-attribute associations while preserving 88% of baseline general knowledge performance in olmo-1b models. the proposed joint loss formulation-simultaneously applying cross-entropy penalties on forget set outputs and adaptive regularization for retain set preservation-achieves a high task aggregate score of 0.973. while the 7b variant maintained comparable forget set removal efficacy (0.964 task score), its 46% mmlu degradation underscores the critical need for more robust mechanistic interventions. these findings suggest three key implications for machine unlearning research: first, that mlp layers in early trans-former blocks serve as preferential targets for factual knowledge modification; second, that output token cross-entropy provides a more surgical intervention than full-sequence loss calculations; and third, that layer freezing thresholds must scale nonlinearly with model depth to maintain utility. future work should investigate constrained gradient updates through kl divergence to reduce the drop in model utility after unlearning. our results cement causal mediation as a vital tool for developing compliant, adaptable llms that meet evolving data privacy requirements without full retraining.