GPT-OSSの埋め込み

埋め込みの続編。gpt-oss-20bの埋め込み空間で king - man + woman をやってみよう。

モデルは mlx-community/gpt-oss-20b-MXFP4-Q8 を使う。ただし、このモデルは推論用には微妙である（そもそもQ8なのにQ4とほとんどサイズが変わらない）。

model, tokenizer = load("mlx-community/gpt-oss-20b-MXFP4-Q8")

prompt = "queen"
token_ids = tokenizer.encode(prompt)  # [153556]
tokens = tokenizer.convert_ids_to_tokens(token_ids)  # ['queen']

トークン列を埋め込みベクトルに変換するための関数を探すために、Pythonに model と打ち込んでみる：

Model(
  (model): GptOssMoeModel(
    (embed_tokens): QuantizedEmbedding(201088, 2880, group_size=64, bits=8, mode=affine)
    (norm): RMSNorm(2880, eps=1e-05)
...以下略...

これから model.model.embed_tokens だと推量する：

embeddings = model.model.embed_tokens(token_ids)
queen = embeddings[0]  # 2880次元のベクトル
mx.linalg.norm(queen)  # array(123, dtype=bfloat16)

同じことを man、woman、king についても行う。

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

X = np.vstack([man.astype(mx.float32),
               woman.astype(mx.float32),
               king.astype(mx.float32),
               queen.astype(mx.float32),
               (king - man + woman).astype(mx.float32)])
cosine_similarity(X)

結果は次の通り：

array([[ 0.9999998 ,  0.43785974,  0.5198675 ,  0.35447446, -0.11281953],
       [ 0.43785974,  0.9999999 ,  0.19880976,  0.4180091 ,  0.6151444 ],
       [ 0.5198675 ,  0.19880976,  1.        ,  0.33824253,  0.4844749 ],
       [ 0.35447446,  0.4180091 ,  0.33824253,  0.99999976,  0.30324462],
       [-0.11281953,  0.6151444 ,  0.4844749 ,  0.30324462,  1.0000004 ]],
      dtype=float32)

king - man + woman は queen より king に近く、それよりさらに woman に近い。