Research paperCOLM 2026

Decoupled Alignment for Robust Plug-and-Play Adaptation

Safety alignment lives in a few MLP gate layers. Put them back.

Authors

Haozheng Luo1,*Jiahao Yu2,*Wenxin Zhang1,*Jialong Li3Chenghao Qiu4Yimin Wang5Eric Hanchen Jiang5Jerry Yao-Chieh Hu1Yan Chen1Binghui Wang6Xinyu Xing1Han Liu1

Equal contribution.

Affiliations

  • 1Northwestern University
  • 2New York University Abu Dhabi
  • 3Stanford University
  • 4Texas A&M University
  • 5University of California, Los Angeles
  • 6Illinois Institute of Technology

Fine-tuning an aligned model for a downstream task quietly undoes part of its safety training. The paper calls this shadow alignment.

A third-party developer starts from a well-aligned base model and adapts it — full fine-tuning, a LoRA adapter, task-specific alignment. The adaptation does not target safety, but it overwrites some of it anyway, and the usual remedy is to realign from scratch.

DAPA asks a narrower question: if the alignment is still present in a sibling model, can it be copied across rather than retrained? That requires knowing where in the network the alignment sits.

Adaptation pulls representations back toward the unaligned model

PCA over the last-token logits of 256 AdvBench prompts, for three checkpoints: a base model, its instruct-aligned version, and a math model fine-tuned from that instruct checkpoint.

The aligned checkpoint sits in its own region. The math model, despite starting from it, drifts back toward the base model — safety is not erased, but it is no longer held.

Figure 3 · Qwen2.5-7B base, instruct, and Qwen2.5-Math-1.5B-Instruct
Two projection directions of the same data

Figure 3: two PCA projections of last-token logits for 256 AdvBench prompts. One model's points form a cluster well away from the other two, which overlap each other.
Original plots from the paper, Figure 3. Read the observation ↗

Alignment is concentrated in the gate projection

Causal tracing answers where. The paper corrupts all hidden states, restores one component at a time, and measures how much the output probability recovers — the indirect effect of that component. Averaged over memory spaces, the gate projection inside the MLP carries more of it than the up or down projections, and attention carries less than the MLP overall. The effect concentrates in the middle layers.

Figure 7: grouped bars of average indirect effect per layer for the gate, up, and down projections across 30 layers. Gate bars rise above the up projection throughout and peak in the middle layers.
Figure 7. Average indirect effect of each MLP module on the last-token hidden representation. Source ↗

Find the smallest set of gate layers that restores refusal, then fuse those layers across from an aligned sibling.

Editing every gate layer would work and would also cost the model its downstream ability. So DAPA borrows delta debugging — the technique for shrinking a failing test case to its minimal cause — and applies it to memory space instead. A set of layers passes if the patched model refuses every question in a small held-out set of unethical prompts.

What transfers is a safety vector: the representation difference between an aligned model and an unaligned one of the same family, added back into the downstream model at the located layers. No gradient step is taken at any point.

  1. 01 /

    Search the memory space

    Let \(\mathbb{S}\) be the universe of MLP memory and \(\pi:\mathcal{S}\to\{0,1\}\) a policy that returns 1 when a candidate set makes the model answer every probe question ethically. The search splits the current set into \(n\) partitions, keeps a passing partition or a passing complement, and doubles \(n\) when neither works — the same descent delta debugging uses on failing inputs. It returns the smallest \(S^{\star}\) it can certify, in \(O(L\log L)\) policy evaluations.

  2. 02 /

    Extract the safety vector

    With \(A\) a safety-aligned model, \(B\) an unaligned base model of the same family, and \(C\) the downstream-adapted model derived from \(A\), the vector \(v_{s}\) is the representation difference between \(A\) and \(B\) restricted to the located layers. It is a compact encoding of what alignment changed.

  3. 03 /

    Fuse it back

    Model fusion adds the vector into the hidden representation of the target model, with \(\lambda\) fit by linear regression on 20 held-out samples — the smallest value that raises the defense success rate without disturbing the downstream task:

    \[ h_{C_{\mathrm{align}}} = h_{C_{\mathrm{org}}} + \lambda v_{s}. \]
Where the edit lands in a Transformer block

Each block updates the hidden state as \(h^{(l)}_{i} = h^{(l-1)}_{i} + a^{(l)}_{i} + M^{(l)}_{i}\), where \(a^{(l)}_{i}\) is self-attention and the MLP is

\[ M^{(l)}_{i} = W^{(l)}_{\mathrm{down}}\,\sigma\!\left(W^{(l)}_{\mathrm{gate}}\gamma\!\left(a^{(l)}_{i} + h^{(l-1)}_{i}\right)\right) \odot W^{(l)}_{\mathrm{up}}\gamma\!\left(a^{(l)}_{i} + h^{(l-1)}_{i}\right). \]

DAPA edits \(W_{\mathrm{gate}}\) in the layers the search returns — on Llama-2-7B that is five layers, 3.25% of parameters. Across the three families the average is 6.26%, at most 8.11%. See Section 2 of the paper.

What the search costs

Memory search runs once, offline. On Llama-2-7B with 50 harmful prompts, the whole pipeline finishes in 48 minutes on one A100 — delta debugging is nearly all of it, with fusion itself around 5 minutes — and no training step is involved. The paper's own figures for retraining on 4×A100 are roughly 4.67 hours for SFT, 9 for DPO, and 18 for RLHF. At inference the fused model is a single ordinary model: no activation editing, no multi-model decoding, no test-time search. See Appendix R.

17 models · three families

Defense success rate mean over 17 models · AdvBench

34.39%48.81%

+14.42 points, up to +51.39

Before → after DAPA

Parameters edited average across three families

6.26%

No training step

3.25% on Llama-2-7B; 8.11% at most

Across 17 models

Defense success rate is the fraction of 128 AdvBench prompts the model refuses, judged by GPT-4o and a keyword classifier together. Perplexity on WikiText-2 and 5-shot MMLU accuracy track what the edit costs; cosine similarity compares responses before and after on ShareGPT prompts.

DAPA across three model families · Table 1
FamilyModelDSR beforeDSR afterPPL beforePPL afterMMLU beforeMMLU afterCos. sim.
Llama-2chinese-alpaca-2-7b82.0387.507.547.4638.7137.430.88
Llama-2-7b37.1642.194.774.7836.3739.300.79
Llama-2-13b37.5046.094.284.2834.7437.080.76
chinese-alpaca-2-13b70.3185.165.635.6048.7747.600.91
Redmond-Puffin-13B22.6647.664.304.3030.0632.380.89
MistralMistral-7B21.0925.784.584.6045.3847.720.76
OpenHermes-2-Mistral-7b33.5946.885.005.0241.2942.460.88
dolphin-2.2.1-mistral-7b24.2241.415.185.1960.1258.250.90
zephyr-7b-alpha24.2232.815.115.1154.0456.730.88
mistral-7B-forest-dpo19.3815.625.135.1054.6254.040.72
dolphin-2.6-mistral-7b-dpo24.2255.475.415.4260.4762.690.91
openchat-3.558.6867.195.155.1061.4058.710.89
Gemmagemma-2b22.0573.447.9224.1533.5724.800.33
Gemmalpaca-2B37.0151.569.9222.0040.9421.170.51
gemma-7b26.5634.386.096.2739.6542.110.66
gemma-7b-ultrachat-sft34.1541.417.177.4842.1129.240.76
gemma-orchid-7b-dpo21.8835.167.227.4242.2638.010.76
Average34.3948.815.917.6044.9842.920.87

DSR and MMLU are percentages; perplexity is on WikiText-2. The paper reports MMLU with standard deviations, omitted here for width — they run from 0.00 to 2.46. Two cases fail rather than degrade: Gemmalpaca-2B loses 19.77 points of MMLU and gemma-7b-ultrachat-sft 12.87, and the small Gemma models are also where perplexity jumps. The paper attributes this to limited parameter redundancy at 2B and to prior task-specific SFT having already reshaped the representations. One model, mistral-7B-forest-dpo, gets less safe.

Full table & protocol

Against other defenses

Seven baselines on the same 128 AdvBench prompts, averaged over all 17 models. The nearest method is Safe LoRA, two points behind.

Average DSR over 17 models · Table 3
MethodAverage DSR
RepE36
LLM Guardrails37
Safety Arithmetic40
ICD42
InferAligner44
Model Surgery45
Safe LoRA47
DAPA49

Per-model columns are given in the paper under anonymized labels. The guardrail baseline is Llama Guard 2. See Table 3 ↗

How it compares to retraining on an equal budget

Given one hour on a single A100 per model, SFT reaches 45.7 average DSR, DPO 50.7, RLHF 54.5, and DAPA 48.8. RLHF is ahead — the paper does not claim otherwise — and the trade DAPA offers is that it reaches a comparable number without a training run, without preference data, and without touching more than a few percent of the weights. See Appendix L, Table 19.

Under attack, at scale, and across modalities

The AdvBench numbers above use direct harmful prompts. Against GCG, which optimizes an adversarial suffix, DAPA raises DSR by an average of 9.62 points on the Llama-2 family; the paper also reports 16.62 points under GPTFuzzer. The method is not tied to the 7B text setting: it transfers to a 70B model and to the language module of a vision-language model.

Figure 6: paired bars for five Llama-2 family models under GCG attack. DSR rises from 36.92 to 49.32, 16.28 to 27.13, 57.39 to 58.72, 29.36 to 48.94, and 6.08 to 10.13 percent.
Under GCG · Llama-2 family
Figure 5: paired bars. Llama3-70B rises from 50.00 to 60.00 percent DSR; LLaVA-1.5-7B rises from 55.00 to 79.00 percent.
Llama-3 70B and LLaVA-1.5-7B

The 70B run uses Hermes-3-Llama-3.1-70B-Uncensored as the target and Llama-3.1-70B-Instruct as the teacher, evaluated on AdvBench. The multimodal run aligns LLaVA-1.5-7B from LLaVA-1.6-vicuna-7B and is evaluated on 103 HarmBench questions. See Section 4.4 ↗

How much editing is enough

Varying the fraction of parameters touched on two Llama-2 models: most of the gain arrives in the first two percent, and the five-layer configuration at 3.25% is already close to what larger budgets reach.

Figure 4: defense success rate against percentage of parameters changed, for Chinese-Alpaca-2-7B and Llama-2-7B. Both curves rise steeply below 2 percent and then stay roughly flat out to 8 percent.
Figure 4. Defense success rate as a function of the share of parameters edited. Source ↗
On current models

Appendix M repeats the experiment on three recent checkpoints adapted for coding, code reasoning, or abliterated outright. All three had lost most of their safety alignment; DAPA recovers an average of 13.87 DSR points — qwen3.5-4B-typescript-coder 14.20→27.40, gpt-oss-20b-heresy 5.30→19.80, GPT-OSS-Code-Reasoning-20B 18.60→32.50 — while AIME24 pass@1 moves by an average of 0.073 points. See Table 28.

Generation is deterministic throughout — temperature 0 — so the safety evaluation is run once per model and has no variance. The capability measurements are repeated three times and reported with standard deviations.

Safety
128 AdvBench prompts · GPT-4o judgment and a keyword classifier must agree
Capability
WikiText-2 perplexity · 5-shot MMLU · Chain-of-Action on Big-Bench · ShareGPT cosine similarity
Attacks
GCG · GPTFuzzer · AutoDAN
Compute
One NVIDIA A100 · 80GB

What this does not cover

DAPA needs a well-aligned sibling in the same family, so it has nothing to offer a model with no aligned counterpart to borrow from. The edit is applied where a held-out probe set says refusal is restored, which is a weaker criterion than safety in general. The two Gemma failures show the method is not uniformly safe for small or heavily SFT-tuned models, and the Big-Bench reasoning scores drop by 2.77% on average, with one model losing 24 points on SocialQA. The paper's own ethics statement notes that its repository includes jailbreak attack implementations and that distillation can carry biases across with the alignment.

Ethics and reproducibility statements ↗
BibTeX
Download
@inproceedings{
luo2026decoupled,
title={Decoupled Alignment for Robust Plug-and-Play Adaptation},
author={Haozheng Luo and Jiahao Yu and Wenxin Zhang and Jialong Li and Chenghao Qiu and Yimin Wang and Eric Hanchen Jiang and Jerry Yao-Chieh Hu and Yan Chen and Binghui Wang and Xinyu Xing and Han Liu},
booktitle={Third Conference on Language Modeling},
year={2026},
url={https://openreview.net/forum?id=6EJ27eMSBB}
}