EasyEdit项目中的Llama2模型PMET方法优化建议

背景介绍

EasyEdit是一个用于大型语言模型(LLM)编辑的开源框架，其中PMET(Prompt-based Model Editing Technique)是一种基于提示的模型编辑方法。在最近的使用过程中，发现当该方法应用于Llama2模型时存在一个需要优化的技术细节。

问题分析

在使用PMET方法编辑Llama2模型时，LlamaTokenizer会自动在目标文本前添加特殊标记""。这一行为源于LlamaTokenizer的设计特性，它会在编码文本时默认添加开始符号(BOS token)。

具体表现为：

1. 在计算zs时，target_new参数会被自动添加""标记
2. 这会导致后续处理过程中查询(query)末尾也被附加了额外的""标记
3. 最终影响PMET方法对zs的优化效果

技术影响

通过对比实验可以观察到：

• 使用原始PMET方法时，优化过程收敛较慢，概率提升效果不明显
• 而MEMIT方法由于正确处理了BOS token，优化效果更好，概率提升更显著

解决方案

借鉴MEMIT方法中的处理方式，建议在PMET方法中添加类似的token处理逻辑：

1. 在计算zs时检查并移除开头的特殊标记
2. 或者更优的方案是统一使用tokenizer.encode(xx, add_special_tokens=False)方法
3. 这样可以避免自动添加不必要的特殊标记

实现建议

具体实现上，可以在compute_zs.py文件中加入以下逻辑：

if target_ids[0] == tok.bos_token_id or target_ids[0] == tok.unk_token_id:
    target_ids = target_ids[1:]

或者更彻底地修改tokenizer的调用方式：

target_ids = tok.encode(request["target_new"], add_special_tokens=False, return_tensors="pt").to("cuda")

总结

这个优化虽然看似是一个小改动，但对于Llama2模型的编辑效果有显著影响。它体现了在使用不同模型架构时，需要特别注意其tokenizer的特殊行为，确保编辑方法的通用性和稳定性。对于EasyEdit框架的用户来说，了解这一细节有助于更好地使用PMET方法进行模型编辑。