NeST: Neuron Selective Tuning for LLM Safety
NeST uses activation probing and cluster-level training for efficient post-hoc safety alignment, reducing jailbreak success from 44.5% to 1.1%.
Key Findings
Methodology
NeST identifies safety-related neurons via activation probing on harmful and benign prompts, then groups neurons with similar activation profiles using clustering algorithms like k-means. It trains shared updates for each cluster while freezing the rest of the model. The training is performed solely on vanilla malicious prompts, avoiding jailbreak-specific data, which enhances robustness against diverse attacks. Post-training, updates are folded into the original weights, resulting in no inference overhead. Extensive experiments across 14 open-weight models demonstrate that NeST reduces attack success rate from 44.5% to 1.1%, with only 0.4 million trainable parameters, outperforming lightweight baselines and approaching full fine-tuning robustness.
Key Results
- NeST reduces average jailbreak attack success rate from 44.5% to 1.1% on text-only models, training only 0.4M parameters. In multimodal models, ASR drops from 55.3% to 1.1%. For downstream fine-tuned variants, safety is restored from 53.8% to 0.8%. These results show that targeted, structure-aware updates can dramatically improve safety while maintaining core capabilities.
- By focusing on localized safety-related neurons, NeST avoids the high computational costs of full fine-tuning and inference-time interventions. Its structure-aware approach ensures that safety improvements are coherent and stable, enabling efficient transfer to downstream tasks.
- Compared to RLHF and LoRA, NeST offers superior safety robustness and parameter efficiency, validating the effectiveness of linking safety behavior to internal neuron clusters. This approach opens new avenues for scalable, maintainable safety in large models.
Significance
This work addresses a critical bottleneck in deploying safe LLMs—balancing safety, efficiency, and maintainability. By localizing safety mechanisms to specific neurons and clusters, NeST provides a scalable solution that can be applied post hoc, reducing costs and complexity. Its structure-aware design enhances robustness against evolving attack strategies, making it highly relevant for real-world deployment. The approach also offers insights into the internal mechanisms of safety behavior, contributing to the broader understanding of model interpretability and control. Overall, NeST paves the way for safer, more reliable large language models, with significant implications for industry and research.
Technical Contribution
NeST introduces a novel framework combining activation probing, neuron clustering, and cluster-level parameter sharing to localize safety-related internal components. It leverages a simple yet effective linear probe to identify neurons associated with harmful prompts, then applies k-means clustering to group neurons with similar activation profiles. During training, shared updates are learned for each cluster, significantly reducing the number of trainable parameters. The updates are subsequently folded into the original model weights, enabling safe deployment without inference overhead. This structure-aware, parameter-efficient approach enhances robustness and transferability, and supports downstream safety hardening. Its core innovation lies in translating the internal safety mechanisms into localized, trainable neuron clusters, enabling targeted, scalable safety adaptation.
Novelty
This is the first work to combine activation probing with clustering-based parameter sharing for post-hoc safety alignment in large models. Unlike LoRA, which applies low-rank updates broadly, NeST focuses on safety-relevant neurons, achieving a high degree of structural specificity. Its approach of localizing safety behavior to neuron clusters and folding updates into the original weights offers a new paradigm for efficient, maintainable safety tuning. This method also demonstrates strong generalization to diverse jailbreaks without attack-specific data, setting a new standard for robustness and transferability in safety alignment.
Limitations
- The effectiveness of NeST depends on accurately identifying safety neurons; if the probing fails to capture all relevant neurons, safety may be compromised. Its performance may vary with different probing datasets and thresholds.
- The clustering process introduces hyperparameters (number of clusters, silhouette thresholds) that require tuning, potentially affecting robustness across models and scenarios. Large models may still incur non-trivial computational costs during training.
- Current evaluations are primarily on static models; adapting NeST to dynamic, continually learning systems remains an open challenge. Further research is needed to extend its applicability to real-time safety updates.
Future Work
Future directions include developing adaptive clustering algorithms that automatically determine optimal cluster numbers, integrating dynamic activation analysis for real-time safety monitoring, and extending the framework to multi-modal and multi-task settings. Exploring federated or distributed training paradigms could enhance scalability. Additionally, combining NeST with other interpretability tools may improve understanding of safety mechanisms, enabling more precise control and verification.
AI Executive Summary
Large Language Models (LLMs) have revolutionized AI applications, enabling tasks from text generation to reasoning. However, their deployment faces significant safety challenges, notably vulnerability to jailbreak attacks that elicit harmful outputs. Traditional safety methods like reinforcement learning from human feedback (RLHF) and full model fine-tuning are effective but costly, difficult to maintain, and often degrade model capabilities. These limitations hinder rapid updates, especially as models evolve or are adapted to new domains.
In response, NeST (Neuron-Selective Tuning) introduces a novel approach that localizes safety mechanisms within specific internal neurons. The framework begins with activation probing, where it identifies neurons in the model’s feed-forward layers that are highly responsive to harmful prompts. Using simple linear probes, NeST scores neurons based on their discriminative power, then groups neurons with similar activation profiles through clustering algorithms such as k-means. During training, only these neuron clusters are updated with shared parameters, while the rest of the model remains frozen. This targeted, structure-aware training significantly reduces the number of trainable parameters—just 0.4 million on average—yet achieves robustness comparable to full fine-tuning.
Experimental results across 14 open-weight models demonstrate that NeST reduces jailbreak success rates from 44.5% to 1.1%, outperforming lightweight baselines like LoRA. The learned updates are folded into the original weights, incurring no inference overhead and enabling easy transfer to downstream models. This approach not only enhances safety but also preserves core reasoning and knowledge capabilities, making it highly practical for real-world deployment.
The core innovation of NeST lies in its ability to connect internal neural structures with safety behavior, enabling efficient, maintainable, and robust safety alignment. Its structure-aware design offers a promising pathway for scalable safety solutions in increasingly complex AI systems. Future work aims to refine neuron identification, extend to dynamic environments, and explore multi-modal applications, further strengthening the safety and reliability of large models in diverse settings.
Deep Dive
Abstract
Safety alignment is essential for the responsible deployment of Large Language Models (LLMs). Yet, existing approaches often rely on heavyweight fine-tuning that is costly to update, audit, and maintain across model families. Full fine-tuning incurs substantial computational and storage overhead, while parameter-efficient methods, e.g., Low-Rank Adaptation (LoRA), trade efficiency for inconsistent safety gains and sensitivity to design choices. Safety intervention mechanisms reduce unsafe outputs without modifying model weights, but do not directly shape or preserve the internal representations that govern safety behavior. We present NeST, a Neuron-Selective Tuning framework for efficient post-hoc safety alignment. NeST identifies safety-relevant feed-forward neurons via activation probing on vanilla harmful and benign prompts, clusters neurons with similar activation profiles, and trains shared cluster-level updates while freezing the rest of the model. Importantly, NeST is trained only on vanilla malicious prompts, without using jailbreak-specific attack data, yet generalizes robustly to diverse jailbreaks. The learned updates are then folded into the original weights, incurring no inference-time overhead. Evaluated on 14 open-weight language and multimodal models, NeST outperforms lightweight baselines and approaches full fine-tuning robustness with significantly fewer trainable parameters. On text-only models, NeST reduces average jailbreak attack success rate from 44.5% to 1.1% while training only 0.4M parameters on average. Across multimodal settings, it reduces ASR from 55.3% to 1.1%, and for downstream fine-tuned variants, it restores safety by reducing ASR from 53.8% to 0.8%. These results show that robust, maintainable safety alignment can be achieved by concentrating adaptation on localized, functionally coherent safety structures.