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

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.

💡 Research Summary

The paper introduces Neuron‑Guided Safe Decoding (NGSD), a lightweight inference‑time safety alignment method for large language models (LLMs) that requires training only a tiny “expert” model and uses a single neuron as a dynamic gating mechanism. Existing safety alignment techniques largely rely on post‑training fine‑tuning (e.g., RLHF, DPO) or inference‑time interventions that either need heavy pre‑computation or apply uniform safety constraints regardless of the model’s intrinsic risk awareness. NGSD bridges this gap by explicitly coupling the model’s internal safety signals with an external expert, while keeping computational overhead minimal and ensuring strong cross‑model generalization.

The approach works as follows. A small expert model (M_e) is fine‑tuned on a safety‑augmented dataset using a low‑cost method such as LoRA. This expert shares the tokenizer and output space with the target base model (M_b), enabling direct transfer across all larger models in the same family. During decoding, the probability distributions over the next token from both models, (p_b) and (p_e), are computed. Their (\ell_1) distance is halved to produce a scalar risk signal (I_t = \frac{1}{2}|p_b - p_e|1). Rather than reacting to each instantaneous discrepancy, NGSD feeds (I_t) into a biologically‑inspired neuron whose membrane potential (V(t)) integrates the signal over time according to (\tau_m \frac{dV}{dt}=-(V-V{rest})+RI(t)). When (V(t)) exceeds a threshold (v_{th}), the neuron spikes, opening a gate that triggers a safety intervention for the current step. The intervention follows the classic SafeDecoding formulation: a candidate token set (C) (union of top‑K tokens from both models) is constructed, and the logits are corrected as (\tilde p(y)=p_b(y)+\alpha(p_e(y)-p_b(y))) for (y\in C). The scalar (\alpha\in