Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Light Alignment Improves LLM Safety via Model Self-Reflection with a Single Neuron

Sicheng Shen, Mingyang Lv, Han Shen, Jialin Wu, Binghao Wang, Zhou Yang, Guobin Shen, Dongcheng Zhao, Feifei Zhao, Yi Zeng

TL;DR

This work tackles the safety-utility trade-off in large language models by introducing NGSD, a lightweight decoding-time alignment method that couples intrinsic risk awareness with external guidance using a single nonlinear neuron as a gating mechanism. Safety signals are derived from a distributional discrepancy between a base model and a small safety expert, and interventions are applied only when the neuron gate fires, via a logit-correction step on a constrained token set. A prompt-level self-reflection stage fixes the safety strength before decoding, enabling adaptive yet low-overhead safety control that transfers across model families within the same lineage. Empirical results across multiple model backbones show strong safety robustness, favorable utility retention, and improved efficiency compared to prior decoding-time defenses, supporting practical deployment of safer LLMs with limited retraining.

Abstract

The safety of large language models (LLMs) has increasingly emerged as a fundamental aspect of their development. Existing safety alignment for LLMs is predominantly achieved through post-training methods, which are computationally expensive and often fail to generalize well across different models. A small number of lightweight alignment approaches either rely heavily on prior-computed safety injections or depend excessively on the model's own capabilities, resulting in limited generalization and degraded efficiency and usability during generation. In this work, we propose a safety-aware decoding method that requires only low-cost training of an expert model and employs a single neuron as a gating mechanism. By effectively balancing the model's intrinsic capabilities with external guidance, our approach simultaneously preserves utility and enhances output safety. It demonstrates clear advantages in training overhead and generalization across model scales, offering a new perspective on lightweight alignment for the safe and practical deployment of large language models. Code: https://github.com/Beijing-AISI/NGSD.

Light Alignment Improves LLM Safety via Model Self-Reflection with a Single Neuron

TL;DR

This work tackles the safety-utility trade-off in large language models by introducing NGSD, a lightweight decoding-time alignment method that couples intrinsic risk awareness with external guidance using a single nonlinear neuron as a gating mechanism. Safety signals are derived from a distributional discrepancy between a base model and a small safety expert, and interventions are applied only when the neuron gate fires, via a logit-correction step on a constrained token set. A prompt-level self-reflection stage fixes the safety strength before decoding, enabling adaptive yet low-overhead safety control that transfers across model families within the same lineage. Empirical results across multiple model backbones show strong safety robustness, favorable utility retention, and improved efficiency compared to prior decoding-time defenses, supporting practical deployment of safer LLMs with limited retraining.

Abstract

The safety of large language models (LLMs) has increasingly emerged as a fundamental aspect of their development. Existing safety alignment for LLMs is predominantly achieved through post-training methods, which are computationally expensive and often fail to generalize well across different models. A small number of lightweight alignment approaches either rely heavily on prior-computed safety injections or depend excessively on the model's own capabilities, resulting in limited generalization and degraded efficiency and usability during generation. In this work, we propose a safety-aware decoding method that requires only low-cost training of an expert model and employs a single neuron as a gating mechanism. By effectively balancing the model's intrinsic capabilities with external guidance, our approach simultaneously preserves utility and enhances output safety. It demonstrates clear advantages in training overhead and generalization across model scales, offering a new perspective on lightweight alignment for the safe and practical deployment of large language models. Code: https://github.com/Beijing-AISI/NGSD.
Paper Structure (43 sections, 10 equations, 3 figures, 9 tables, 2 algorithms)

This paper contains 43 sections, 10 equations, 3 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Jailbreak Attack
  4. Decoding-time Safety Methods
  5. Leveraging Intrinsic Model Capabilities for Safety
  6. Preliminary
  7. Safety-Aware Decoding
  8. Nonlinear Neuron as Temporal Gating Mechanism
  9. Methodology
  10. Overview
  11. Neuron-Guided Decoding
  12. Distributional Discrepancy as Risk Signal.
  13. Neuron-Inspired Gating.
  14. Selective SafeDecoding Intervention.
  15. Prompt-Level Self-Reflection
  16. ...and 28 more sections

Figures (3)

  • Figure 1: The overview of Neuron Guided Safe Decoding (NGSD) pipeline.
  • Figure 2: Trade-off between safety and utility under different fixed values of $\alpha$.
  • Figure 3: Neuron is robust among gating mechanism against attack methods.