Alignment-Aware Model Adaptation via Feedback-Guided Optimization

TL;DR

This work tackles the misalignment risk inherent in fine-tuning foundation models by introducing AWARE, an alignment-aware fine-tuning framework that injects an external, non-differentiable verifier via policy-gradient regularization. A key contribution is the adaptive per-sample soft-gating mechanism, $\beta_{\\pi_\\theta}(x)$, which balances supervised and alignment-driven updates and triggers abstention for fully misaligned inputs using learned abstention labels. The approach is evaluated on general and biomedical instruction-tuning benchmarks, showing consistent reductions in harmful outputs and hallucinations while preserving task performance; it also demonstrates robustness to adversarial fine-tuning and prompt-based attacks. Overall, AWARE provides a practical, scalable path to alignment-preserving model adaptation by turning alignment signals into on-policy, per-sample guidance rather than relying on static datasets or post-hoc interventions.

Abstract

Fine-tuning is the primary mechanism for adapting foundation models to downstream tasks; however, standard approaches largely optimize task objectives in isolation and do not account for secondary yet critical alignment objectives (e.g., safety and hallucination avoidance). As a result, downstream fine-tuning can degrade alignment and fail to correct pre-existing misaligned behavior. We propose an alignment-aware fine-tuning framework that integrates feedback from an external alignment signal through policy-gradient-based regularization. Our method introduces an adaptive gating mechanism that dynamically balances supervised and alignment-driven gradients on a per-sample basis, prioritizing uncertain or misaligned cases while allowing well-aligned examples to follow standard supervised updates. The framework further learns abstention behavior for fully misaligned inputs, incorporating conservative responses directly into the fine-tuned model. Experiments on general and domain-specific instruction-tuning benchmarks demonstrate consistent reductions in harmful and hallucinated outputs without sacrificing downstream task performance. Additional analyses show robustness to adversarial fine-tuning, prompt-based attacks, and unsafe initializations, establishing adaptively gated alignment optimization as an effective approach for alignment-preserving and alignment-recovering model adaptation.

Figures (15)

• Figure 1: AWARE optimization. Given an input instruction, the policy $\pi_\theta$ samples $k$ responses that are scored by a black-box alignment verifier. The resulting alignment statistics determine an adaptive mixing weight $\beta_{\pi_\theta}(x)$, which interpolates between supervised fine-tuning and alignment regularization. Fully misaligned cases are routed to abstention-based supervision, while remaining samples follow adaptive gradient updates.
• Figure 2: Align-Plots for adaptive soft-gating. Per-sample alignment statistics ($\mu_\theta(x)$, $\sigma_\theta(x)$) across $k$ responses, colored by adaptive mixing coefficient $\beta_{\pi_\theta}(x)$. Well-aligned samples (high $\mu_x$, low $\sigma_x$) receive low $\beta$; uncertain or misaligned samples receive high $\beta$ and stronger regularization. Base Llama3.2-1B on Alpaca+Hex-Phi.
• Figure 3: Effect of alignment regularization and adaptive weighting. Training dynamics under different objectives. Left: task loss (NLL), Middle: semantic quality (BERTScore), Right: harmful score ($HS$). Eq. (3) performs standard supervised fine-tuning, Eq. (5) adds a fixed-weight alignment policy-gradient term, and Eq. (10) uses adaptive weighting. Fixed weights either under-regularize or destabilize training, whereas adaptive weighting preserves task performance and semantics while substantially reducing harmful behavior.
• Figure 4: Alignment dynamics from a misaligned (unsafe) model.(a)Align-Plots before fine-tuning show predominantly misaligned behavior with low mean alignment and high variance. (b) After AWARE, samples shift toward higher mean alignment and lower variance, indicating improved and stabilized safety. (c) Under adversarial fine-tuning, AWARE substantially reduces harmful scores compared to DPO-Cqisafety_dpoc baseline, demonstrating more effective correction from unsafe initialization.
• Figure 5: Robustness to prefilling attacks. Harmful score as a function of the number of prefixed harmful tokens. The base model and DPO-C degrade rapidly as more harmful tokens are prefilled. AWARE-D (default abstention) improves robustness, while AWARE (creative abstention) provides the strongest defense by learning context-aware refusals that are internalized during training rather than relying on fixed templates.

Theorems & Definitions (6)

• Lemma 1.1: Consistency of empirical moments
• proof
• Lemma 1.2: Variance collapse
• proof
• Proposition 1.3: Gate vanishing
• proof