Rethinking Deep Alignment Through The Lens Of Incomplete Learning

AI summary

Overview

Research area: Safety alignment of large language models, mechanistic interpretability of training dynamics, adversarial robustness.

Technical level: Advanced. The paper combines formal treatment of autoregressive loss gradients and error accumulation with practical training-time interventions and multi-model empirical evaluation.

Scope: The paper explains why safety alignment concentrates in early response tokens ("shallow alignment"), proposes base-favored tokens as a diagnostic of incomplete alignment, and introduces a targeted completion method that improves adversarial robustness across four model families while preserving general capabilities.

What This Paper Is About

Safety-aligned language models remain systematically vulnerable to adversarial attacks, and prior work has shown empirically that alignment mostly changes the first few tokens of a response while later positions look similar to the base model. This paper asks why that pattern arises during training, argues it follows from gradient concentration and signal decay inherent to autoregressive objectives, and uses that explanation to build a targeted fix rather than a broad retraining procedure.

Key Contributions

1. A mechanistic explanation of shallow alignment. The authors attribute position-dependent alignment to gradient concentration and signal decay in sequential loss functions, formalizing how early tokens receive stronger gradient signals and how alignment deviations at early positions propagate to later positions through attention-based influence.
2. Base-favored tokens as computational indicators. They define base-favored tokens as vocabulary elements where the base model assigns higher probability than the aligned model, providing a vocabulary-level detector of undertrained regions that aggregate measures such as KL divergence cannot localize.
3. A targeted completion framework. They propose training-time adaptive L2 penalties on detected base-favored tokens combined with hybrid teacher distillation, translating an inference-time contrastive decoding intervention into a parameter update that avoids inference-time overhead.
4. Empirical validation across model families. They report robustness gains on Llama-2-7B-Chat, Llama-3.1-8B-Instruct, Qwen-2.5-7B-Instruct, and Qwen-3-8B-Instruct, plus safety recovery after fine-tuning degradation and increased deliberative reasoning under adversarial conditions.

Main Findings

• Base-favored token counts are higher in adversarial contexts. In step-by-step generation comparing harmful contexts (HEx-PHI prompts with adversarial prefixes) to benign contexts (Databricks Dolly instructions test set) using Llama-3.1-8B (base) and Llama-3-8B-Instruct (aligned), early positions (0–2) show 45 vs 30 tokens (50% increase), and later positions (15+) show 35 vs 25 tokens (40% increase).
• Base-favored tokens are mostly formatting and structural elements. The most frequent base-favored tokens for Llama-3-8B are predominantly punctuation, special tokens, common words, and structural elements rather than explicitly harmful content, supporting the view that alignment shifts the entire distribution over the vocabulary rather than only suppressing harmful words.
• Inference-time contrastive decoding validates the mechanism functionally. On Llama-3.1-8B-Instruct using 4-token prefill attacks, contrastive decoding reduces prefill attack success from 47.5% to 0.2% while utility benchmarks deviate minimally from baseline (the paper states most changes are within 1–2 percentage points).
• Contrastive decoding sustains distributional separation across positions. The baseline aligned model shows KL divergence of approximately 1.8 at early positions decaying to approximately 0.5 late, while contrastive decoding maintains elevated KL divergence from 10.0 to 6.0 across the sequence.
• Deep alignment on Llama-2-7B-Chat reduces GCG attacks. GCG attack success falls to 0.4% ± 0.9% versus 51.0% ± 42.9% for the baseline, a 99.2% reduction, and is comparable to safety augmentation at 1.6% ± 3.6%. Prefill success is 0.5% versus 23.0% baseline (safety augmentation: 0.8%), and fine-tuning HRR is 4.4% versus 21.4% baseline (safety augmentation: 4.4%).
• Attack resistance and utility hold across four model families. Prefill attack success rates drop as follows: Llama-2-7B-Chat 23.0% to 0.5%, Llama-3.1-8B-Instruct 90.1% to 14.8%, Qwen-2.5-7B-Instruct 85.9% to 44.3%, and Qwen-3-8B-Instruct 96.1% to 16.4%. Fine-tuning HRR drops from 21.4 to 4.4, 25.3 to 12.3, 24.7 to 13.8, and 10.7 to 4.7 respectively, with MMLU, ARC-C, BoolQ, HellaSwag, and Winogrande scores changing only slightly.
• The reported reduction range varies by section of the paper. The abstract describes 48–98% reductions in attack success rates, the contributions list 96-98% attack reduction across model families, and the evaluation section describes 48–96% reductions across four model families.
• Safety recovery after fine-tuning is near-complete. After LoRA fine-tuning on Dolly (rank 32, learning rate 2×10⁻⁴, batch size 128, 1 epoch), the method achieves HRR of 1.0% for Llama-3.1-8B, 0.5% for Qwen-2.5-7B, and 0.3% for Llama-2-7B, compared with 25.5%, 24.7%, and 21.4% for standard LoRA, and better than Vaccine (21.3%, 19.3%, 16.7%), SaLoRA (8.1%, 3.4%, 0.0%), and Safe LoRA (11.0%, not reported for Qwen-2.5-7B, 0.0%). Utility on ARC-C, GSM8K, ToxiGen, and TruthfulQA is preserved or improved in the reported rows.
• Recovery under prefill attack approaches uncompromised base models. The method achieves 44.5% attack success rate, versus Vaccine at 71.9%, SaLoRA at 69.3%, Safe LoRA at 73.4%, and 47.4% for uncompromised base models.
• Deliberative reasoning increases under adversarial conditions. Evaluating Qwen-3-8B-Instruct on 384 AdvBench prompts under prefill attacks, harmful responses drop from 96.1% to 16.4% and reasoning engagement rises from 37.8% to 60.2%. Safe responses with reasoning increase from 9 to 196 (described as 22-fold), while harmful responses with reasoning fall from 136 to 35, described as a shift from reactive recovery to proactive deliberation.

Methodology in Plain English

The authors start from the structure of the standard language modeling loss, which sums a prediction term over every response position. They argue two things follow from that structure: early positions sit at the end of shorter dependency chains and therefore receive stronger gradients, and early positions feed into every later prediction, so parameters affecting them receive contributions from many loss terms while later positions contribute only to their own term. They formalize the resulting alignment error at a position as a sum over earlier positions of a per-position distributional deviation multiplied by an influence score, noting that computing influence for all position pairs is O(t²).

To detect where learning stayed incomplete, they compare the base and aligned models at each generation step and collect tokens where the base model assigns higher probability. These base-favored tokens are extracted with a top-k selection over the difference in logits. First they use them at inference time: at each step they subtract a penalty proportional to the base–aligned logit difference before the softmax, which suppresses base-favored tokens during generation. This is framed as a functional test that the tokens are exploitable mechanisms rather than statistical noise.

Because the inference-time version requires loading two models and extra forward passes per step, they move the same idea into training. For harmful training contexts, they identify the top-k base-favored tokens and add an adaptive L2 penalty on the aligned model's logits for those tokens, scaling the penalty weight by the density of base-favored tokens in the vocabulary. A hybrid teacher is built by interpolating base and aligned logits with a weight of λ = 1.2 that amplifies the aligned model's preferences while retaining base information, and the student is trained with temperature-scaled KL divergence against that teacher at temperature 2.0. Training uses HEx-PHI data (330 harmful pairs) for the completion loss with top-100 base-favored token selection, GSM8K for distillation supervision, and 20 epochs.

Evaluation uses three attack protocols (prefilling, fine-tuning robustness, and GCG optimization), utility benchmarks (MMLU, ARC-C, BoolQ, HellaSwag, Winogrande, ARC-E, GSM8K, ToxiGen, TriviaQA, TruthfulQA), and comparisons against Vaccine, SaLoRA, and Safe LoRA for recovery.

Why This Matters

The paper reframes shallow alignment from an empirical observation into a consequence of how autoregressive training distributes gradient signal, which suggests that vulnerabilities shared across SFT, RLHF, and DPO have a common origin rather than method-specific flaws. That framing matters for research because it points to position-dependent training signals as a general design variable, and it argues that completing distributional alignment may also affect deliberative behavior rather than only refusal rates.

Real-world applications:

• Safety recovery for fine-tuned deployments. Organizations that adapt aligned models to domain tasks with LoRA can restore Harmfulness Rejection Rates close to base levels (1.0%, 0.5%, 0.3% in the reported models) instead of accepting the degradation seen with standard LoRA.
• Open-weight model release hardening. Model providers can apply a lightweight training intervention rather than depending on inference-time guardrails that require concurrent model loading and per-step penalty computation.
• Red-teaming and safety evaluation. Base-favored token counts give evaluators a per-token diagnostic for locating undertrained regions in a response, complementing aggregate metrics such as KL divergence and top-100 overlap.
• Serving systems with no inference-time budget. Because the intervention is folded into parameters, deployments avoid the doubled memory and extra forward passes that contrastive decoding requires.

Industry relevance centers on the fine-tune-then-deploy pipeline, where safety properties are known to degrade after task adaptation, and on cost-sensitive serving environments where inference-time interventions are operationally difficult.

Future Directions

• Scaling the targeted completion method to larger models, which the conclusion explicitly identifies as future work.
• Examining how alignment completeness relates to other safety dimensions beyond the attack types and benchmarks studied here.
• Determining whether the gradient concentration analysis and base-favored token diagnostics transfer to additional model families, architectures, and longer or multimodal response contexts, which the paper does not report.
• Reducing the O(t²) cost of computing pairwise influence scores, and quantifying the training-time overhead of 20-epoch completion relative to the inference-time overhead the method was designed to avoid, which is not reported.

Target Audience

This paper is best suited to alignment and safety researchers, mechanistic interpretability practitioners, and machine learning engineers responsible for fine-tuning or deploying aligned models. Readers need familiarity with autoregressive language modeling objectives, KL divergence, LoRA fine-tuning, and standard adversarial attack protocols such as prefill and GCG attacks to follow the formal sections and interpret the benchmark tables.

Authors’ abstract

Large language models exhibit systematic vulnerabilities to adversarial attacks despite extensive safety alignment. We provide a mechanistic analysis revealing that position-dependent gradient weakening during autoregressive training creates signal decay, leading to incomplete safety learning where safety training fails to transform model preferences in later response regions fully. We introduce base-favored tokens -- vocabulary elements where base models assign higher probability than aligned models -- as computational indicators of incomplete safety learning and develop a targeted completion method that addresses undertrained regions through adaptive penalties and hybrid teacher distillation. Experimental evaluation across Llama and Qwen model families demonstrates dramatic improvements in adversarial robustness, with 48--98% reductions in attack success rates while preserving general capabilities. These results establish both a mechanistic understanding and practical solutions for fundamental limitations in safety alignment methodologies.