LLM 读心术 (1): 通过 hook 机制记录 LLM 的“脑活动”

当大语言模型 (LLMs) 回答 "Sure, I can help you with that" 时, 它的内部机理是如何运作的? 如果能够深入理解这个问题, 不仅可以帮助我们更清楚地了解 LLM 的运行机制, 甚至可以让我们干预和调整模型行为, 使其更加可控. 这一点类似于 neuroscientists 梦寐以求的目标: 当试图探究大脑活动与行为之间的关联时, 他们使用各种 "探针" 来 "读取" 大脑, 如 single unit、array、fMRI、EEG、钙成像等. 然而, 这些手段在时间或空间分辨率方面都存在一定局限性, 难以实现全脑实时高分辨率成像. 而对于 LLM, 这些局限并不存在, 因为其整个运行机制对我们来说是 "透明" 的, 我们可以随时获取任何所需的 LLM"脑活动" 数据, 甚至修改它们 (而目前 DL 可解释工作方面的现状也间接说明, 读取了大脑的所有活动!= 理解大脑如何工作... 参见 Could a Neuroscientist Understand a Microprocessor? Anyway 这是后话, 我们以后再探讨 ).

在这里, 本系列博客介绍一些相关的 LLM"读心术" 基础知识和最近研究成果, 这些知识或将对准备从事 LLM 可解释性相关工作的人有所帮助. 我们将从读取模型的内部活动开始, 到对内部活动进行一系列分析和干预, 从而实现一定程度上“脑控”模型结束.

其实, 现在已经存在一些库, 比如 TransformerLens 可以用于此类工作了. 但是作为研究者, 我们仍然希望自己实现一个最小的读取机制. 这样一来可以方便后续的修改, 二来可以在此类库 支持的模型 之外进行研究.

而在本文中, 我们将主要关心基于 pytorch 实现的的 huggingface 的 transformers 库所提供的预训练模型. 这个系列的模型也是目前最广泛使用的模型, 因此具有一定的代表性.

加载模型

首先加载一个模型. 我们这里使用一个小模型进行概念验证即可:

from transformers import AutoTokenizer, GPTNeoXForCausalLM  # , AutoModelForCausalLM

model_name = "EleutherAI/pythia-70m-deduped"
revision = "step143000"
cache_dir = f"./data/pythia-70m-deduped/{revision}"
model = GPTNeoXForCausalLM.from_pretrained(
    model_name,
    revision=revision,
    cache_dir=cache_dir,
)

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    revision=revision,
    cache_dir=cache_dir,
)

打印一下模型结构

GPTNeoXForCausalLM((gpt_neox): GPTNeoXModel((embed_in): Embedding(50304, 512)
    (emb_dropout): Dropout(p=0.0, inplace=False)
    (layers): ModuleList((0-5): 6 x GPTNeoXLayer((input_layernorm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(post_attention_layernorm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)(post_attention_dropout): Dropout(p=0.0, inplace=False)
        (post_mlp_dropout): Dropout(p=0.0, inplace=False)
        (attention): GPTNeoXAttention((rotary_emb): GPTNeoXRotaryEmbedding()(query_key_value): Linear(in_features=512, out_features=1536, bias=True)
          (dense): Linear(in_features=512, out_features=512, bias=True)
          (attention_dropout): Dropout(p=0.0, inplace=False)
        )(mlp): GPTNeoXMLP((dense_h_to_4h): Linear(in_features=512, out_features=2048, bias=True)
          (dense_4h_to_h): Linear(in_features=2048, out_features=512, bias=True)
          (act): GELUActivation())))(final_layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True) )
  (embed_out): Linear(in_features=512, out_features=50304, bias=False)
)

可以看到, 这是一个 6 层中间层的 transformer 模型. 接着我们尝试往里面送一点输入看看结果. 其实一般来说, 我们通过 model(**inputs) 已经可以获得 hidden_states 的内容了. 但是这个内容只是每个 hidden layer 的整体输出结果, 如果想要知道模型更加具体的每个 layer 内部细节甚至干预它们, 则需要利用 pytorch 的 forward hook 机制对各个 layer 的 forward 方法进行 hook.

在这里注册之前, 首先我们得知道哪些层可以被注册. 于是我们先看看各层的名字:

def gather_names(module, module_names, name=""):
    if isinstance(module, nn.Module):
        if name:
            module._name = name
        module_names.append(module._name)
        for child_name, child in module.named_children():
            gather_names(child, module_names, f"{name}.{child_name}" if name else child_name
            )module_names = []
gather_names(model, module_names, model_name)
module_names

输出

['EleutherAI/pythia-70m-deduped',
 'EleutherAI/pythia-70m-deduped.gpt_neox',
 'EleutherAI/pythia-70m-deduped.gpt_neox.embed_in',
 'EleutherAI/pythia-70m-deduped.gpt_neox.emb_dropout',
 'EleutherAI/pythia-70m-deduped.gpt_neox.layers',
 'EleutherAI/pythia-70m-deduped.gpt_neox.layers.0',
 ...,
 'EleutherAI/pythia-70m-deduped.gpt_neox.final_layer_norm',
 'EleutherAI/pythia-70m-deduped.embed_out']

"植入" hook

上面将该模型下所有的 nn.Module 和对应的名字记录在了 module_names 中, 接下来可以利用名字有选择性地对要记录的 layer 进行注册钩子:

activities = defaultdict(list) # 用于记录的字典


def forward_hook(module, input, output):
    activities[module._name].append(output)
    # print(f"Inside {module._name} forward")# print(f"Input: {input}")# print(f"Output: {output}")names2reg = ["EleutherAI/pythia-70m-deduped.gpt_neox.layers.5.mlp.dense_4h_to_h"] # 这是我们要记录的 layer


def register_hooks(module, all_hooks):
    if isinstance(module, nn.Module):
        for child_name, child in module.named_children():
            register_hooks(child, all_hooks)
        if any([part in module._name for part in names2reg]):# 这里后面可以换成其他的判别标准, 可以参考 peft 的参数传递方式.
            print(f"registering {module._name}")all_hooks.append(module.register_forward_hook(forward_hook))


hooks = []

register_hooks(model, hooks)

记录活动

# 在运行前先清空活动
for k in activities.keys():
    activities[k] = list()inputs = tokenizer("Hello, I am your", return_tensors="pt")
print(inputs["input_ids"].shape)tokens = model.generate(**inputs)
tokenizer.decode(tokens[0])
print(len(tokens[0]))

输出

torch.Size([1, 6])
20

接着检查一下记录到的活动:

layer = "EleutherAI/pythia-70m-deduped.gpt_neox.layers.5.mlp.dense_4h_to_h"
print(len(activities[layer])) # 14

在上面的例子中, 一开始我们送入 6 个 token, 总共输出 20 个 token, 那么 forward 的次数就是 14, 得到的 activities 长度也验证了这一点. 接着我们对 activities 进一步处理:

activities[layer] = torch.cat(activities[layer], dim=1) # 在时间轴上合并

# 这里我们假设所有的 batch 内都是同一个类型的文本需要分析,  所以进一步地, 我们在 batch dim 上进行平均
activities[layer] = activities[layer].mean(dim=0)

接着我们对这部分的“脑活动”进行一个粗糙的可视化:

fig, ax = plt.subplots(figsize=[10,2])
ax.imshow(activities[layer])
ax.set_xlabel('neurons')
ax.set_ylabel('time steps')
fig.tight_layout()plt.show()

得到

不难看出, 绝大多数的 neuron 在这个句子里还是比较“懒惰”的.

想法

上面为一个非常粗糙的 LLM 内部活动读取过程. 在这个系列的后面, 我们将把这个方法发展的好用一点, 使其能够:

1. 自定义 hook 不同 layer 的判别标准
2. 支持 neuroscience 里类似的跨 trial 的数据分析
3. 支持修改 neuron 的 activation 操作