XPhoneBERTの埋め込みを試す

はじめに

音韻特徴を捉えた埋め込み（「みず-ゆ」よりは「みず-きず」のほうが似ている、みたいな）を得る方法を試行錯誤しています。
今回はBERTを音素列で学習したXPhoneBERTを試してみます。

XPhoneBERT

XPhoneBERTはTTS向けに音素列で学習されたBERTです。
モデルはhuggingfaceから利用することができます。

TTS用途ということである程度は音韻の特徴を学習できていそうなので、試してみます。ただ、学習データは「330M phoneme-level sentences from nearly 100 languages and locales」であり、sentenceのみで音声波形を使っていなさそうなので、それでどのくらい性能がでるかは試してみないとわかりません。

検証

まずはhuggingfaceのコードを実行してみます。
単語分割済みの入力文を音素列に変換し、さらにそれを音素のID列に変換したものをモデルにつっこむと、特徴ベクトルが得られるようです。

from transformers import AutoModel, AutoTokenizer
from text2phonemesequence import Text2PhonemeSequence
import torch
# Load XPhoneBERT model and its tokenizer
xphonebert = AutoModel.from_pretrained("vinai/xphonebert-base")
tokenizer = AutoTokenizer.from_pretrained("vinai/xphonebert-base", add_prefix_space=True)

# Load Text2PhonemeSequence
# text2phone_model = Text2PhonemeSequence(language='eng-us', is_cuda=True)
text2phone_model = Text2PhonemeSequence(language='jpn', is_cuda=False)

# Input sequence that is already WORD-SEGMENTED (and text-normalized if applicable)
# sentence = "That is , it is a testing text ."  
sentence = "これ は 、 テスト テキスト です ."

input_phonemes = text2phone_model.infer_sentence(sentence)

input_ids = tokenizer(input_phonemes, return_tensors="pt")

with torch.no_grad():
    features = xphonebert(**input_ids)

print(features)

BaseModelOutputWithPoolingAndCrossAttentions(last_hidden_state=tensor([[[ 0.0635, -0.1972, -0.0265,  ...,  0.0503,  0.6004, -0.1093],
         [-0.3267, -0.0715, -0.2766,  ..., -0.3860,  0.1892,  0.4146],
         [-0.3546,  0.0751, -0.1748,  ..., -0.4973,  0.5676, -0.1788],
         ...,
         [-0.0703, -0.1634, -0.1806,  ...,  0.0655,  0.0962, -0.2656],
         [-0.3719, -0.4315,  0.1166,  ..., -0.3314,  0.1720, -0.5358],
         [ 0.0131, -0.1108, -0.1122,  ...,  0.0747,  0.3016, -0.2153]]]), pooler_output=tensor([[ 2.5155e-02, -3.6114e-03,  3.9136e-02, -1.2485e-01, -1.4914e-01,
         ...
          1.4286e-01,  6.0723e-03, -2.2473e-01]]), hidden_states=None, past_key_values=None, attentions=None, cross_attentions=None)

featuresにはlast_hidden_stateとpooler_outputの２つのテンソルが含まれます。
last_hidden_stateは各音素トークンに対応する最終層のテンソル、pooler_outputはおそらく[CLS]に相当するテンソルです。
系列全体のテンソルとしては、last_hidden_stateの平均値か、pooler_outputのどちらかを使えば良さそうです。
とりあえずpooler_outputで試してみます。

単語分割済みのsentenceを入力としてpooler_outputを返す処理を関数化します。

def get_feature_vector(sentence):
    input_phonemes = text2phone_model.infer_sentence(sentence)
    input_ids = tokenizer(input_phonemes, return_tensors="pt")
    with torch.no_grad():
        features = xphonebert(**input_ids)
    return features.pooler_output

# 例として、関数を使って特徴ベクトルを取得
example_sentence = "これ は 、 テスト テキスト です ."
feature_vector = get_feature_vector(example_sentence)

print(feature_vector)

tensor([[ 2.5155e-02, -3.6114e-03,  3.9136e-02, -1.2485e-01, -1.4914e-01,
          2.4427e-01, -9.8719e-02,  1.2688e-01, -8.5747e-02,  8.8691e-03,
          1.0907e-01, -7.9545e-02, -8.9626e-02,  9.1621e-02, -3.0248e-02,
         -2.1478e-01,  7.4607e-02, -2.7360e-02,  9.9830e-03,  2.1486e-02,
         -1.1748e-01,  4.0039e-02, -6.2457e-02,  7.9601e-02,  1.6807e-02,
          3.1896e-02,  2.1359e-01,  9.9437e-02, -3.3402e-02, -1.1814e-03,
          ...
          1.4286e-01,  6.0723e-03, -2.2473e-01]])

いい感じです。

「みず」「きす」「おゆ」の３つの単語の類似度を計算してみます。
日本人的な感覚としては母音が重視されてほしいので、「みず-きす」の類似度が高くなってほしいです。

from sklearn.metrics.pairwise import cosine_similarity

word1 = "きす"
word2 = "みず"
word3 = "おゆ"

embeddings1 = get_feature_vector(word1)
embeddings2 = get_feature_vector(word2)
embeddings3 = get_feature_vector(word3)

cos_sim12 = cosine_similarity(embeddings1.cpu().numpy(), embeddings2.cpu().numpy())
cos_sim13 = cosine_similarity(embeddings1.cpu().numpy(), embeddings3.cpu().numpy())
cos_sim23 = cosine_similarity(embeddings2.cpu().numpy(), embeddings3.cpu().numpy())

print(f"Cosine Similarity between '{word1}' and '{word2}': {cos_sim12[0][0]}")
print(f"Cosine Similarity between '{word1}' and '{word3}': {cos_sim13[0][0]}")
print(f"Cosine Similarity between '{word2}' and '{word3}': {cos_sim23[0][0]}")

Cosine Similarity between 'きす' and 'みず': 0.9219405651092529
Cosine Similarity between 'きす' and 'おゆ': 0.9367311000823975
Cosine Similarity between 'みず' and 'おゆ': 0.9013646841049194

「きす-おゆ」の類似度が最も高いという結果でした。
あまり感覚とは一致しません。

あ段のかなの類似度を計算してみます。

# ひらがなのあ段のペアの類似度を降順にソートして出力

# あ段のひらがなリスト
a_dan_hiragana = "あ か さ た な は ま や ら わ が ざ だ ば ぱ".split(" ")

# ペアのリストを作成
pairs = [(a_dan_hiragana[i], a_dan_hiragana[j]) for i in range(len(a_dan_hiragana)) for j in range(i + 1, len(a_dan_hiragana))]

# 類似度を計算して保存するリスト
similarities = []

for word1, word2 in pairs:
    # 音素列に変換
    embeddings1 = get_feature_vector(word1)
    embeddings2 = get_feature_vector(word2)
    # コサイン類似度を計算ｒ
    cos_sim = cosine_similarity(embeddings1.cpu().numpy(), embeddings2.cpu().numpy())
    similarities.append((word1, word2, cos_sim[0][0]))

# 類似度を降順にソート
sorted_similarities = sorted(similarities, key=lambda x: x[2], reverse=True)

# 結果を出力
for word1, word2, sim in sorted_similarities:
    print(f"Cosine Similarity between '{word1}' and '{word2}': {sim}")

Cosine Similarity between 'た' and 'ば': 0.9860794544219971
Cosine Similarity between 'は' and 'ら': 0.9813922047615051
Cosine Similarity between 'た' and 'ぱ': 0.9802631139755249
Cosine Similarity between 'か' and 'ぱ': 0.9800434112548828
Cosine Similarity between 'か' and 'や': 0.9789953827857971
Cosine Similarity between 'ば' and 'ぱ': 0.9787014722824097
Cosine Similarity between 'か' and 'ら': 0.978649377822876
Cosine Similarity between 'や' and 'ら': 0.9785498380661011
Cosine Similarity between 'か' and 'ば': 0.9784786105155945
Cosine Similarity between 'か' and 'た': 0.9766682386398315
...

最も類似度が高いペアは「た-ば」でした。よく「ば」と「ま」は互いに変換されやすく類似した音であると言われたりしますが、そういった一般的に類似するといわれる子音のペアはそこまで上位にはきていなさそうです。

pooler_outputではなくlast_hidden_stateのmean_poolingでも実験してみましたが、似たような結果で、あまり感覚とは一致しませんでした。

おわりに

XPhoneBERTの埋め込みにもとづく類似度をいくつかのサンプルで眺めてみましたが、あまり直感とは合わない結果でした。
なにかモデルに対する勘違いや実装ミスがある可能性もありますが、一旦このくらいで検証を終わろうと思います。