Selective Gradient Masking
Free while signed in. Answers cite the passages they came from.
Anthropic researchers present Selective Gradient Masking (SGTM), a technique that removes dangerous capabilities like CBRN knowledge from language models during pretraining while preserving general capabilities. Unlike data filtering, SGTM localizes target knowledge into dedicated "forget" parameters that can be zeroed out after training. - **Absorption mechanism:** SGTM splits parameters into forget and retain components, with gradients masked so only forget parameters update on labeled dangerous content. Unlabeled dangerous content naturally gravitates toward forget parameters through self-reinforcing "absorption," providing robustness to imperfect labeling. - **Recovery resistance:** Traditional unlearning methods (RMU) recovered removed biology knowledge in 50 fine-tuning steps. SGTM required 350 steps - 7x more resistant than RMU and matching the robustness of models trained with perfect data filtering. - **Retain/forget trade-offs:** On Wikipedia biology experiments with 254M parameter models, SGTM achieved superior trade-offs compared to both weak and strict data filtering, retaining more knowledge from adjacent fields like medicine and chemistry with only 5% compute penalty. - **Mechanistic validation:** Gradient analysis on bilingual data showed forget parameters develop higher gradient norms for forget-domain content while retain parameters specialize for general content, with this localization strengthening at larger scales.
Absorption mechanism: SGTM splits parameters into forget and retain components, with gradients masked so only forget parameters update on labeled dangerous content. Unlabeled dangerous content naturally gravitates toward forget parameters through self-reinforcing “absorption,” providing robustness to imperfect labeling.
Recovery resistance: Traditional unlearning methods (RMU) recovered and removed biology knowledge in 50 fine-tuning steps. SGTM required 350 steps - 7x more resistant than RMU and matching the robustness of models trained with perfect data filtering.
Retain/forget trade-offs: On Wikipedia biology experiments with 254M parameter models, SGTM achieved superior trade-offs compared to both weak and strict data filtering, retaining more knowledge from adjacent fields like medicine and chemistry with only 5% compute penalty.
Mechanistic validation: Gradient analysis on bilingual data showed forget parameters develop higher gradient norms for forget-domain content while retain parameters specialize for general content, with this localization strengthening at larger scales.
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.
Get next week’s papers.
The same picks and the same summaries, in your inbox. Free, and 176 issues deep.
Subscribe on Substack