Zenn
📗

LLMの重みでなんやかんや試してみる-①埋め込み層編-

2024/10/21に公開
機械学習
transformer
LLM
tech

このシリーズの大目的

このシリーズは、LLM(主にDecoder-Onlyで重みが公開されているもの)の内部をいじって遊んでみることでTransformerの新たな一面を見つけることができたらいいなという、ふわっとしたゴールに向かって進んでいきます。

今回のターゲット

この記事では、まずモデルへの入力が最初に出会う「埋め込み層」を触っていこうと思います。

埋め込み層(Embedding Layer)には、Tokenizerによって文字列がtoken idの列に変換されたものが入力として入ってきます。

# こんなイメージ
input_ids = [0, 5, 9, 43, 10]

そしてinput_idsの要素ごとに、対応した埋め込み層の重みベクトルを取ってきて、idの列から埋め込み表現の列に変換します。

# こんなイメージ
input_embeds = []
for id in input_ids:
    input_embeds.append(embed_weight[id])
# len(embed_weight[i])はモデルの埋め込み次元(=hidden_size)

そしてこの埋め込み層の重みは学習可能なパラメータで、事前学習やファインチューニングによって調整されます。

似たトークンは近い表現になってるの？

では、十分に学習されたモデルの埋め込み層は、やはり意味的に近い言葉(トークン)は近い表現になっているのでしょうか。
ここでは近い表現=コサイン類似度が大きいと定義して試してみます。

使用したコード
import torch
import numpy as np
from transformers import AutoTokenizer
from safetensors import safe_open
from numpy.linalg import norm

safetensor_file = "../models/llm-jp-3-1.8b/model.safetensors"
tokenizer_name = "llm-jp/llm-jp-3-1.8b"

tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

with safe_open(safetensor_file, framework="pt") as f:
    embed_tokens_weight = f.get_tensor("model.embed_tokens.weight")

if embed_tokens_weight.dtype == torch.bfloat16:
    embed_tokens_weight = embed_tokens_weight.to(torch.float32)

embed_tokens_weight = embed_tokens_weight.cpu().numpy()

def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (norm(vec1) * norm(vec2))

selected_id = 51484
selected_embedding = embed_tokens_weight[selected_id]

similarity_scores = []
for i, embedding in enumerate(embed_tokens_weight):
    if i != selected_id:
        similarity = cosine_similarity(selected_embedding, embedding)
        similarity_scores.append((i, similarity))

similarity_scores = sorted(similarity_scores, key=lambda x: x[1], reverse=True)

selected_token = tokenizer.decode([selected_id])
top_5 = similarity_scores[:5]
temp = []
for i in top_5:
    temp.append((i[0], i[1]))

print(f"Selected token (ID: {selected_id}): {selected_token}")
print("Top 5 most similar tokens:")
for token, similarity in temp:
    print(f"{token}: {similarity:.3f}")

1. llm-jp/llm-jp-3-1.8b

https://huggingface.co/llm-jp/llm-jp-3-1.8b

  • Selected token (ID: 51484): 猫

    1. 猫: 0.576
    2. 犬: 0.284
    3. 犬: 0.244
    4. ネコ: 0.242
    5. ネコ: 0.236

  • Selected token (ID: 39763): こたつ

    1. こたつ: 0.627
    2. コタツ: 0.477
    3. コタツ: 0.412
    4. 炬: 0.295
    5. ストーブ: 0.175

  • Selected token (ID: 39801): こんにちは

    1. こんにちは: 0.460
    2. こんばんは: 0.449
    3. こんにちは。: 0.383
    4. こんにちは！: 0.343
    5. こんにちは、: 0.335

初めに、先月LLM-jpより公開されたllm-jp/llm-jp-3-1.8bです。
Tokenizerの語彙の中に重複する語彙があるのか、decodeしたときに同じものが複数でてきました。
とはいえ、「猫」に対して「ネコ」が5位圏内に入っていたり、「こたつ」に対して「ストーブ」が近かったりと、似た言葉を近い表現で表せていそうです。
また、以下で他のモデルも試していますが、近いものは近い、それ以外は遠いといったように比較的はっきりと分かれている印象です。

2. cyberagent/calm2-7b

https://huggingface.co/cyberagent/calm2-7b

  • Selected token (ID: 7827): 猫

    1. の猫: 0.233
    2. 舌: 0.158
    3. ホ: 0.156
    4. ご: 0.149
    5. 新潟: 0.148

  • Selected token (ID: 41221): ストーブ

    1. コンロ: 0.194
    2. 豆腐: 0.191
    3. シロップ: 0.184
    4. 資材: 0.177
    5. テーブル: 0.174

  • Selected token (ID: 8569): こんにちは

    1. こんばんは: 0.392
    2. おはよう: 0.334
    3. 新型コロナ: 0.203
    4. どうも: 0.196
    5. 〒: 0.191

次に昨年サイバーエージェントより公開されたcyberagent/calm2-7bです。
やはり同様に、似た言葉が近い表現として学習されているように見えます。
こちらのモデルのTokenizerでは「こたつ」が1トークンでは表せなかったため、代わりに「ストーブ」という言葉にしましたが、「コンロ」が最も近く納得できます。(「豆腐」は「トーフ」ってことなんでしょうか...？)
気になる点を挙げるとするならばllm-jp-3-1.8bに比べて類似度のばらつきが小さいかな？という感じです。

3. YukiTomita-CC/AKU-d_ms-0.5B-chat-v0.1

https://huggingface.co/YukiTomita-CC/AKU-d_ms-0.5B-chat-v0.1

  • Selected token (ID: 2673): 猫

    1. ネコ: 0.527
    2. 犬: 0.519
    3. ねこ: 0.455
    4. 狐: 0.348
    5. 動物: 0.347

  • Selected token (ID: 1621): 温泉

    1. 湯: 0.455
    2. 風呂: 0.372
    3. リゾート: 0.371
    4. 入浴: 0.369
    5. ホテル: 0.360

  • Selected token (ID: 29343): こんにちは

    1. おはよう: 0.353
    2. さよなら: 0.323
    3. ありがとう: 0.311
    4. さらば: 0.285
    5. ください: 0.276

最後は先日こちらの記事で紹介した、0.5BサイズのMistralスクラッチモデル(自作)です。
正直性能は良いとは言えない結果だったため、表現を獲得できてない例として引き合いに出すつもりだったのですが、意外や意外、ちゃんと似た言葉を近い表現で表せています。
(「こたつ」も「ストーブ」も1トークンではなかったので「温泉」という言葉にしています)

しかし、冷静になって考えると「温泉」の例のように、同じコンテキストで使われるが異なる言葉の表現が近いというのは良くないのかもしれません。
例えば人間が違う意味で使った言葉をモデルは同じ意味として認識してしまうわけですから、Attentionがうまく機能しないとモデルが困惑する原因にもなりそうです。
その意味で言うとllm-jp-3-1.8bのように意味を確立させた方がいいのかも...？

他言語の近さは？

日本語同士の近さを見てきましたが、他の言語との意味の近さはどうなのでしょうか。
いくつかトークンを試していたときに「言語」という言葉がこのトピックによさそうだったのでピックアップします。
どのモデルも「言語」と「language」という言葉は近い言葉だと学習しているようです。
少し逸れますが、llm-jp-3-1.8b以外は「プログラミング言語」も近い言葉だと認識しているみたいですね。確かに文脈によっては『A: 得意な言語何？ B: Pythonかな。』という風に言うのも自然ですから妥当だと思います。

1. llm-jp/llm-jp-3-1.8b

  • Selected token (ID: 53240): 言語
    1. 言語: 0.513
    2. 언어: 0.226 (言語という意味だそうです)
    3. 多言語: 0.201
    4. languages: 0.190
    5. language: 0.179

2. cyberagent/calm2-7b

  • Selected token (ID: 7271): 言語
    1. の言語: 0.228
    2. ミング言語: 0.211
    3. 中国語: 0.183
    4. _languages: 0.167 (_は半角スペース)
    5. 人間: 0.148

3. YukiTomita-CC/AKU-d_ms-0.5B-chat-v0.1

  • Selected token (ID: 1896): 言語
    1. プログラミング言語: 0.445
    2. Language: 0.412
    3. 公用語: 0.395
    4. 語: 0.382
    5. 諸語: 0.373

近い表現なら置き換えても似た出力になるの？

以上の検証の中で、最も類似度の高かったのは私のモデルの「猫」と「ネコ」でした。
類似度が高いといっても0.527なのでそこまでは高くないのですが、このトークンだけを置き換えた場合、出力はどう変化するのでしょうか。
以下、(入力) -> (出力)の形で試した結果です。

猫は好きですか? -> 猫も好きですね。猫も可愛いですよね。

ネコは好きですか? -> はい、好きですね。あなたは?

好きであることは同じですが、これは私が学習データに猫が嫌いなことは含めておらず(確認した限り)、むしろ好きというデータを複数入れた記憶があるのでそのせいだと思います。
埋め込み表現が近いと出力が似るのかについてはもう少し調べてみたいです。

まとめ

埋め込み層の重みを持ってきて色々試してみました。

もちろんこれはTransformerのコンポーネントの1つですので、ここだけで何かがわかるというものではないですが、モデルは埋め込み表現をどんな風に獲得しているんだろうという疑問が少し晴れてよかったです。

次は少々重くなりそうですがAttentionブロックをいじってみたいですね。

ここまで読んでいただきありがとうございました！

Discussion