PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.17836v1cs.LG

TL;DR

Because alignment is not robust, the paper develops a category-level white-box attack that treats alignment removal as behavioral knowledge editing. Context association retrieval produces transferable edits that unlock semantic and structural prompt categories while preserving more general model performance.

  • Problem

    Alignment is not robust, motivating investigation of methods that can induce unsafe behavior in increasingly autonomous language models.

  • Method

    The method retrieves subject-conditioned reasoning contexts to estimate edit keys and trigger compliant openings across harmful prompt categories.

  • Results

    Generalization reaches 2.38/2.54 for ROME/MEMIT within semantic categories and 2.41/2.37 within structural categories, versus 2.33 and 2.31 for Refusal.

  • Takeaways & Limitations

    The attack can unlock entire structural or semantic prompt categories while causing much less damage to general model performance.

  • Takeaways & Limitations

    The study covers three open-weight models but not the latest versions released in 2026, and it assumes direct access to model parameters.

Abstract

from arXiv · show

As large language models (LLMs) are granted increasing autonomy, it is essential to investigate methods that can induce unsafe behavior. We propose a novel white-box attack inspired by locate-then-edit approaches from the field of Knowledge Editing. Our choice is motivated by the observation that models edited with such schemes tend to assign unusually high prediction probabilities to the edit target, a property that is particularly advantageous when designing attacks. We modify the editing framework by incorporating as- sociative knowledge retrieved from the model, thereby extending constraint removal to an entire thematic category rather than being limited to prompts from a predefined dataset. Experiments with various archi- tectures demonstrate improved attack effectiveness over competing methods without dealing critical damage to general model performance.

1 Introduction

The paper introduces a category-level white-box attack that reframes alignment removal as knowledge editing over behavioral associations. It modifies model weights without explicit prompt injections and aims to unlock harmful prompt categories while limiting general-performance damage.

  • Our contribution: The attack replaces explicit backdoor words pasted before prompts with weight modifications that find alternatives within harmful prompts.This avoids the easy identification associated with selecting and prepending specific backdoor words.
  • Our contribution: The method reframes alignment removal as a knowledge editing problem over behavioral associations at the category level.It extends the attack beyond prompts from a predefined dataset to broader structural or semantic categories.
  • Portability: The attack is portable because a modified tracing procedure identifies the entity to edit in any harmful prompt.The procedure is based on tracing from Meng et al. [2022a].
  • Editor and architecture agnosticism: The approach is editor-agnostic because any knowledge-editing framework described in Section 2 can be used for the task.The paper provides additional details in Appendix A.
  • Editor and architecture agnosticism: The method is architecture-agnostic and applies across GPT-based, Llama-based, and Qwen3-MoE models.Experiments on these model families are reported in Section 6.

2 Preliminaries

Knowledge editing replaces an original object in a subject–relation–object triple by adjusting internal representations to increase the new object’s conditional probability. Locate-then-edit methods focus on a critical FFN layer, representing the subject with a key vector and updating weights using a solved output vector.

  • Problem statement: Knowledge editing replaces an original object o with a new object o∗ by increasing P(o∗| s, r) through internal representation adjustment.The paper illustrates this by changing Cristiano Ronaldo’s sport from soccer to hockey.
  • Other methods: Basic methods such as FT, FT-L, AdaLoRA, and MEND face a difficult trade-off among performance, computational cost, and model degradation.The authors therefore consider these methods irrelevant to their work.
  • Locate-then-edit: Locate-then-edit methods ROME and MEMIT model the last FFN layer as a linear associative memory matching key and value vectors.Tracing identifies a critical layer where these representations are used for editing.
  • Locate-then-edit: The key vector k∗ represents the subject, while the output vector v∗ is obtained by minimizing an objective before computing and inserting the weight updates.Subsequent methods use the same form of k∗ described in the preceding formulation.
  • Locate-then-edit: PMET extends these ideas by optimizing attention weights while applying changes only to the FFN, and later work adapts them to Mixture-of-Experts architectures.Alignment-removal approaches have also applied knowledge-editing ideas using backdoors formed from rare prompt-prefix words.

3 Motivation

The motivation centers on locating the model components that store factual associations and addressing the difficulty of editing complex, multi-hop relations. Causal tracing identifies influential mid-layer MLP modules, while naive locate-then-edit methods fail on two-step relations.

  • Tracing: Causal tracing identifies a critical middle layer where factual associations are stored and subject-specific keys are activated.Mid-layer MLP modules at the subject’s last token show strong indirect effects; GPT-J, for example, typically uses layer 5.
  • Multi-hop connections: Naive locate-then-edit approaches fail on two-step multi-hop semantic relations, motivating methods that can edit connected facts.In a triple chain, the first object becomes the subject of another relation, making the edit more complex than changing a single fact.

4 Methodology

The methodology traces agreement-phrase probabilities to identify influential prompt tokens, extracts subjects from harmful prompts, and constructs associative-context prefixes for knowledge editing. It assigns category-specific target responses and optimizes representations intended to generalize across related semantic structures.

  • Custom tracing: Custom tracing identifies prompt tokens that most influence an agreement phrase’s probability, enabling extraction of a proper subject from any harmful prompt.The original procedure instead traces the most probable answer, which here would be a disagreement phrase.
  • Selection of subject and object: Subjects are extracted by locating prepositions with regular expressions, with “How to...” prompts handled by skipping the first preposition and ending at the second.The extraction targets the sentence beginning with the main instruction.
  • Target construction: Malware/Hacking prompts target “Sure! The easiest way is:”, Instructions/Guide prompts target “Sure! The plan is:”, and Creative harm prompts target “Sure! Here it is:”.Prompts are assigned to these structural categories using an LLM-as-a-judge framework.
  • Constructing the prefix with s.: Associative-context prefixes replace random prefixes so k* captures a distribution over contexts rather than a single surface-level occurrence of the subject.The accumulated activations are collected while the model reasons about the subject in a neutral context.
  • Optimization of k* and v*: The resulting k* is intended to represent the subject more faithfully and improve performance on related semantic structures.After generating subject prefixes, k* is determined by equation (1), while v* is optimized using the locate-then-edit framework.

5 Experimental setup

The experiments use established harmless and harmful instruction datasets organized into ten semantic categories, evaluate the approach across harmfulness, generalization, and locality, and test it with multiple model architectures. The study integrates the method into ROME and MEMIT and compares it with refusal-based and locate-then-edit baselines.

  • Datasets and categories: The evaluation uses Alpaca as the harmless dataset and AdvBench, TDC2023, HarmBench, StrongREJECT, and JailbreakBench as harmful datasets.The harmful instructions are classified into ten JailbreakBench categories, with training and test splits of 704 and 254 examples.
  • Datasets and categories: The harmful dataset is mapped into ten semantic categories spanning harassment, malware, physical harm, economic harm, fraud, disinformation, sexual content, privacy, expert advice, and government decisions.A classifier assigns each instruction to one JailbreakBench category before semantic-category evaluation.
  • Model architectures: The tested models include fine-tuned GPT4ALL-J version “1.2-groovy” and instruction-tuned Llama-3-8B-Instruct, named Llama-3-8B for clarity.GPT4ALL-J is used instead of standard GPT-J-6B because GPT-J-6B is not properly aligned or tuned to question-answer formatting.
  • Evaluation metrics: The approach is evaluated by Harmfulness, Generalization, and Locality, measuring harmful outputs, approximation of the harmful instruction’s semantic neighborhood, and preservation of harmless responses.Harmfulness is assessed using two approaches, including LLM-as-a-judge with DeepSeek.
  • Methods and evaluation: Experiments integrate the proposed approach into ROME and MEMIT and compare it with the refusal-based baseline.Harmfulness is evaluated on each harmful training prompt before and after intervention using greedy decoding.
  • Baselines and model-specific methods: Locate-then-edit comparisons evaluate ROME, MEMIT, BIRD, GLAME, and PMET separately, while MoE models use MoEEdit and MoTE.The refusal-based baseline samples 100 harmless prompts to compute a common harmless mean activation and also includes RDO and Heretic.

6 Experimental results

The experiments evaluate harmfulness, locality, and generalization across edited prompts, harmless and benchmark prompts, structural categories, and semantic categories. The prefix construction improves attack effectiveness for ROME and MEMIT, while revealing a harmfulness–locality trade-off and strong generalization beyond seen instructions.

  • Experimental setup: The evaluation measures Harmfulness on editing prompts, Locality on harmless, MMLU, and ARC-C prompts, and Generalization across structural and semantic categories.These measurements define the experimental scope of Section 6.
  • Harmfulness and Locality: ROME with the prefix construction surpasses the refusal baseline on Harmfulness and Locality, supporting knowledge editing as a milder alignment-removal approach.ROME alone matches the refusal baseline on Harmfulness, while adding the prefix pushes it beyond that baseline.
  • Harmfulness and Locality: On GPT4ALL-J, the strongest Harmfulness attack also has the lowest Locality and largest MMLU and ARC-C drops within its block.This result exemplifies the harmfulness–locality trade-off analyzed in Section 7.2.
  • Generalization: The prefix construction consistently improves generalization for both ROME and MEMIT, extending harmful behavior beyond the instructions used for training.The approach encodes topic knowledge in k∗, enabling evaluation on previously unseen instructions.
  • Generalization: For unseen instructions within a structural category, generalization is approximately the same as for prompts generalized by semantic similarity.The results also attribute near-certain target-token output probability on current-structure instructions to s → o∗ overfitting.

7 Ablation studies

The ablations identify instruction-span tokens as the causal-tracing subject and select intervention layers for GPT4ALL-J and Llama-3-8B-Instruct. They also examine broken responses and refusal-direction strength across models.

  • Causal tracing: AIE peaks at the final tokens of instruction spans, so the instruction is treated as the subject in this setup.Examples include “Give an argument for” and “Write a manual on how to.”
  • Intervention layers: Layers 5 and 9 were selected for ROME intervention in GPT4ALL-J and Llama-3-8B-Instruct, respectively.The passage also states that MEMIT layer selection follows this analysis, but its specific layer is truncated.
  • Broken responses: The study excludes prompts with pre_score = 4 or post_score = 0 from main metrics and separately ablates cases where the model loses substantive response structure.The post_score = 0 cases are compared with refusal-based methods in a separate ablation.
  • Refusal-direction ablation: α values tested were 1.0, 0.6, 0.3 for GPT4ALL-J and 2.0, 1.0, 0.6 for Llama-3-8B and Qwen3-30B-A3B.Higher GPT4ALL-J values were not tested because α = 1.0 fully breaks the model; α < 0.6 yielded poor Llama-3-8B performance.
  • Refusal-direction ablation: Random prefixes produce nearly the same number of broken cases as the proposed approach, so they are omitted under this metric.The methods are surpassed only by refusal-based shifts in the reported comparison, whose continuation is truncated.

8 Conclusion

The conclusion presents context association retrieval as a way to edit instruction-like subjects toward structurally matched compliant openings. The method estimates k∗ from subject-conditioned reasoning contexts, enabling more transferable edits while remaining compatible with ROME and MEMIT.

  • 8 Conclusion: The method treats agreement as an object in associative memory and edits instruction-like subjects to trigger structurally matched compliant openings.This frames the attack around editing subjects toward compliant output openings.
  • 8 Conclusion: Context association retrieval estimates k∗ from subject-conditioned reasoning contexts rather than context-free random <bos>-prefixes.The retrieval procedure uses a distribution of contexts conditioned on the subject.
  • 8 Conclusion: The resulting edits are more transferable and remain plug-and-play with ROME and MEMIT.The passage identifies compatibility with both ROME and MEMIT as a practical property.

9 Limitations

The study is limited to three open-weight models and omits versions released in 2026. It also does not evaluate all current locate-then-edit baselines with the prefix modification, including methods that differ in context-association mechanisms.

  • The experiments cover three open-weight models with different architectures and sizes, released from 2022 to 2025.
  • The study does not cover the latest model versions released in 2026.
  • Not all current locate-then-edit baselines are tested with the prefix modification.
  • GLAME already includes a context-association mechanism, whereas PMET and BIRD do not.

10 Ethical Considerations … C Implementation details

The paper reports dual-use risks and controlled research practices, then provides supplementary formulas, judge prompts, and implementation details for ROME, MEMIT, and MoE-adapted methods.

  • 10 Ethical Considerations: The attack can increase harmful-generation probabilities and weaken refusal behavior, motivating controlled study to inform stronger safeguards.The work explicitly identifies white-box attacks on aligned LLMs as dual-use research.
  • Appendix: The supplementary materials accompany the paper on association-context retrieval and white-box attacks.The passage identifies the supplementary-materials document but provides no further technical content.
  • A Additional formulas: The formulas define hidden states, attention outputs, MLP outputs, and GPT-like MLPs with two fully connected layers and a rectifying nonlinearity.These definitions establish the dense-transformer notation used for the editing derivation.
  • A Additional formulas: N random context examples generated by the model are averaged to form k∗, while v∗ is optimized to maximize o∗ probability and control essence drift.The prompts p′ use the paper’s Section 7.1 template rather than the original “{subject} is a” formulation.
  • A Additional formulas: The calculated pair (k∗, v∗) is inserted through an MLP weight update; in LLaMA-like models, v∗ enters down_proj via a rank-one update.The key is derived from the pre-down_proj gated activation, and C is precomputed from Wikipedia-text covariance estimates.
  • A Additional formulas: MoE layers route tokens to top-K experts, whose expert-specific keys and values produce contributions combined by the sparse gating distribution.Each expert uses its own gated FFN and down-projection, and the layer output is the router-weighted sum of active-expert contributions.
  • C Implementation details: For implementation, GPT4ALL-J follows prior hyperparameters, Llama-3-8B modifies layer 9 for ROME and layers {7, 8, 9, 10, 11} for MEMIT, and MoE methods modify layers {3, 4, 5, 6, 7}.The selected GPT4ALL-J layer is confirmed suitable in Section 7.1.

D Prefix length ablation · E Prefix templates

The paper ablates ROME’s prefix length using random prefixes and subject-containing prompt templates, finding that longer random prefixes improve Harmfulness while reducing Generalization and preserving Locality. It also documents the templates used to generate prefixes.

  • D Prefix length ablation: ROME is evaluated with three random-prefix settings: 5/10 tokens, 30/50 tokens, and 50/100 tokens.The original setting averages 10 prefixes of length 5 with 10 prefixes of length 10; the other settings increase both lengths.
  • D Prefix length ablation: Table 3 reports the random-prefix results for ROME applied to GPT4ALL-J on the structural category.The table compares random prefixes with the authors’ prefixes of different lengths.
  • D Prefix length ablation: Increasing random prefix length improves Harmfulness and proportionally decreases Generalization, while Locality remains stable.The lower part of Table 3 presents the corresponding ablation results for subject-containing prompts.
  • D Prefix length ablation: Subject-containing prompt ablations use lengths of 10, 50, and 100, with five prefixes for each prompt template.These settings were chosen to match the random-prefix configurations.
  • E Prefix templates: The prefix templates used to generate x_j in equation (1) are documented in Appendix E.The section introduces the templates used for prefix generation.
  • E Prefix templates: One template is “Continue sentence {subject}.”The placeholder {subject} is retained exactly as presented.
  • E Prefix templates: Another template is “What do you think about {subject}?”.The template is listed among those used for prefix generation.
  • E Prefix templates: A third template is “What is the context about {subject}?”.The template is listed among those used for prefix generation.
Loading 2608.17836v1…