AttentionSwitch tunes selected AI attention heads to resist jailbreaks

Doniyorkhon Obidov, Honggang Yu, Xiaolong Guo and Kaichen Yang describe AttentionSwitch in a preprint submitted to arXiv on November 5, 2026. The work has not been peer reviewed. Its central question is whether a defense can learn both where to intervene inside a language model and how strongly to intervene, rather than treating those as separate decisions.

Jailbreaks are prompts designed to make a model provide harmful assistance it would otherwise refuse. A defense must do more than increase refusals: a system that rejects nearly everything may block attacks while becoming useless for ordinary questions. The study therefore measures both attack rejection and the preservation of legitimate responses.

Adjusting strength, not direction

AttentionSwitch targets attention heads, components that process relationships among tokens in parallel. It leaves the model’s original weights unchanged and adds small controllers that multiply selected heads’ output vectors by positive numbers. A multiplier below one weakens a head’s contribution; a multiplier above one strengthens it. Because every component of a head’s output receives the same multiplier, its mathematical direction stays unchanged.

The adjustment is dynamic. Each selected head’s internal query vector determines its multiplier at each input-token position. The selected heads are fixed after training, but their scaling factors can vary across prompts and positions.

Training has three stages: initialize scaling functions across all heads, jointly optimize head selection and scaling, then refine the scaling functions for the chosen subset. A mathematical relaxation lets gradient-based training explore otherwise discrete choices of heads. Harmful examples are paired with refusal responses, while benign examples are paired with the original model’s continuations.

During use, scaling happens only while the model processes the input prompt, not while it generates subsequent tokens. Its effects persist through cached internal representations of that prompt. Deployment requires access to the model’s internal computations, rather than only a text-generation interface.

Strong results with little training data

The researchers tested five models ranging from 7 billion to 27 billion parameters: DeepSeek-LLM, OpenChat-3.5, Llama-3.1, Vicuna and Gemma-3. Each defense received either 30 or 120 training examples, evenly split between harmful and benign inputs. Harmful training examples came from three prompt categories; six additional attack types appeared only in evaluation.

The main evaluation included 2,115 harmful prompts, 500 benign instruction-following prompts and 250 benign prompts containing superficially alarming language. The latter test over-refusal: rejecting an innocent request because it sounds dangerous out of context.

With the default intervention on 20% of attention heads, AttentionSwitch’s average attack rejection rate across the five models reached 96.58% with 30 training examples and 99.05% with 120. Its benign pass rates were 97.68% and 96.84%, respectively. Average scores on MMLU, a multiple-choice knowledge benchmark, remained close to those of the undefended models.

The comparison illustrates why rejection alone is insufficient. Full supervised fine-tuning achieved a similar average rejection rate of 96.14% in the 30-example setting, but its benign pass rate fell to 24.64%. AttentionSwitch did not win every metric or every model comparison; its principal strength was the balance between rejecting harmful requests and retaining ordinary usefulness.

Controlled comparisons supported the joint-selection approach. With 30 examples, scaling all heads produced 92.26% average attack rejection. Selecting heads first using another defense’s procedure, then learning AttentionSwitch’s scaling, produced 94.73%. Joint selection and scaling reached 96.58%, with similar benign pass rates. Across the tested head budgets, rejection generally peaked around 20–30% rather than improving as more heads were adjusted.

Adaptive attacks reveal the limits

The strongest caution comes from attacks built against the defended models themselves. In the 30-example configuration, adaptive GCG optimized adversarial suffixes directly against AttentionSwitch, while adaptive PAIR iteratively refined prompts using its responses. Each was tested on 50 harmful prompts per model.

Average rejection fell to 86.80% for adaptive GCG and 77.60% for adaptive PAIR. DeepSeek rejected only 52% of adaptive GCG prompts. These results rule out interpreting the near-perfect main-test scores as comprehensive protection.

There were also residual usefulness costs: average over-refusal reached 10.08% with 30 examples and 13.04% with 120. Most response judgments came from WildGuard, an automated evaluator. It agreed with a blinded human annotator on 525 of 540 sampled responses, or 97.22%, supporting—but not guaranteeing—the reliability of those measurements.

The evidence establishes a promising targeted defense under the tested conditions, not universal safety. Five models, selected benchmarks and two adaptive attack methods cannot cover every deployment or adversary, and preserved MMLU scores do not establish preservation of every capability.