ATTENTIONSWITCH: Jailbreak Defense for LLMS via Joint Learning of Sparse Attention-Head Selection and Scaling

Large language models have achieved remarkable performance across a wide range of tasks, yet they remain susceptible to jailbreak attacks that circumvent safety mechanisms and elicit harmful or unintended outputs. Although safety can be improved through model fine-tuning, such approaches can be computationally expensive and may degrade previously acquired capabilities through catastrophic forgetting. Recent jailbreak defenses have therefore explored targeted interventions on specific model components associated with unsafe behavior. However, existing approaches typically decouple two inherently related decisions: where to intervene and how to modify the selected components. Such sequential optimization may be suboptimal because the optimal intervention location can depend on the intervention applied there, and vice versa. In this work, we introduce ATTENTIONSWITCH, the first jailbreak defense to jointly optimize intervention sites and strengths within a frozen LLM. Specifically, ATTENTIONSWITCH learns a sparse set of attention heads together with query-conditioned multiplicative scaling fac- tors, allowing the outputs of the selected heads to be either attenuated or amplified without altering the directions of their representations. Across popular language models and jailbreak attacks, ATTENTIONSWITCH outperforms existing defenses, with particularly strong results when training data is scarce.