PMET Llama No Difference In Output

[uygarkurt]

Hi, thank you for this amazing work!

I'm trying to edit NousResearch/Llama-2-7b-chat-hf using PMET. I'm following the samples you provided. However it gives exactly the same answer for both the original and the edited model. I use the example which you change profession of Grant Hill from basketball player to football player.

When I ask Who is Grant Hill?, the original model provides the following answer:

Grant Hill is an American former professional basketball player and current sports analyst. He played in the National Basketball Association (NBA) for 19 seasons, primarily with the Detroit Pistons, Orlando Magic, and Phoenix Suns. Hill was a three-time NBA All-Star and was named the NBA's Sixth Man of the Year twice. After retiring from basketball, Hill became a sports commentator for ESPN and has appeared on various television shows, including NBA Countdown and SportsCenter.

The edited model gives the following answer:

Grant Hill is an American former professional basketball player and current sports analyst. He played in the National Basketball Association (NBA) for 19 seasons, primarily with the Detroit Pistons, Orlando Magic, and Phoenix Suns. Hill was a three-time NBA All-Star and was named the NBA's Sixth Man of the Year twice. After retiring from basketball, Hill became a sports commentator for ESPN and has appeared on various television shows, including NBA Countdown and SportsCenter.

I was able to edit GPT-2XL using this. I was also able to edit Llama using MEMIT. However, PMET with Llama didn't work. II'm attaching the codes I used.

edit.py

def test_PMET():

    prompts = ['Ray Charles, the',
               'Grant Hill is a professional',
               'The law in Ikaalinen declares the language'
               ]
    ground_truth = ['piano',
                    'basketball',
                    'Finnish'
                    ]
    target_new = ['violin',
                  'soccer',
                  'Swedish'
                  ]
    subject = ['Ray Charles',
               'Grant Hill',
               'Ikaalinen'
               ]

    locality_inputs = {
        'neighborhood':{
            'prompt': ['Joseph Fischhof, the', 'Larry Bird is a professional', 'In Forssa, they understand'],
            'ground_truth': ['piano', 'basketball', 'Finnish']
        },
        'distracting': {
            'prompt': ['Ray Charles, the violin Hauschka plays the instrument', 'Grant Hill is a professional soccer Magic Johnson is a professional', 'The law in Ikaalinen declares the language Swedish In Loviisa, the language spoken is'],
            'ground_truth': ['piano', 'basketball', 'Finnish']
        }
    }
    portability_inputs = {
        'synonym':{
            'prompt': ['Ray Charles, the', 'Grant Hill is a professional', 'The law in Ikalis declares the language'],
            'ground_truth': ['violin', 'soccer', 'Swedish']
        },
        'one_hop':{
            'prompt': ['Ray Charles, the', 'Grant Hill is a professional', 'The law in Ikalis declares the language'],
            'ground_truth': ['violin', 'soccer', 'Swedish']
        }
    }

    hparams = PMETHyperParams.from_hparams('./hparams/PMET/llama-7b')
    editor = BaseEditor.from_hparams(hparams)
    metrics, edited_model, _ = editor.edit(
        prompts=prompts,
        ground_truth=ground_truth,
        target_new=target_new,
        subject=subject,
        locality_inputs=locality_inputs,
        portability_inputs=portability_inputs,
        keep_original_weight=True
    )

    return metrics, edited_model

    _, model = test_PMET()
    model.generation_config.temperature = None
    model.generation_config.top_p = None
    print("RETURNED MODEL")
    model.save_pretrained("models/llama-7b-edited/")

llama-7b.yaml

alg_name: "PMET"
model_name: "NousResearch/Llama-2-7b-chat-hf"
stats_dir: "./data/stats"
device: 0
layers: [4, 5, 6, 7, 8]
clamp_norm_factor: 1
layer_selection: "all"
fact_token: "subject_last"
v_num_grad_steps: 20
v_lr: 5e-1
v_loss_layer: 31
v_weight_decay: 0.5
kl_factor: 0.0625
mom2_adjustment: true
mom2_update_weight: 15000
rewrite_module_tmp: "model.layers.{}.mlp.down_proj"
rewrite_module_tmps: ["model.layers.{}.mlp.down_proj"]
layer_module_tmp: "model.layers.{}"
mlp_module_tmp: "model.layers.{}.mlp.down_proj"
attn_module_tmp: "model.layers.{}.self_attn.o_proj"
ln_f_module: "model.norm"
lm_head_module: "lm_head"
mom2_dataset: "wikipedia"
mom2_n_samples: 100000
mom2_dtype: "float32"
model_parallel: false
nll_loss_factor: 1

At least I get the results with the following code

from transformers import pipeline, LlamaTokenizerFast, LlamaForCausalLM, AutoModelForCausalLM
import torch

tokenizer = LlamaTokenizerFast.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
prompts = [
           "Who is Grant Hill?",
           ]

for prompt in prompts:
    model_orig = AutoModelForCausalLM.from_pretrained("NousResearch/Llama-2-7b-chat-hf", local_files_only=True).to("cuda")
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model_orig.generate(**inputs.to("cuda"),
                                  max_new_tokens=512,
                                  num_beams=10,
                                  early_stopping=False,
                                  no_repeat_ngram_size=2,
                                  )

    res = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    print(f"Original Model Output: {res[0]}")

    del model_orig
    torch.cuda.empty_cache()
    model_edited = AutoModelForCausalLM.from_pretrained("models/llama-7b-edited", local_files_only=True).to("cuda")
    outputs = model_edited.generate(**inputs.to("cuda"),
                                    max_new_tokens=512,
                                    num_beams=10,
                                    early_stopping=False,
                                    no_repeat_ngram_size=2,
                                    )

    res = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    print(f"Edited Model Output: {res[0]}")

    del model_edited
    torch.cuda.empty_cache()

[XeeKee]

Thank you for your attention. We are currently busy with the NeurIPS submission, we will address this matter by the end of the month.

[pengzju]

The issue you are concerned with is actually the generalization of editing, from simple editing prompts to the generation of complex queries. This is a problem that many current works have not yet resolved (including the PMET you are using). For more details, you can refer to this paper: https://arxiv.org/abs/2402.09394.