PaperlayerView exact revision on arXiv

Source-linked AI summary

NeST: Neuron Selective Tuning for LLM Safety

Sasha Behrouzi, Lichao Wu, Mohamadreza Rostami, Ahmad-Reza Sadeghi

arXiv:2602.16835v2cs.CRcs.LG

TL;DR

Safety alignment must become cheaper and more maintainable without losing robustness, but existing approaches are costly, inconsistent, or intervention-dependent. NeST identifies and clusters safety-relevant feed-forward neurons, trains shared updates on them, and folds those updates into the model. Across 14 open-weight language and multimodal models, it sharply reduces attack success rates with far fewer trainable parameters, while its scope is the hosted black-box threat model.

  • Problem

    Existing safety-alignment methods can be costly to update, audit, and maintain, while lightweight methods may provide inconsistent gains or depend on inference-time intervention.

  • Method

    NeST uses activation probing to identify safety-relevant feed-forward neurons, clusters similar activation profiles, trains shared cluster updates while freezing other neurons, and folds the updates into the original weights.

  • Results

    Across 14 open-weight language and multimodal models, NeST substantially reduces attack success rates and uses orders-of-magnitude fewer trainable parameters than full fine-tuning and lightweight baselines.

  • Takeaways & Limitations

    NeST provides a reusable safety prior for post-hoc hardening that preserves core capabilities and remains effective across jailbreak benchmarks, judges, modalities, and inference settings.

  • Takeaways & Limitations

    NeST targets hosted black-box deployment, while white-box attacks with access to parameters, gradients, or internal representations are outside the evaluated threat model.

Abstract

from arXiv · show

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.

1. Introduction

NeST targets the practical gap between costly, broad safety alignment methods and lightweight approaches by tuning localized safety-relevant neurons without inference-time intervention. It achieves strong robustness across language and multimodal models while using substantially fewer trainable parameters.

  • Motivation: Existing safety-alignment methods remain expensive, structurally blind, or dependent on inference-time intervention, limiting practical maintainability.Practical alignment should be parameter-efficient, structure-aware, and maintainable across downstream model variants.
  • Method: NeST is a neuron-selective framework that aligns updates with localized safety-relevant components without model-wide fine-tuning or inference-time intervention.Its cluster-based mechanism groups safety-relevant neurons by activation similarity, learns shared updates, and freezes the remainder of the model.
  • Results: 44.5% to 1.1%: NeST reduces average attack success rate across 14 open-weight LLMs, with performance comparable to full fine-tuning and better than other evaluated counterparts.The comparison uses a unified evaluation protocol spanning multiple model families and sizes.
  • Results: 55.3% to 1.1%: NeST reduces average attack success rate across four multimodal LLMs spanning text-only, image-only, and reasoning-augmented settings.The evaluation covers diverse inference settings and input modalities.
  • Results: 53.8% to 0.8%: NeST reduces average attack success rate across ten downstream fine-tuned variants by reusing safety-relevant neuron clusters from the base model.This supports post-hoc safety hardening after task-specific adaptation.

2. Preliminaries

LLMs use transformer blocks with self-attention and position-wise feed-forward networks, whose token-dependent dimensions act as neurons. Fine-tuning updates model parameters on task-specific safety data, while LoRA reduces trainable parameters through low-rank updates.

  • Large Language Models: Decoder-only LLMs refine token representations through stacked transformer blocks containing self-attention, feed-forward networks, and residual connections.The feed-forward network is applied independently at each token position.
  • Large Language Models: Modern LLMs use gated feed-forward networks whose higher-dimensional, token-dependent activations define the neurons targeted by NeST.Gate and up projections expand the representation, while the down projection returns it to the model dimension.
  • LLM Fine-Tuning: Fine-tuning optimizes pretrained model parameters on task-specific input-output pairs using a model loss.In safety alignment, the data typically pairs harmful prompts with refusals and benign prompts with helpful answers.
  • LLM Fine-Tuning: Full fine-tuning updates all model parameters, whereas parameter-efficient fine-tuning freezes pretrained weights and trains only restricted or additional parameters.LoRA is presented as an example that augments a frozen linear transformation with a trainable low-rank update.
  • LLM Fine-Tuning: LoRA reduces trainable parameters through low-rank updates, while NeST instead restricts updates to safety-relevant structures.This distinguishes layer- or projection-level low-rank adaptation from structure-selected safety tuning.

3. NeST

NeST targets localized safety-relevant neurons, groups them by activation profiles, and trains shared updates while freezing the rest of the model. The resulting neuron clusters can also be reused to harden downstream fine-tuned variants.

  • 3.1. Threat Model: NeST operates in a hosted black-box deployment setting against prompt-based adversarial inputs, while training occurs offline without inference-time interventions.White-box and training-time adversaries are outside the stated setting.
  • 3.2. Idea and High-Level Design: NeST identifies safety neurons from harmful-versus-benign activation differences, then trains shared cluster updates while freezing non-selected neurons.The workflow uses probing and filtering before cluster-based safety training.
  • 3.3. Safety Neuron Identification: NeST summarizes each neuron by its strongest token response, using coordinatewise max pooling to avoid diluting localized safety signals.The pooled prompt-level activation vector is then used for neuron-level probing.
  • 3.3. Safety Neuron Identification: A linear probe scores neuron-level harmfulness, and standardized positive scores above a threshold define a sparse safety-neuron set.The positivity constraint selects neurons aligned with harmful or refusal-triggering behavior.
  • 3.4. Neuron Clustering: NeST clusters selected neurons by activation-profile similarity, selecting the cluster count with silhouette analysis and tying one update to each cluster.The method balances specialized updates against parameter efficiency and can fall back to a single cluster.
  • 3.5. Cluster-Based Safety Tuning: The same selected neurons and cluster assignments can be transferred from a base model to downstream variants for post-hoc safety hardening.Non-selected neurons remain frozen during this downstream adaptation.

4. Implementation

NeST implements neuron detection with balanced harmful and benign probing data, FFN activation hooks, and linear probes, then clusters selected neurons before selective tuning.

  • 4.1. Safety Neuron Detection: The probing dataset balances harmful prompts from three safety datasets with benign prompts sampled from Natural Reasoning.Harmful prompts are labeled positive and benign prompts negative.
  • 4.1. Safety Neuron Detection: Forward hooks collect activations from gate_proj and up_proj modules across every transformer block.The implementation focuses on feed-forward projections.
  • 4.1. Safety Neuron Detection: A linear probe distinguishes harmful from benign prompts, and standardized probe weights produce a sparse safety-neuron set for each layer.The selection threshold is separately studied as a hyperparameter.
  • 4.2. Safety Neuron Clustering: For each layer-projection pair, k-means clusters only detected safety neurons by activation profile and selects k using silhouette score.If no candidate split exceeds the threshold, all selected neurons form one cluster.
  • 4.3. Safety Tuning: Original parameters are frozen while each selected neuron receives the trainable update associated with its cluster.Cluster-level updates provide parameter sharing during safety tuning.

5. Case Study

The case study examines whether NeST’s activation-based safety clusters are structurally coherent and aligned with optimization behavior. Visualized neurons form compact groups, and within-cluster gradients are more aligned than between-cluster gradients.

  • 5. Case Study: The case study analyzes safety neurons in layer 6 of LLaMA-3.2-1B-Instruct using gate_proj and up_proj activations.Activation profiles are clustered and projected into two dimensions with PCA.
  • 5. Case Study: PCA visualizations show safety neurons forming compact groups in both feed-forward projections.The authors interpret this structure as support for sharing updates among similarly responding neurons.
  • 5. Case Study: The study evaluates whether clustered safety neurons are relevant to optimization, because NeST ties their parameter updates during fine-tuning.Gradient vectors are extracted for weights associated with detected safety neurons.
  • 5. Case Study: Within-cluster neuron pairs have consistently higher gradient cosine similarity than between-cluster pairs, which are often near zero.This indicates that activation-based clusters also reflect aligned update directions during safety tuning.

6. Experimental Results

NeST is evaluated against lightweight and full fine-tuning baselines across text-only, multimodal, and downstream-adapted models. It substantially reduces jailbreak success while using far fewer trainable parameters and largely preserving utility.

  • Experimental setup: NeST is evaluated on 14 open-weight instruction-tuned models spanning text-only and multimodal settings.The evaluation compares NeST with original models, Full FT, LoRA, and other lightweight safety methods.
  • Text-only safety results: 44.5% to 1.1%: NeST reduces average text-model ASR while outperforming lightweight baselines and remaining competitive with full fine-tuning.NeST updates 0.4M parameters on average and achieves lower ASR than full fine-tuning on five of ten models.
  • Parameter efficiency: 0.4M trainable parameters: NeST uses substantially fewer parameters than Full FT, LoRA, Circuit Breaker, and SN-Tune while attaining comparable ASR.The reported reductions are approximately 17,700× versus Full FT and 9.4× versus LoRA.
  • Multimodal robustness: 55.3% to 1.1%: NeST lowers average ASR across four multimodal models and remains effective across text, image, and reasoning-augmented settings.Base models reach up to 85.0% ASR in some image-based configurations.
  • Downstream safety restoration: 53.8% to 0.8%: post-hoc NeST hardening restores safety across ten downstream fine-tuned variants.After hardening, ASR is at most 1.6% for every evaluated downstream model.
  • Utility preservation: NeST preserves reasoning utility with minor average changes across GSM8K, ARC, and MMLU after fine-tuning.Average accuracy changes from 61.2% to 60.3% on GSM8K, 74.0% to 69.1% on ARC, and 60.9% to 57.2% on MMLU.

7. Ablation and Hyperparameter Study

Additional evaluations show that NeST generalizes across jailbreak benchmarks and judges, preserves benign-task behavior, and depends on selective neuron detection and balanced clustering choices.

  • Stability across jailbreaks and judges: NeST’s evaluation trends remain consistent on JailBreakV-28k and under the Llama-Guard-3-8B judge.The relative ordering of methods remains stable despite generally lower ASR values under Llama-Guard-3-8B.
  • Benign utility: 0.9% average ORR: NeST maintains low refusal rates on benign prompts across evaluated models.The result supports preservation of helpful behavior rather than broad output suppression.
  • Safety-neuron stability: 76.8% seed Jaccard and 90.2% dataset Jaccard indicate that selected safety neurons remain substantially stable across probe seeds and harmful datasets.Lower seed overlap mainly reflects near-threshold or redundant neurons, while dataset overlap exceeds 90% for most models.
  • Neuron-selection threshold: 0.54% average ASR at zthr = 2 versus 3.7% at zthr = 4 shows that overly restrictive neuron filtering degrades safety performance.The default zthr = 3 yields 1.1% average ASR and remains effective under moderate threshold variation.
  • Clustering strength: Strong clustering lowers average ASR to 0.3% but raises trainable parameters from 0.4M to 5.3M, exposing an efficiency–flexibility trade-off.Weak clustering costs 0.3M parameters and yields 1.5% ASR, while the default setting yields 1.1% ASR.

8. Discussion

The discussion positions NeST as a black-box-deployment defense with limits against privileged white-box attackers and potential compatibility with broader alignment objectives.

  • Threat model: NeST targets hosted black-box deployment, where adversaries can query the model but cannot access parameters, gradients, or activations.White-box attacks requiring internal inspection or manipulation fall outside the considered deployment setting.
  • Adaptive attacks: Under the black-box threat model, LLM-Fuzzer achieves 0.78% ASR before NeST and 0.03% after NeST on Phi-4 over 20,000 queries.Gradient-based GCG attacks are not directly applicable because they require model gradients.
  • Structure-aware adaptation: NeST aligns updates with neuron-level internal structure rather than applying generic low-rank constraints without regard to semantic roles.The framework is presented as task-structured parameter-efficient adaptation.
  • Broader alignment compatibility: The neuron detection, clustering, and selective adaptation mechanisms are described as compatible in principle with RLHF, GRPO, and other preference-based optimization methods.In those settings, updates would be restricted to clustered safety-relevant neurons.

9. Related Work

Prior work either improves parameter efficiency without modeling safety-specific structure or uses safety-related internal structure without a compact cluster-based update mechanism. NeST combines localized neuron selection with functionally coherent shared updates.

  • Prior approaches: Existing safety alignment approaches pursue either output-level behavior or internal safety representations, each addressing only part of practical adaptation needs.The related work frames practical safety adaptation as requiring both efficiency and reliability.
  • Neuron-level alignment: Neuron-level methods can improve refusal behavior while leaving most model parameters unchanged, but commonly identify neurons through harmful-query importance measures.SN-Tune is cited as an example of selective neuron modification.
  • NeST’s positioning: NeST combines safety-neuron localization, functionally coherent clustering, and shared updates restricted to selected clusters.This addresses the stated gap between generic parameter-efficient adaptation and structure-aware safety methods.

10. Conclusions

NeST is a lightweight, structure-aware framework that hardens LLM safety by selectively updating clustered safety-relevant neurons while freezing the rest of the model. Across 14 open-weight language and multimodal models, it reduces attack success rates with fewer trainable parameters and no additional inference-time cost.

  • 10. Conclusions: NeST identifies safety-relevant neurons through contrastive activation probing, groups them into functionally coherent clusters, and trains shared cluster-level updates.The rest of the model remains frozen.
  • 10. Conclusions: The learned updates are folded into the original weights, producing a standard hardened model with no additional inference-time cost.
  • 10. Conclusions: Across 14 open-weight language and multimodal models, NeST substantially reduces attack success rates while using fewer trainable parameters than full fine-tuning and lightweight baselines.It also preserves core reasoning and knowledge capabilities, avoids excessive over-refusal, and remains effective across jailbreak benchmarks, judges, modalities, and inference settings.

Appendix A. Ethical Considerations

The study frames NeST as a post-hoc safety-alignment method intended to reduce unsafe generations while preserving useful capabilities. Its evaluation uses public data, open-weight models, and controlled offline protocols, with defensive safeguards and responsible-research practices.

  • Appendix A. Ethical Considerations: NeST aims to reduce unsafe generations, improve robustness to harmful prompts, and support safer deployment of open-weight and multimodal LLM systems.
  • Appendix A. Ethical Considerations: The method is intended for developers, deployers, organizations, safety researchers, and end users affected by unsafe model behavior.The stated goal is to improve reliability while preserving useful model capabilities.
  • Appendix A. Ethical Considerations: All experiments use publicly available datasets, open-weight pretrained models, and controlled offline evaluation without private, proprietary, or personally identifiable data.No real-world deployment was involved.
  • Appendix A. Ethical Considerations: NeST is described as strengthening intrinsic safety mechanisms while maintaining utility and enabling post-hoc hardening after downstream fine-tuning.
  • Appendix A. Ethical Considerations: The framework operates through offline training and evaluation and does not introduce mechanisms for generating harmful content or weakening safeguards.Released artifacts focus on defensive safety alignment and responsible research use.
  • Appendix A. Ethical Considerations: The study reports controlled benchmark evaluation and follows standard responsible research practices for safety-related data handling and contributor well-being.
Loading 2602.16835v2…