Mask2Shield: Strengthening LLM Safety against Neuron-Pruning Attacks

2026-07-24

AI Summary

Despite safety alignment, large language models remain vulnerable to neuron pruning attacks, where removing only a few critical units can compromise their ability to refuse harmful requests. To address this, this work proposes Mask2Shield (M2S), the first method to integrate mask-based training into the safety alignment process. M2S performs masked forward propagation under simulated functional pruning and employs teacher–student distillation, compelling the student model to recover safe refusal capabilities using the remaining computational resources while preserving general performance through a frozen, intact teacher model. Experiments across ten model configurations show that successful attacks drop from 80–279 to 1–44 out of 313 adversarial prompts, with minimal degradation on four standard capability benchmarks, demonstrating significantly enhanced robustness against structured pruning attacks.

Abstract

Large language models (LLMs) are safety-aligned before deployment to reduce harmful content generation. Yet neuron-level pruning attacks show that refusal can depend on a small set of removable units: disabling them can remove safety behavior while leaving much of the model usable. To address this problem, we introduce Mask2Shield (M2S), a masked-forward alignment method that trains a model under this functional pruning. The masked student must recover a safe refusal through the remaining computation, while a frozen, unmasked teacher supplies complete benign answers to limit capability drift. Across ten model configurations, M2S reduces successful recomputed pruning attacks from 80--279 to 1--44 out of 313 prompts while generally preserving four capability benchmarks. We also evaluate M2S with TwinBreak, which uses a different neuron-selection rule and iterative pruning procedure. Together, these results show that M2S makes targeted pruning less effective by reducing reliance on a small, removable safety-neuron set.

Problem

• LLM safety
• neuron-pruning attacks
• safety alignment
• refusal behavior
• model robustness

Innovation

• neuron-pruning attacks
• masked-forward alignment
• safety alignment
• large language models
• adversarial robustness