大規模言語モデルの「脳波」をとって言葉を生成しているときにどこが活動しているのか見てみる

ChatGPTなんかの大規模言語モデルが言葉を生成しているときに、どういう反応が起きているのか気になりますよね。きっと気になる。
ということで、手元で動かせる言語モデルのニューラルネット各層での出力を表示してみました。

GPTにはGPTブロックが複数あって、それぞれのブロックが２層のニューラルネットレイヤを持っています。

モデルを読み込んだあとでこういうコードを動かしてニューラルネットの出力の二乗を足していきます。

for idx, elm in enumerate(model.transformer.h):
    elm.ln_1.index = idx * 2
    elm.ln_2.index = idx * 2 + 1
    elm.ln_1.old_forward = elm.ln_1.forward
    elm.ln_2.old_forward = elm.ln_2.forward
    def new_forward(self, x):
        result = self.old_forward(x)
        ar = result.detach().numpy()
        summary[self.index] += ar[0][0] ** 2
        return result
    elm.ln_1.forward = new_forward.__get__(elm.ln_1)
    elm.ln_2.forward = new_forward.__get__(elm.ln_2)

    tensor1 = elm.ln_1.weight.data
    tensor2 = elm.ln_2.weight.data
    summary.append(np.zeros(len(tensor1), dtype=float))
    summary.append(np.zeros(len(tensor2), dtype=float))

あとはテキスト生成したあとでグラフ表示

prompt = input("prompt: ")
if prompt == "exit":
    break
generated_text = pipe(prompt, do_sample=True, use_cache=True, **opts)[0]   
print(Fore.YELLOW + prompt + Fore.WHITE + generated_text['generated_text'][len(prompt):])     
plt.imshow(np.clip(np.sqrt(summary), 0, range_max), cmap=cmap, aspect='auto')
plt.colorbar()
plt.show()