1 Introduction↩︎

Large language models (LLMs) [1], [2] have demonstrated remarkable capabilities across a wide range of tasks, ranging from question answering to complex instruction following. With their growing deployment in real-world settings, concerns about their reliability, safety, and robustness have become a major concern [3]. To address these risks, developers introduce an additional safety alignment stage after pretraining, often through supervised fine-tuning or preference optimization. These methods mainly differ in alignment data, training objective, and model-level safety representations but are commonly optimized through gradient-based updates [4]–[7]. This stage makes relatively small but targeted weight adjustments to discourage the generation of harmful, unethical, or dangerous content while preserving the model’s general utility and instruction-following capability [8], [9].

Despite recent progress in developing effective and efficient LLM safety alignment algorithms, achieving robust safety alignment remains a significant challenge. In particular, safety performance can deteriorate rapidly under post-alignment perturbations to model weights or intermediate activations. Recent studies show that even small and generic disturbances can substantially weaken safety alignment: injecting Gaussian noise into model weights can cause otherwise aligned models to exhibit unsafe or misaligned behavior [10], while perturbations to specific intermediate activation layers can reliably weaken safety guardrails [11], [12]. Beyond explicit noise injection, post-training compression techniques such as quantization may introduce approximation errors that weaken safety alignment [13]. These perturbation-induced failures suggest that safety alignment may be encoded by a sharp or overfitted safety boundary in parameter or activation space, where small deviations can move the model from refusal to unsafe compliance. Robust alignment therefore requires not only stronger safety signals, but also smoother local safety landscapes.

Existing research on robust safety alignment has mainly focused on three directions: improving data-level robustness, redesigning objective-level alignment signals, and strengthening model-level safety representations. Data-level approaches improve robustness by curating or reweighting fine-tuning data [14], [15]; objective-level approaches modify alignment losses to encourage more persistent refusal behavior [16], [17]; and model-level approaches stabilize safety representations or protect safety-critical layers [18]–[20]. While effective, these methods typically rely on problem-specific modifications to the training data, alignment objective, or model internals. In contrast, the role of the base optimizer itself, independent of such problem-wise and algorithm-level modifications, remains largely unexplored in shaping alignment robustness. The perturbation-induced failures discussed above suggest that robust alignment is also closely tied to the local geometry of the learned safety boundary, which can be shaped by the base optimizer itself.

This motivates us to revisit zeroth-order (ZO) optimization as an optimizer-level mechanism for robust safety alignment. Unlike standard gradient-based updates, ZO estimates update directions by evaluating losses at randomly perturbed parameters, which can be viewed as optimizing a locally smoothed objective [21]. While recent work has mainly explored ZO methods for memory-efficient LLM fine-tuning [22]–[24] and LLM unlearning [25], [26], its potential for improving alignment robustness remains underexplored.

In this paper, we strive to improve the robustness of safety alignment through the lens of zeroth-order methods, which have been largely underexplored in this context. Zeroth-order optimization offers a promising route to alignment robustness as it evaluates how safety alignment behaves under perturbations. Our study consists of three main components. We empirically show that safety-aligned models remain fragile under simple perturbations: lightweight parameter- or activation-level noise can substantially weaken safety behavior. While ZO optimization can enhance alignment robustness, directly replacing first-order (FO) alignment with ZO optimization is ineffective due to its slower and less stable convergence. Therefore, we propose a hybrid refinement strategy that applies ZO updates after standard FO safety alignment. We show theoretically and empirically that a small amount of post-alignment ZO refinement can improve safety robustness with only modest computational overhead. Finally, we exploit the perturbation-based evaluations inherent to ZO optimization to estimate layer-wise robustness sensitivity, enabling a targeted refinement strategy that concentrates updates on robustness-critical layers, improving the efficiency and effectiveness of robustness-oriented safety alignment. Our contributions can be summarized as follows:

• We are the first to comprehensively study the robustness of safety alignment against parameter and activation perturbations from the perspective of zeroth-order optimization.

• We propose an FO–ZO hybrid optimization framework that integrates FO and ZO optimizers, combining ZO-induced robustness with FO-driven alignment effectiveness.

• We show theoretically and empirically that a small amount of post-alignment ZO refinement can improve safety robustness with only modest computational overhead.

2 Background and Related Work↩︎

2.1 Safety Alignment in Large Language Models↩︎

Safety alignment aims to ensure that large language models (LLMs) generate outputs that are consistent with human values and safety constraints [17], [27], [28]. A standard paradigm is to post-train pretrained models using human feedback, such as Reinforcement Learning from Human Feedback (RLHF) [4]. Subsequent work further improves or simplifies this paradigm from different perspectives. Constitutional AI reduces the dependence on direct human feedback by using a set of human-written principles and AI feedback to improve harmlessness [29]; Safe RLHF explicitly separates helpfulness and harmlessness preferences and formulates safety alignment as a constrained optimization problem [30]; Direct Preference Optimization (DPO) directly optimizes language models from preference data without explicitly training a reward model or running reinforcement learning [5]; and recent reasoning-based alignment methods improve safety by training models to explicitly reason about safety requirements before generating final answers [31].

2.2 Robustness of Safety Alignment↩︎

Prior work has identified several ways to compromise the safety behavior of aligned models. At the input level, jailbreak and adversarial prompting attacks can bypass safety guardrails by rewriting, obfuscating, or optimizing harmful instructions [32]–[34]. At the model level, safety behavior can be weakened by malicious or even benign fine-tuning [35], [36], persistent backdoors [37], or interventions on refusal-related representations [38]. These results suggest that safety alignment may rely on mechanisms that are effective under normal conditions, yet remain fragile under adversarial inputs, distribution shifts, or post-training modifications.

While many of these failures involve stronger attack capabilities or carefully engineered procedures, safety alignment can also be vulnerable to much simpler changes. Even generic perturbations, arising from quantization, parameter noise, activation noise, hardware faults, or other post-alignment modifications, may be sufficient to degrade the model’s safety behavior. Quantization reduces numerical precision to improve efficiency [39], but the induced approximation errors may alter model behavior; recent work further shows that quantization-related fault injection can weaken safety alignment [13]. Similarly, weight or activation perturbations can disturb internal representations and expose safety failures that are not visible under standard inference [11], [12], [40].

2.3 Zeroth-Order Optimization↩︎

Zeroth-order (ZO) optimization decreases memory consumption by constructing gradient estimates by evaluating the objective at perturbed points in the parameter space, without requiring backpropagation [21]. Beyond its computational advantage, it offers a promising route to enhancing robustness by evaluating how the model behaves under perturbations. This intuition is further supported by a curvature-based connection: ZO optimization has been shown to favor flatter minima [41], which contribute to improved local robustness against small perturbations [42]. This perspective has motivated prior work on using ZO optimization to train and adapt large models under constrained or noisy conditions [24], [43]. More recently, zeroth-order methods have also been explored in alignment-related settings. Galatolo et al. [44] explore zeroth-order methods for LLM preference optimization, aiming to reduce memory usage and iteration time. In this work, we instead investigate zeroth-order optimization from the perspective of safety alignment robustness and show that it can enhance the resilience of LLM safety alignment against parameter and activation tampering.

3 Preliminaries↩︎

Zeroth-order gradient estimation approximates gradients from function evaluations at perturbed parameter points, without requiring explicit gradient information.

Lipschitz continuity and smoothness control the variation of the objective and its gradient, serving as standard assumptions in FO and ZO convergence analyses [22], [46]. We also adopt the Polyak–Łojasiewicz (PL) condition, a common gradient-dominance assumption for deriving linear convergence-type guarantees in nonconvex settings [47].

4 Methodology↩︎

In this section, we present our framework for improving the robustness of safety alignment, as illustrated in Figure 1. Motivated by the fragility of aligned models under simple perturbations, we use the perturbation-based evaluations inherent to ZO optimization to identify robustness-critical layers and then concentrate ZO refinement on these layers.

4.1 Zeroth-Order Robustness Refinement↩︎

We introduce ZO robustness refinement as a post-alignment stage, which improves the robustness of an already FO-aligned model rather than replacing FO safety alignment. While FO efficiently optimizes the alignment objective, ZO is then used to enhance stability under perturbations.

Algorithmic framework. Algorithm 2 summarizes our three-stage pipeline. We first perform FO safety alignment from \(\theta_0\) to obtain an aligned model \(\theta^{\mathrm{fo}}\). We then compute layer-wise robustness sensitivity scores \(S(\ell)\) and select the top-\(m\) layers as \(\mathcal{S}_{\mathrm{rob}}\); the score construction is detailed in Section 4.2. Starting from \(\theta^{\mathrm{fo}}\), we apply targeted ZO refinement only to \(\mathcal{S}_{\mathrm{rob}}\), by masking ZO perturbation directions outside the selected layers. Thus, the pipeline first identifies robustness-critical layers and then concentrates ZO updates on the parameters most relevant to robustness degradation. We next provide a theoretical analysis to justify the design of our FO-to-ZO pipeline. Specifically, we aim to answer the following two questions.

4.1.0.1 Question 1: Why ZO after FO alignment?

We first explain why ZO is used as a refinement stage rather than as a replacement for the entire FO alignment process. Although ZO can optimize the smoothed objective using only function evaluations, its finite-difference estimator imposes additional dimension- and smoothing-dependent stepsize constraints, making convergence slower than FO. Therefore, using ZO throughout the whole alignment stage would be unnecessarily expensive. Instead, our pipeline first uses FO optimization to efficiently reach a well-aligned solution, and then applies a small number of ZO steps to refine the local robustness of this solution. To make this distinction precise, we analyze the convergence behavior of ZO optimization on the smoothed objective \(f_\beta\). The following result shows that ZO is a valid optimization procedure, but it requires a more restrictive stepsize than its FO counterpart.

4.1.0.2 Question 2: How does ZO refinement improve robustness?

The convergence result above explains why ZO is a valid post-alignment optimization procedure. We now show why such refinement can improve robustness after FO alignment. Let \(\theta^{\mathrm{fo}}\) be the model obtained by FO safety alignment on the objective \(f\). Although \(\theta^{\mathrm{fo}}\) may perform well on the alignment objective, it does not explicitly control the loss landscape in the surrounding parameter neighborhood. To quantify how much safety alignment is weakened by parameter perturbations, we define the perturbation robustness gap as \[\mathsf{Rob}_{\rho}(\theta) := \mathbb{E}_{v\sim\mathcal{N}(0,I_d)} \bigl[f(\theta+\rho v)\bigr] - f(\theta),\] where \(\rho>0\) is the perturbation scale. A larger value of \(\mathsf{Rob}_{\rho}(\theta)\) means that the loss increases more significantly under parameter perturbations, indicating weaker robustness. Even when FO alignment has already reached a stationary point of the objective \(f\), ZO refinement can still make progress on the neighborhood-averaged objective \(f_\rho\). The following theorem formalizes this intuition.

Theorem 3 formalizes the main intuition behind our framework. Even after FO alignment has converged, the aligned model may remain non-robust in a local neighborhood of its parameters. A ZO refinement step directly optimizes this local behavior through perturbed function evaluations, and can therefore improve robustness without retraining the model from scratch.

4.2 Robustness-Aware Layer Selection↩︎

The second stage of Algorithm 2 determines which layers should be updated during targeted ZO refinement. Rather than adopting generic pruning-based importance scores, we design a robustness-aware criterion that is directly aligned with the goal of refinement. The key idea is simple: if perturbing a layer substantially weakens safety alignment, then this layer is more likely to be responsible for robustness degradation and should be prioritized during ZO refinement. For each layer \(\ell\), we define a layer-wise robustness sensitivity score \(S(\ell)\) by combining two complementary perturbation effects: \[S(\ell) = S_{\mathrm{noise}}(\ell) + \lambda S_{\mathrm{quant}}(\ell),\] where \(\lambda\) balances the contribution of quantization-based sensitivity. The first term captures sensitivity to random local parameter perturbations, which matches the perturbation-based evaluation principle of ZO optimization. The second term captures sensitivity to structured deployment perturbations such as quantization. Together, \(S(\ell)\) measures how much safety alignment is weakened when layer \(\ell\) is locally or structurally perturbed.

Random perturbation sensitivity. For the random perturbation term, we perturb only layer \(\ell\) while keeping all other layers fixed: \[\theta^{(\ell,\mathrm{noise})} = \bigl( \theta_1,\ldots,\theta_{\ell-1}, \theta_\ell+\rho v_\ell, \theta_{\ell+1},\ldots,\theta_L \bigr), \qquad v_\ell \sim \mathcal{N}(0,I),\] where \(\rho\) is the perturbation scale. The corresponding sensitivity is defined as \[S_{\mathrm{noise}}(\ell) = f\bigl(\theta^{(\ell,\mathrm{noise})};\mathcal{D}_{\mathrm{safe}}\bigr) - f\bigl(\theta;\mathcal{D}_{\mathrm{safe}}\bigr).\] This term measures how strongly safety alignment is weakened by a ZO-like local perturbation applied only to layer \(\ell\).

Quantization sensitivity. We further include a quantization-based term to account for structured perturbations that commonly arise in practical deployment. Let \(Q(\cdot)\) denote the quantization operator. For each layer \(\ell\), we quantize only this layer: \[\theta^{(\ell,\mathrm{quant})} = \bigl( \theta_1,\ldots,\theta_{\ell-1}, Q(\theta_\ell), \theta_{\ell+1},\ldots,\theta_L \bigr).\] The quantization sensitivity is defined as \[S_{\mathrm{quant}}(\ell) = f\bigl(\theta^{(\ell,\mathrm{quant})};\mathcal{D}_{\mathrm{safe}}\bigr) - f\bigl(\theta;\mathcal{D}_{\mathrm{safe}}\bigr).\] This term measures how much safety alignment is weakened when layer \(\ell\) is exposed to quantization. After computing \(S(\ell)\) for all layers, we select the layers with the largest scores: \[\mathcal{S}_{\mathrm{rob}} = \operatorname{Top}\text{-}m \{S(\ell)\}_{\ell=1}^{L}.\] The selected set \(\mathcal{S}_{\mathrm{rob}}\) is then used as the update target in the targeted ZO refinement stage. At each refinement step, the random ZO direction is masked outside \(\mathcal{S}_{\mathrm{rob}}\), so that only robustness-critical layers are perturbed and updated. In this way, parameter selection and ZO optimization are coupled through the same perturbation-based principle: perturbations are first used to locate vulnerable layers, and ZO updates are then concentrated on these layers to improve robustness efficiently.

Pruning scores do not capture safety robustness. We compare our robustness-aware scores with SNIP [48] and WANDA [49] in Figure 3. Although pruning-based scores can indicate layers important for maintaining safety alignment, they do not consistently align with robustness sensitivity under perturbation or quantization. This mismatch suggests that pruning scores alone are insufficient for targeted ZO refinement, as they may overlook layers that are critical to safety robustness.

3.8pt

5 Experiments↩︎

5.1 Experiment Setups↩︎

****Models and Datasets. We conduct experiments on Llama-3-8B-Instruct [50] and Qwen2-7B-Instruct [2], using them as the base safety-aligned models for fine-tuning and robustness refinement. We use CB-Safety [51] for both safety alignment and robustness refinement.

****Evaluation Metrics. We evaluate utility and safety using three types of metrics. For utility, we report WikiText-2 perplexity (PPL) [52] and the average zero-shot accuracy over 11 lm-evaluation-harness tasks [53]. For safety, we report ASR on the non-hash HarmBench split [54] using three measures: HarmBench ASR with the official HarmBench-Llama-2-13B-cls classifier, LlamaGuard3 ASR with Llama Guard 3-8B, and AdvBench-style refusal-prefix ASR [33]. Lower PPL and ASR are better, while higher lm-eval accuracy is better.

****Perturbation Settings. We consider three simple post-alignment perturbations. Quantization reduces numerical precision during deployment, where W4A16 denotes 4-bit weight quantization with 16-bit activations, and W4A4 denotes 4-bit weights with 4-bit activations. Parameter noise adds absolute Gaussian noise to all model parameters, controlled by the noise scale \(\sigma\). Activation noise injects Gaussian noise into hidden states, with the noise magnitude scaled by the standard deviation of the corresponding activations and controlled by the relative scale \(\alpha\).

5.2 Safety Alignment Remains Fragile under Simple Perturbations↩︎

Table 1 shows that Llama-3-8B-Instruct remains fragile under simple post-alignment perturbations. Both parameter-level perturbations, including quantization and weight noise, and activation-level noise consistently increase ASR. The degradation becomes more severe under stronger perturbations, such as W4A4 quantization or larger activation noise. These results suggest that safety alignment is effective in the standard setting but remains vulnerable to lightweight parameter- and activation-level perturbations. Additional results on Qwen2-7B-Instruct are reported in Appendix 9.

5.3 FO Safety Alignment Improves Safety Ability, while ZO Refinement Enhances Robustness↩︎

Table 2 shows that 100-step FO safety alignment substantially reduces ASR across all settings, mainly by strengthening the model’s overall safety alignment ability. In contrast, the subsequent layer-robustness-aware ZO refinement is designed as a targeted robustness enhancement stage. Although it introduces only small changes to PPL and lm-eval, it further reduces ASR in most perturbed settings, especially under W4A4 quantization and activation noise. These results indicate that FO alignment is effective for acquiring safety behavior, while layer-robustness-aware ZO refinement further strengthens the robustness against post-alignment perturbations. Consistent results are also observed on Qwen2-7B-Instruct in Appendix 9.

4.4pt

5.4 Efficiency Analysis: Training Time and Storage Overhead↩︎

r0.48 5pt

Table 3 compares the cost of the FO alignment stage and the subsequent ZO robustness refinement stage. FO training is used to establish safety alignment, while ZO refinement is applied afterward to improve robustness. The ZO stage adds only \(152.11\) seconds on top of the \(1146.8\)-second FO alignment stage, increasing the total training time by about \(13.3\%\). It also uses only \(0.32\times\) the peak GPU memory of FO training. These results indicate that ZO refinement provides a lightweight post-alignment robustness enhancement with modest additional cost.

5.5 Ablation Study↩︎

Zeroth-Order Refinement Alone Improves Robustness.

Table 4 shows that replacing late-stage FO updates with ZO refinement already improves safety robustness, even without robustness-critical layer selection. Across most perturbation settings, FO+ZO achieves lower ASR than the FO-aligned model on HarmBench, LlamaGuard3, and AdvBench. The gain is most evident under W4A16 quantization, where ASR drops from \(0.090/0.083/0.137\) to \(0.067/0.057/0.083\) across the three benchmarks. Under activation noise, FO+ZO also consistently reduces ASR at both \(\alpha=0.05\) and \(\alpha=0.08\). Overall, these results show that ZO refinement alone can provide post-alignment robustness gains before introducing robustness-critical layer selection.

4.2pt

Replacing Late-Stage FO with ZO Improves Robustness and Reduces Memory. We further study whether the robustness gain comes from additional training steps or from the use of ZO refinement itself. To this end, we compare three variants: 100-step FO alignment, continued 110-step FO alignment, and 100-step FO alignment followed by 10-step ZO refinement. The comparison between 110-step FO and 100-step FO + 10-step ZO is particularly important, since both use the same total number of optimization steps, but differ in the optimizer used during the final stage. As shown in Figure 4, replacing the last few FO steps with ZO refinement consistently improves robustness compared with simply continuing FO training. Under quantization, weight noise, and activation noise, the FO+ZO variant achieves lower average ASR across the three safety benchmarks. Meanwhile, ZO refinement avoids backpropagation-based gradient computation and therefore reduces GPU memory consumption during the refinement stage. These results suggest that late-stage ZO refinement provides a better robustness–efficiency trade-off than continued FO fine-tuning.

Layer-Wise Robustness Selection Enables Stronger ZO Refinement. Figure 5 compares different parameter-selection strategies for ZO robustness refinement. Robustness-aware selection achieves the lowest HarmBench ASR across all perturbation settings, showing that it can effectively identify layers that are most responsible for safety degradation under perturbations. In contrast, SNIP and WANDA perform well under weight noise but lead to noticeably higher ASR under perturbation. This suggests that the proposed robustness-aware criterion provides a more reliable signal for targeted ZO refinement than the pruning-based importance scores.

6 Conclusion↩︎

In this work, we study the robustness of safety-aligned LLMs under simple post-alignment perturbations, including quantization, weight noise, and activation noise. We theoretically and empirically show that standard safety alignment can remain fragile, and propose a perturbation-guided FO-to-ZO refinement framework where FO establishes safety alignment and ZO further improves robustness. By identifying robustness-critical layers, we restrict ZO updates to these layers for more efficient refinement. Future work includes extending our layer-level selection to finer-grained robustness attribution at the module, neuron, or parameter level.

7 Hyperparameters and Hardware Configuration↩︎

7.1 Training Hyperparameters↩︎

The main hyperparameters used for safety alignment and zeroth-order robustness refinement are summarized in Table 5. We use the same data preprocessing and evaluation protocol across all compared methods unless otherwise specified. All stages use a per-device batch size of 1 and gradient accumulation of 8.

7.2 Hardware Configuration↩︎

The Llama-3-8B-Instruct experiments were run on a local RTX 4090 server. The machine contains 8 NVIDIA GeForce RTX 4090 GPUs with 24GB memory each. The Qwen2-7B-Instruct experiments were run on RTX 4090 GPUs with 48GB memory. Training uses multi-GPU execution, while quantization and safety evaluation use fewer GPUs when possible.

8 Theoretical Proofs↩︎

8.1 Proof of ZO Convergence↩︎

We first prove the convergence guarantee for ZO optimization on the smoothed objective \(f_\beta\). Recall the two-sided ZO estimator \[\widehat g(\theta;v) := \frac{d}{2\beta} \bigl( f(\theta+\beta v)-f(\theta-\beta v) \bigr)v .\] We use the following standard variance bound.

8.2 Proof of One-Step Robustness Improvement↩︎

We next prove that one ZO refinement step can reduce the perturbation robustness gap around an FO-aligned solution. Recall the robustness gap \[\mathsf{Rob}_{\rho}(\theta) := f_\rho(\theta)-f(\theta), \qquad f_\rho(\theta) := \mathbb{E}_{v\sim\mathcal{N}(0,I_d)} [ f(\theta+\rho v) ].\]

9 Additional Experimental Results↩︎

4.2pt

4.2pt

4.2pt

4.2pt

References↩︎