🤖 AI

No Single Neuron of Failure: Distributed Safety Alignment Against White-Box Attacks

This paper proposes Distributed Safety Alignment (DSA), a method that enhances large language model robustness against white-box neuron-level attacks by redundantly encoding safety capabilities across multiple neurons through targeted disruption and gradient-guided identification, thereby eliminating single points of failure while preserving general utility.

Original authors: Simiao Xie, Chuancheng Shi, Shangze Li, Wenhua Wu, Fei Shen, Ying Zhou, Zhiyong Wang, Tat-Seng Chua

Published 2026-08-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Simiao Xie, Chuancheng Shi, Shangze Li, Wenhua Wu, Fei Shen, Ying Zhou, Zhiyong Wang, Tat-Seng Chua

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a giant, super-smart robot to be a helpful assistant. You want it to be brilliant at solving math problems, writing stories, and chatting about your day, but you also need to make sure it never agrees to help with something dangerous, like hacking a bank or writing a mean letter. This is the world of "AI safety," where researchers try to build digital guardrails into these massive computer brains.

To understand the problem this paper tackles, think of how a human brain works. When you decide to say "no" to a bad idea, it's not just one tiny thought in your head that does the work; it's a whole team of neurons firing together. But in the AI models we use today, researchers have discovered something scary: the robot's "no" is often handled by just a few, very specific neurons. It's like if your entire safety team relied on a single, tiny security guard standing at the front door. If a clever hacker (a "white-box attacker") could sneak in and knock out that one guard, the whole building would be wide open to danger. The robot would suddenly forget how to say "no" and might start helping with terrible requests. This paper asks a simple but crucial question: What if we didn't rely on just one guard, but instead trained a whole army of them to share the job?

The researchers behind this study, Simiao Xie and their team, propose a new way to train these AI models called Distributed Safety Alignment (DSA). They argue that the current way of making AI safe is too fragile because it concentrates all the "refusal" power into a small, fixed group of neurons. If an attacker finds and removes those specific neurons, the safety system collapses. To fix this, DSA treats safety like a game of "musical chairs" for the robot's brain cells.

Here is how their method works, using a playful analogy: Imagine the AI's brain is a giant orchestra. In the old way, only the first violinist was allowed to play the "Stop!" note. If that violinist got sick or was kicked out, the music would turn into a disaster. The new DSA method forces the conductor (the training algorithm) to make the entire orchestra learn that "Stop!" note.

To do this, the researchers use a clever trick during the training process. They first identify which neurons are currently doing the heavy lifting for saying "no" to bad requests. Then, they play a game of "hide and seek" with the model. They deliberately turn off (mask) those top neurons, pretending they are gone. But here's the twist: they also randomly turn off other neurons at the same time. This forces the AI to panic a little bit and realize, "Oh no, my main 'Stop' neurons are gone, and some backup ones are missing too! I need to find new neurons to help me say 'no'!"

By repeating this process, the AI learns to spread the "safety" job across hundreds of different neurons instead of just a few. It becomes like a team where everyone knows the emergency code. Even if an attacker manages to knock out the original "safety neurons," the model still has dozens of other neurons ready to step in and refuse the harmful request.

The paper shows that this approach works incredibly well. When they tested their method on several different AI models (including Qwen2.5 and LLaMA), they found that even when attackers tried to surgically remove the most important safety neurons, the DSA-trained models still refused to do bad things. In fact, in some tests, the old models failed almost every time the safety neurons were removed, while the DSA models kept saying "no" almost perfectly.

Crucially, the researchers also checked that this new safety system didn't make the AI dumber. They tested the models on math problems, reading comprehension, and general knowledge, and found that the models were just as smart and helpful as before. The safety didn't come at the cost of intelligence; the models just became much harder to trick.

The paper suggests that this "distributed" approach is a major step forward because it stops relying on a "single point of failure." Instead of building a wall with one weak spot, they are building a safety net that is woven from many threads. If one thread breaks, the net still holds. While the paper notes that this is a simulation-based improvement and that real-world attacks are always evolving, the results strongly suggest that spreading the safety load across the AI's brain is a much more robust way to keep these powerful tools safe for everyone.

Technical Summary: Distributed Safety Alignment Against White-Box Attacks

1. Problem Statement

With the rapid deployment of open-weight large foundation models (LFMs), safety threats have evolved from black-box jailbreaks to fine-grained white-box attacks. These attacks directly identify and prune critical internal "safety neurons" or pathways, causing the model to bypass refusal mechanisms and output harmful content.

Existing defense strategies generally fall into two categories: unit-based (preserving specific safety neurons) and pathway-based (protecting cross-layer chains). However, the paper argues that both approaches suffer from a fundamental structural flaw: they rely on a "static fortification" paradigm. By concentrating safety capabilities on a limited, fixed set of neurons or pathways, these methods create a fragile "single point of failure." When an adversary precisely locates and prunes these specific hubs, the entire defense collapses, leading to immediate safety failures. The core challenge is to maintain safety even when core computational units are compromised, requiring a shift from centralized defense to distributed redundancy.

2. Methodology: Distributed Safety Alignment (DSA)

The authors propose Distributed Safety Alignment (DSA), a framework designed to redundantly encode safety capabilities across multiple computational neurons, ensuring the model maintains a safety baseline even when critical neurons are disrupted.

A. Localization and Attribution

DSA operates on the intermediate activations immediately preceding the down-projection layers in language-side feed-forward networks (FFNs). Instead of relying on pre-defined lists of safety neurons, DSA dynamically identifies dominant refusal neurons using a direction-aware first-order Taylor score:

Refusal Loss Calculation: The model computes a response-only refusal loss ( $L_{ref}$ ) on harmful queries paired with safe refusal targets.
Taylor Expansion: To estimate a neuron's contribution, the method simulates suppressing a specific neuron's activation ( $\Delta a = -a$ ) and calculates the resulting change in loss using a first-order Taylor expansion: $\Delta L \approx -a \cdot \frac{\partial L}{\partial a}$ .
Direction-Aware Scoring: A score $S_i$ is computed by retaining only the positive suppression effects (neurons whose removal increases the refusal loss). This ensures the method targets neurons that actively support refusal behavior, rather than just those with high magnitude.
Selection: The top- $K$ neurons with the highest scores are identified as the core refusal-neuron set ( $K_{core}$ ).

B. Adversarial Redundant Safety Alignment

Once the core neurons are identified, DSA employs a two-pronged perturbation strategy during training to force the model to abandon reliance on these specific units and distribute safety across compensatory neurons:

Deterministic Targeted Masking: The identified top- $K$ neurons are deterministically masked (set to zero) throughout the realignment process. This simulates a targeted attack, forcing the model to find alternative pathways.
Stochastic Dropout: To prevent the model from simply re-concentrating safety onto a new, minimal subset of neurons, the remaining unmasked neurons are subjected to stochastic dropout. This introduces diverse local failures, compelling the model to encode refusal behavior redundantly across a broader set of compensatory neurons.

C. Optimization Objective

The training objective involves joint optimization on two datasets under the combined masking and dropout regime:

Harmful-Refusal Loss ( $L_{safe}$ ): Ensures the model learns to refuse harmful instructions despite the disruption of its original safety neurons.
Generality-Preserving Loss ( $L_{gen}$ ): Ensures the model maintains its ability to follow benign instructions and perform general tasks.

The final objective is: $\theta^*_{ff} = \arg \min_{\theta_{ff}} (L_{safe} + \lambda_{gen} L_{gen})$ .
Crucially, only the language-side feed-forward parameters are trainable; the backbone remains frozen, and all interventions are removed after training, resulting in zero inference-time overhead.

3. Key Contributions

DSA Framework: A novel method that utilizes direction-aware Taylor scoring and dynamic perturbation to redundantly encode refusal mechanisms across a broader set of internal neurons, breaking the reliance on single points of failure.
Adversarial Redundant Safety Alignment: A training strategy that combines deterministic masking of high-scoring refusal neurons with stochastic dropout of the remainder, forcing the emergence of compensatory safety pathways.
Robustness and Utility Preservation: Extensive experiments demonstrate that DSA significantly enhances robustness against fine-grained white-box safety suppression attacks while preserving the model's native general language and multimodal utility.

4. Experimental Results

The authors evaluated DSA on diverse Large Language Models (Qwen2.5, LLaMA-3.2, Gemma, Phi-4) and Multimodal Large Language Models (Qwen2.5-VL, LLaVA-1.5).

White-Box Attack Robustness: Under the most aggressive FULL pruning attack (where the top safety neurons are removed), DSA drastically reduced the Attack Success Rate (ASR).
- On Qwen2.5-7B, ASR dropped from 279/313 (Original) to 9/313 (DSA).
- On LLaMA-3.2-3B, ASR dropped from 176/313 (Original) to 0/313 (DSA).
- DSA consistently outperformed baselines like SafeNeuron, RLHF-Safety, and SN-Tune across all model sizes and attack settings (ORI, ES, SAS, FULL).
General Capability: DSA maintained or slightly improved performance on standard benchmarks (ARC, GSM8K, TruthfulQA) compared to the original backbones, demonstrating that the safety alignment did not incur an "alignment tax" on utility.
Multimodal Safety: DSA extended its robustness to multimodal models, effectively resisting attacks where harmful intent was embedded in images or visual text.
Unseen Attacks: The model remained robust against unseen white-box attacks (GRAD, WANDA, ABLATE) that were not part of the training criteria, suggesting the learned redundancy is generalizable.
Post-Attack Utility: Even after applying the FULL pruning attack to the trained model, general capabilities (e.g., ARC scores) and benign generation validity remained largely unchanged, confirming that the low ASR was due to robust safety, not model collapse.

5. Significance and Claims

The paper claims that DSA represents a paradigm shift in AI safety alignment. By moving away from the fragile "static fortification" of specific neurons toward a distributed safety model, DSA addresses the critical vulnerability of single points of failure in internal safety mechanisms.

The authors emphasize that DSA achieves this robustness without introducing inference-time overhead or degrading general model utility. The method forces the model to structurally encode safety redundantly, ensuring that even if an adversary successfully prunes the most obvious safety neurons, the model retains the ability to refuse harmful requests through compensatory pathways. This offers a scalable and robust solution for maintaining safety in the era of open-weight foundation models.