Beyond Data Filtering: Knowledge Localization for Capability Removal in LLMs

2025-12-05

AI Summary

Dual-use risks of large language models (LLMs) are increasingly prominent, yet data filtering during pretraining faces fundamental challenges due to high annotation costs and extreme sensitivity to label noise—where even a few mislabeled examples can induce harmful capabilities. This paper proposes Selective GradienT Masking (SGTM), a gradient-masking-based method for knowledge localization and localized parameter updating that enables precise isolation and removal of specific harmful knowledge (e.g., bilingual, encyclopedic, or biomedical capabilities) during pretraining—without requiring high-quality annotations. SGTM enhances gradient routing to achieve knowledge-level controllable forgetting. It demonstrates strong robustness under both label noise and adversarial fine-tuning: it achieves more complete forgetting than data filtering and Rank-1 Model Unlearning (RMU); moreover, adversarial recovery requires seven times more optimization steps, significantly improving model security.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Natural Language Processing: Safety and Robustness
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Large Language Models increasingly possess capabilities that carry dual-use risks. While data filtering has emerged as a pretraining-time mitigation, it faces significant challenges: labeling whether data is harmful is expensive at scale, and given improving sample efficiency with larger models, even small amounts of mislabeled content could give rise to dangerous capabilities. To address risks associated with mislabeled harmful content, prior work proposed Gradient Routing (Cloud et al., 2024) -- a technique that localizes target knowledge into a dedicated subset of model parameters so they can later be removed. We explore an improved variant of Gradient Routing, which we call Selective GradienT Masking (SGTM), with particular focus on evaluating its robustness to label noise. SGTM zero-masks selected gradients such that target domain examples only update their dedicated parameters. We test SGTM's effectiveness in two applications: removing knowledge of one language from a model trained on a bilingual synthetic dataset, and removing biology knowledge from a model trained on English Wikipedia. In both cases SGTM provides better retain/forget trade-off in the presence of labeling errors compared to both data filtering and a previously proposed instantiation of Gradient Routing. Unlike shallow unlearning approaches that can be quickly undone through fine-tuning, SGTM exhibits strong robustness to adversarial fine-tuning, requiring seven times more fine-tuning steps to reach baseline performance on the forget set compared to a finetuning-based unlearning method (RMU). Our results suggest SGTM provides a promising pretraining-time complement to existing safety mitigations, particularly in settings where label noise is unavoidable.

Problem

• Removing dual-use capabilities from LLMs safely
• Localizing harmful knowledge into removable parameters
• Improving robustness to labeling errors in unlearning

Innovation

• Selective GradienT Masking (SGTM) zero-masks selected gradients
• SGTM localizes target knowledge into dedicated model parameters
• SGTM provides robustness to label noise and adversarial fine-tuning