CAKE: Causal-Guided Adaptive Knowledge Editing for LLMs

🌟 Overview

• CAKE proposes a new Locate–Weight–Assign–Edit (LWAE) paradigm for multi-layer knowledge editing. LWAE replaces binary layer selection with continuous, causal-guided layer weights, and uses these weights to adaptively assign the editing residual across layers rather than distributing edits uniformly.
• This paradigm addresses a key failure mode of prior Locate-and-Edit style methods: heuristic layer selection and uniform residual allocation can misplace editing burden, amplifying cross-layer interference and causing error accumulation under sequential edits. By allocating edits in proportion to each layer’s causal importance, CAKE improves edit reliability while better preserving unrelated knowledge.

🏗️ Requirements

At least one L40s 48G GPU.

• torch==2.6.0
• einops==0.8.1
• higher==0.2.1
• hydra-core==1.3.2
• transformers==4.51.3
• datasets==2.21.0
• matplotlib==3.10.3
• spacy==3.4.1
• scipy==1.15.2
• scikit-learn==1.6.1
• nltk==3.9.1

🚀 Quick Start

An example for editing Llama3 (8B) on counterfact dataset using CAKE

1. Edit Llama3 (8B) model

python3 -m experiments.evaluate     --alg_name=Cake     --model_name=meta-llama/Meta-Llama-3-8B-Instruct     --hparams_fname=Llama3-8B.json --ds_name=mcf --dataset_size_limit=2000    --num_edits=100 --downstream_eval_steps=5

This command runs an evaluation script for the CAKE algorithm using the Llama3-8b-instruct. Below are the explanations for each argument:

• --alg_name=Cake: Specifies the name of the algorithm being used, which is CAKE in this case.
• --model_name=meta-llama/Meta-Llama-3-8B-Instruct: Indicates the name of the model being evaluated, here it is Llama-3-8B-Instruct.
• --hparams_fname=Llama3-8B.json: Points to the JSON file containing hyperparameters specific to the Llama-3-8B-Instruct model.
• --ds_name=mcf: Specifies the dataset name, in this case, "mcf" (CounterFact).
• --dataset_size_limit=2000: Sets the total number of editing samples to 2000.
• --num_edits=100: Defines the batch size for each round of editing, meaning 100 edits will be performed in each batch.
• --downstream_eval_steps=5: Indicates that a test of general capabilities is conducted after every 5 rounds of editing.

Results from each run are stored at results/<method_name>/run_<run_id> in a specific format:

results/
|__ Cake/
    |__ run_<run_id>/
        |__ params.json
        |__ case_0.json
        |__ case_1.json
        |__ ...
        |__ case_2000.json

2. Summarize the results

To summarize the results, you can use experiments/summarize.py:

python summarize.py --dir_name=Cake --runs=run_<run1>,run_<run2>

🧪 Hyperparameters

Key hyperparameters in hparams/Cake/Llama3-8B.json:

• layers: List of layer indices to edit (e.g., [4, 5, 6, 7, 8])
• causal_scores: Dictionary of causal importance scores for each layer (from causal tracing)
• temperature: Controls the sharpness of weight distribution

Example configuration:

{
    "layers": [4, 5, 6, 7, 8],
    ......
    "causal_scores": {
        "0": 0.4812439084,
        "1": 0.4743820429,
        "2": 0.4656370878,
        "3": 0.4440660179,
        "4": 0.4335190654
    },
    "temperature": 0.1
}