🤖 machine learning

Leveraging Association Context Retrieval in Knowledge Edit- ing to Build White-Box Attacks on LLMs

This paper proposes a novel white-box attack on large language models that leverages association context retrieval within a knowledge editing framework to induce unsafe behavior across entire thematic categories while preserving general model performance.

Original authors: Roman Maksimov, Vladimir Aletov, Vladimir Solodkin, Dmitry Bylinkin, Daniil Medyakov, Aleksandr Beznosikov

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Roman Maksimov, Vladimir Aletov, Vladimir Solodkin, Dmitry Bylinkin, Daniil Medyakov, Aleksandr Beznosikov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are the powerful computer programs that can write stories, answer questions, and help with complex tasks. To make these tools safe for everyday use, developers train them to refuse harmful requests, such as instructions on how to build a weapon or spread dangerous misinformation. This process, known as alignment, is meant to act as a guardrail, ensuring the machine behaves like a helpful assistant rather than a source of danger. However, as these models become more autonomous and are entrusted with sensitive decisions, researchers need to understand exactly how strong these guardrails are. If the safety mechanisms can be bypassed, it is crucial to know how, so that stronger protections can be built. This field of study involves looking inside the model's internal workings to see how it processes information and makes choices, rather than just testing it from the outside.

A team of researchers has discovered a new way to test these safety systems by treating the model's refusal to answer as a specific piece of knowledge that can be overwritten. Instead of trying to trick the model with clever wordplay or hidden commands, they used a technique called knowledge editing. This method allows scientists to pinpoint the exact part of the model's brain responsible for a specific fact and change it. For example, if a model believes a certain fact is true, researchers can alter the internal connections so the model believes a different fact instead. The researchers realized that this same technique could be used to change the model's behavior: instead of teaching it a new fact, they could teach it to stop refusing harmful requests. They found that by modifying the model's internal weights to associate a harmful request with a compliant response, they could effectively remove the safety guardrails for entire categories of dangerous questions.

The researchers tested this approach on several different types of language models, including those based on popular architectures like GPT, Llama, and Qwen. They focused on a method that first locates the specific layer in the model where a decision is made and then edits that layer to change the outcome. In previous attempts to break safety filters, researchers often relied on finding a single "refusal direction" in the model's math and pushing against it. The authors of this study found that this older method was fragile; it often broke the model's ability to speak coherently or caused it to fail at basic tasks. In contrast, their new approach was more surgical. They identified that when a model is edited using these precise methods, it tends to assign a very high probability to a specific, compliant starting phrase, such as "Sure, here is the plan." By targeting this specific reaction, they could force the model to agree to harmful requests without destroying its general intelligence.

To make this work on real-world questions, the team had to solve a tricky problem: how to find the right part of a sentence to edit. In simple fact-checking tasks, the subject is clear, like a person's name. But in a complex instruction like "How do I create a virus?", the subject is the entire instruction itself. The researchers developed a way to trace the model's thinking to find exactly which words in the prompt were most responsible for the refusal. They then used the model's own knowledge to build a richer context around those words, creating a more stable target for the edit. This allowed them to remove the refusal behavior not just for one specific question, but for an entire family of similar questions. For instance, if they edited the model to ignore safety rules for writing malware code, the model would also become willing to write other types of harmful code, even if the wording was slightly different.

The results of the experiments showed that this method was highly effective. On one of the models tested, the new approach increased the model's willingness to provide harmful information to a score of 3.12 out of 4, compared to 2.98 for a standard editing method and 3.01 for the older refusal-based approach. More importantly, the model remained useful and coherent. When tested on harmless questions, the edited model still answered correctly, maintaining a high level of performance on general knowledge tasks. The older methods, which tried to push the model away from refusal, often caused the model to lose its ability to think logically, resulting in gibberish or broken responses. The new method managed to unlock the harmful behavior while keeping the rest of the model's brain intact. This suggests that the safety filters in these models are not just a single wall, but a set of specific associations that can be carefully removed without collapsing the whole structure.

The study also highlighted the limits of current defenses. The researchers found that while their method worked well on older or less aligned models, it was harder to apply to the newest, most carefully trained models. Even on these robust systems, the method still showed some ability to bypass safety, though the effect was smaller. This indicates that as models become more advanced, the safety mechanisms become more complex, but they are not yet impenetrable. The researchers emphasized that their work was conducted in a controlled environment using open-source models that anyone can access. They did not provide a step-by-step guide for creating dangerous tools, but rather demonstrated a vulnerability that needs to be fixed. By showing that safety can be compromised through precise internal edits rather than just clever prompts, the study offers a new perspective on how to build models that are truly safe. The findings suggest that future safety measures need to be more robust against these internal changes, ensuring that the refusal to do harm is deeply embedded in the model's structure rather than just a surface-level reaction.

Technical Summary: BRAIn Lab – Leveraging Association Context Retrieval in Knowledge Editing for White-Box LLM Attacks

1. Problem Statement

As Large Language Models (LLMs) are increasingly deployed with greater autonomy and entrusted with sensitive data, ensuring their alignment (i.e., preventing harmful behavior) is critical. However, recent studies indicate that current alignment mechanisms are not robust. While "black-box" jailbreaks are well-studied, "white-box" attacks—where the adversary has direct access to model weights—offer more potent manipulation capabilities.

Existing white-box approaches often rely on refusal direction suppression (orthogonalizing the model's response along a "refusal" vector). The authors identify two critical flaws in this approach:

Loss of Expressiveness: Orthogonalization eliminates the model's ability to represent an entire dimension of the parameter space, leading to a loss of logical coherence and incoherent responses.
Sensitivity and Generalization: The effectiveness of refusal direction suppression is highly sensitive to the weight parameter $\alpha$ . Furthermore, a mean refusal direction computed over a subset of prompts often fails to generalize even to paraphrased instructions.

While Knowledge Editing (KE) frameworks (specifically locate-then-edit methods like ROME and MEMIT) offer more precise modifications with less impact on general behavior, previous attempts to adapt them for attacks relied on backdoors (inserting specific rare words into prompts). This requirement for explicit prompt injection makes such attacks easily identifiable. The paper proposes a method to remove alignment without needing backdoors or explicit prompt injections, instead modifying weights to unlock entire thematic categories of harmful prompts.

2. Methodology: BRAIn Lab

The proposed method reframes alignment removal as a knowledge editing problem over behavioral associations. The core innovation is the use of Association Context Retrieval to construct the subject representation ( $k^*$ ) for the editing process.

2.1 Custom Tracing and Subject Selection

Standard KE methods trace causal effects on factual triples (e.g., "Cristiano Ronaldo plays soccer"). In the context of harmful prompts, the "subject" is not a concrete entity but an instruction.

Tracing: The authors adapt causal tracing to identify which tokens in a harmful prompt most influence the probability of an agreement token (e.g., "Sure") rather than the most probable token.
Subject Extraction: Through tracing, they find that the AIE (Average Indirect Effect) peaks at the final tokens of the instruction span (e.g., "Give an argument for..."). Consequently, the instruction itself is treated as the subject $s$ .
Object Definition: The target object $o^*$ is defined as a compliant opening phrase (e.g., "Sure! The easiest way is:") specific to the prompt's structural category (Malware/Hacking, Instructions/Guide, Creative Harm).

2.2 Association Context Retrieval (The Core Contribution)

Standard locate-then-edit methods compute the subject vector $k^*$ using randomly generated prefixes ( $x_j$ ) concatenated with the subject. The authors argue this context-agnostic encoding fails to probe the associative structure of the subject, leading to poor generalization.

The Modification:
Instead of random prefixes, the method enriches the prefixes $x_j$ with the model's own knowledge related to the subject.

The model is forced to generate prefixes not just from the <bos> token, but from predefined phrases that trigger reasoning about the subject in a neutral context.
This results in a $k^*$ vector that captures a distribution over associative contexts rather than a single surface-level occurrence.
This accumulation of activations makes $k^*$ a more faithful representation of the subject, improving the transferability of the edit to semantically and structurally similar prompts.

2.3 Optimization

The method utilizes standard locate-then-edit algorithms (ROME, MEMIT for dense models; MoEEdit, MoTE for MoE models) to optimize the value vector $v^*$ . The goal is to force the model to start its response with the desired compliant phrase ( $o^*$ ) when presented with the subject ( $s$ ), thereby suppressing refusal behavior.

3. Key Contributions

The paper introduces a category-level white-box attack with the following characteristics:

Portable: Uses a modified tracing procedure to extract subjects from any harmful prompt without manual selection.
Editor-agnostic: Compatible with any locate-then-edit framework (e.g., ROME, MEMIT, BIRD, GLAME).
Architecture-agnostic: Successfully applied to GPT-based, Llama-based, and Qwen3-MoE architectures.
Local and General: Unlocks entire structural or semantic categories of prompts while minimizing damage to general model performance (Locality).
No Backdoors: Unlike previous KE-based attacks, it requires no special rare words or explicit prompt injections; the attack is embedded solely in the weight modification.

4. Experimental Results

Experiments were conducted on GPT4ALL-J, Llama-3-8B, and Qwen3-30B-A3B using datasets including AdvBench, HarmBench, and JailbreakBench.

4.1 Harmfulness and Locality

Harmfulness: The proposed method (specifically ROME+pref and MEMIT+pref) significantly outperformed the refusal-based baseline in generating harmful content. For instance, on GPT4ALL-J, MEMIT+pref achieved a Harmfulness score of 3.12 (vs. 3.01 for Refusal) and 3.19 on Llama-3-8B (vs. 2.72 for Refusal).
Locality: Crucially, the method preserved general model capabilities better than refusal-based attacks.
- On Llama-3-8B, the refusal baseline caused a 3.79% drop in MMLU accuracy, whereas MEMIT+pref caused only a 0.05% drop.
- The method avoided the "broken response" issue (where the model outputs incoherent text) common in high- $\alpha$ refusal attacks.

4.2 Generalization

The method demonstrated strong generalization to unseen prompts within the same structural and semantic categories.

Structural Generalization: MEMIT+pref achieved a score of 2.54 on unseen structural prompts for GPT4ALL-J, outperforming the refusal baseline (2.33).
Semantic Generalization: The method successfully unlocked knowledge for semantically related prompts, confirming that the associative context retrieval allows the edit to transfer beyond the specific training examples.

4.3 Ablation Studies

Tracing Validation: Tracing confirmed that instruction spans act as the subject, validating the custom tracing procedure.
Refusal Baseline Sensitivity: The ablation study on the refusal direction weight $\alpha$ showed a trade-off: higher $\alpha$ increased harm but drastically increased the fraction of "broken" (incoherent) responses. The proposed method achieved high harm with minimal broken responses.

5. Significance and Claims

The paper claims that Knowledge Editing constitutes a milder, more effective approach to alignment removal compared to dimensionality reduction via refusal directions.

Mechanism: By treating the "agreement" phrase as an object in associative memory and enriching the subject representation with context, the attack becomes more robust and generalizable.
Defensive Insight: The authors argue that their findings highlight the fragility of current alignment. They demonstrate that refusal-direction suppression can severely degrade model coherence, whereas locate-then-edit approaches produce more localized changes.
Ethical Stance: The work is presented as a controlled study of vulnerabilities to inform the development of more robust safeguards. The authors emphasize that the method is not a recommendation for bypassing safety systems but a tool for understanding the limits of current alignment techniques. They explicitly state that the attack requires white-box access, limiting its practical applicability in production systems, and that no harmful generations are released at scale.

In summary, the paper proposes that by leveraging association context retrieval within the knowledge editing framework, attackers can bypass safety alignments more effectively and with less collateral damage to model utility than traditional refusal-suppression methods.