Wav2vec2 CTCモデルのトークンの表現が音韻特徴を捉えているか眺めてみる

はじめに

意味ではなく音韻的な特徴を反映した単語のembeddingを作りたいです。
例えば「いね-むぎ」より「いぬ-いね」のほうがcos類似度が高くなるようなembeddingがほしいです。

Phonetic Word Embedding (Sharma 2021)という論文を読みました。結構いい感じの結果に見えたのですが、ベースとなる類似度のルールやパラメータは人間が与えていたので、もうちょっと、音声波形と音素の対応みたいなものからembeddingを作れたほうがいいのではないかと思いました。

素のwav2vec2は音声波形のembeddingを学習していますが、音素や単語の情報は持っていません。wav2vec2のCTCモデルであれば、素のwav2vec2に対して、波形と書き起こしを対応させたfine-tuningを行っている（ざっくりとしか理解していません）ので、文字トークンの表現をモデルに内包していることが期待されます。

ChatGPTにきいてみたところ、CTCモデルの最終層はトークンごとの出力スコアを計算するための重み行列なので、その各行はトークンの表現になっている可能性があるとのことでした。

そこで、wav2vec2のCTCモデルの最終層からトークンに対応するベクトルを取得して、トークン間の類似度が感覚とあうか眺めてみます。

実装

日本語で学習されたwav2vec2のCTCモデルであるreazon-research/japanese-wav2vec2-base-rs35khを使ってみます。

準備

まず仮想環境を作って必要なライブラリをインストールします。

mkdir phonetic-embedding
cd phonetic-embedding
uv init
uv add transformers ipykernel scikit-learn sudachipy sudachidict-core
uv add "numpy==1.*" # torchの依存関係でv1しか使えないため
uv add torch==2.2 # uvだと最新のtorchをaddできないため

以降の作業はjupyter notebook上で実施しています。

モデルとトーカナイザをダウンロードします。

from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer

model_name = "reazon-research/japanese-wav2vec2-base-rs35kh"
model = Wav2Vec2ForCTC.from_pretrained(model_name)
tokenizer = Wav2Vec2Tokenizer.from_pretrained(model_name)

最終層の重みを取得します。

# 最終線形層の重みを取得
embedding_matrix = model.lm_head.weight.detach().cpu().numpy()
# embedding_matrix.shape -> (vocab_size, hidden_size)

以下のように語彙に対応するembeddingを取得できます。

vocab = tokenizer.get_vocab()  # {token_str: token_id}のdict

for word, id in vocab.items():
    print(f"{word}: {id}")
    print(embedding_matrix[id])
    break

!: 22
[-3.95605564e-02  2.10904609e-02 -1.12539329e-01  1.37193069e-01
 -4.24791165e-02 -2.15949565e-02 -3.85517180e-02 -3.01879961e-02
  6.02556244e-02  2.45358348e-02 -4.71599884e-02 -3.43878902e-02
 -3.73844765e-02 -5.06201945e-02 -1.55719995e-01  1.90920517e-01
  1.05109043e-01  1.17944255e-01 -5.48516214e-02  3.88123654e-02
 -3.85209471e-02  8.99153762e-04  9.86133888e-02  6.18048683e-02
  1.64006185e-02  1.06509164e-01  3.13028432e-02 -3.38351540e-02
...

cos類似度

トークンのembedding同士のcos類似度を計算し、類似度の高いペアを眺めてみます。

上位の類似ペア

まずはtop10くらいのペアを眺めます。

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

# コサイン類似度行列を計算
cos_sim_matrix = cosine_similarity(embedding_matrix)

# 類似度の高いペアtop100を取得
similar_pairs = []
vocab_size = embedding_matrix.shape[0]

for word, id in vocab.items():
    for word2, id2 in vocab.items():
        if id < id2:
            similar_pairs.append((word, id, word2,id2, cos_sim_matrix[id, id2]))

# 類似度の降順にソートしてtop100を取得
similar_pairs.sort(key=lambda x: x[-1], reverse=True)

# 結果を出力
for word, id, word2, id2, cos_sim in similar_pairs[:10]:
    print(f"{word} and {word2}: {cos_sim}")

▁その and ▁今: 0.9703211784362793
▁これ and ▁はい: 0.96944260597229
▁その and ▁あ: 0.9693166017532349
▁はい and ▁あ: 0.9692965745925903
▁その and ▁お: 0.969156801700592
▁その and ▁そして: 0.9690490961074829
▁この and ▁その: 0.9689103364944458
▁その and ▁でも: 0.9686933755874634
▁そして and ▁そう: 0.968663215637207
▁これ and ▁今: 0.9686161279678345

あまり音韻が近い感じがしません。
先頭にアンダースコアがついているのはおそらくサブワードです。
サブワードだから結果が変なのかもしれません。
サブワードをフィルタして上位のペアを眺めてみます。

# 類似度の降順にソートしてtop100を取得
similar_pairs.sort(key=lambda x: x[-1], reverse=True)

# 結果を出力
count = 0
for word, id, word2, id2, cos_sim in similar_pairs:
    if count > 20:
        break
    if not word.startswith("▁") and not word2.startswith("▁"):
        print(f"{word} and {word2}: {cos_sim}")
        count += 1

</s> and <pad>: 0.965054452419281
<s> and </s>: 0.9638963341712952
<s> and <pad>: 0.9630551338195801
ほか and 他: 0.8556908369064331
というのは and っていうのは: 0.8166471719741821
んですけど and んですけれども: 0.7813968062400818
更 and さらに: 0.7781423926353455
んですね and んですよね: 0.7694100141525269
今日 and きょう: 0.7511346340179443
している and してる: 0.746308445930481
っている and ってる: 0.7452219724655151
んですけど and んですが: 0.7213988304138184
見て and を見て: 0.7186458706855774
んですが and んですけれども: 0.7181704044342041
ですよね and んですよね: 0.7167232632637024
つけ and 付け: 0.715129017829895
んですよ and んですよね: 0.7140146493911743
んです and んですよ: 0.709439218044281
になります and になりました: 0.7088898420333862
けど and けれども: 0.7087136507034302
何か and なんか: 0.7048472166061401

サブワードに比べると、ある程度音韻が近いトークンペアになっているような気はします。
ただ、純粋な発音の近さだけでなく意味的な近さも影響していそうです（になります-になりました、けど-けれども）など。

名詞の類似ペア

上記は助詞や接続詞的なもの多いので、名詞でも試してみます。

「犬」「稲」「麦」でみてみます。
意味的には稲と麦が近いですが音韻的には犬と稲のほうが近いと期待しています。

# 「犬」と「麦」と「稲」の類似度を計算して出力
words_to_compare = ["犬", "麦", "稲"]
word_ids = [vocab[word] for word in words_to_compare]

for i in range(len(word_ids)):
    for j in range(i + 1, len(word_ids)):
        word1, word2 = words_to_compare[i], words_to_compare[j]
        id1, id2 = word_ids[i], word_ids[j]
        cos_sim = cos_sim_matrix[id1, id2]
        print(f"{word1} and {word2}: {cos_sim}")

犬 and 麦: 0.14278215169906616
犬 and 稲: 0.26119738817214966
麦 and 稲: 0.1056247130036354

犬と稲の類似度が一番高いので、これは良い結果に見えます。

別の単語群でも試してみます。
水、氷、湯、傷、鉄、です。水、傷、鉄の類似度が高くなることを期待しています。

vocabにこれらの語が存在するか確認します。語彙が3003個しかないため、意外とないことがあります。

words = ["水", "湯", "氷", "傷", "鉄"]
for word in words:
    if word in vocab:
        print(f"{word} is in vocab")
    else:
        print(f"{word} is not in vocab")

水 is in vocab
湯 is in vocab
氷 is in vocab
傷 is in vocab
鉄 is in vocab

OKです。
類似度を見てみます。

# 「犬」と「麦」と「稲」の類似度を計算して出力
words_to_compare = ["水", "湯", "氷", "傷", "鉄"]
word_ids = [vocab[word] for word in words_to_compare]

for i in range(len(word_ids)):
    for j in range(i + 1, len(word_ids)):
        word1, word2 = words_to_compare[i], words_to_compare[j]
        id1, id2 = word_ids[i], word_ids[j]
        cos_sim = cos_sim_matrix[id1, id2]
        print(f"{word1} and {word2}: {cos_sim}")

水 and 湯: 0.12245313823223114
水 and 氷: 0.11520753055810928
水 and 傷: 0.17717738449573517
水 and 鉄: 0.16831175982952118
湯 and 氷: 0.15924611687660217
湯 and 傷: 0.11892639100551605
湯 and 鉄: 0.14361493289470673
氷 and 傷: 0.08718910068273544
氷 and 鉄: 0.09668049216270447
傷 and 鉄: 0.15622802078723907

水-傷、水-鉄、傷-鉄、がtop3です。感覚と一致します。
名詞に限定して、上位の類似度ペアを見てみます。
まずtokenが名詞かどうか判定する関数を作ります。

from sudachipy import tokenizer
from sudachipy import dictionary

tokenizer_obj = dictionary.Dictionary(dict_type="full").create()

# SplitMode.Cはなるべく長い単語に分割するモード
mode = tokenizer.Tokenizer.SplitMode.C

def is_all_nouns(text):
    # 辞書オブジェクトを取得
    tokens = tokenizer_obj.tokenize(text, mode)
    
    # すべてのトークンが名詞かどうかをチェック
    for token in tokens:
        # 品詞情報を取得
        pos = token.part_of_speech()
        # 名詞でない場合はFalseを返す
        if pos[0] != "名詞":
            return False
    return True

名詞の上位ペアを出力します。

# 結果を出力
count = 0
for word, id, word2, id2, cos_sim in similar_pairs:
    if count > 20:
        break
    if not word.startswith("▁") and not word2.startswith("▁") and is_all_nouns(word) and is_all_nouns(word2):
        print(f"{word} and {word2}: {cos_sim}")
        count += 1

ほか and 他: 0.8556908369064331
今日 and きょう: 0.7511346340179443
齋 and 齊: 0.7019595503807068
沢 and 澤: 0.699118971824646
邊 and 邉: 0.6921684741973877
状況 and 状態: 0.6744210720062256
斎 and 齋: 0.6667966842651367
辺 and 邉: 0.645785927772522
島 and 嶋: 0.6441478729248047
桜 and 櫻: 0.6343605518341064
子 and 子ども: 0.6231734752655029
時 and とき: 0.6193947792053223
富 and 冨: 0.6037569642066956
ご覧 and 覧: 0.6016298532485962
兄 and 姉: 0.6015802621841431
女性 and 男性: 0.6007189750671387
辺 and 邊: 0.5988811254501343
謗 and 誹: 0.598186731338501
介 and 輔: 0.5960387587547302
祐 and 佑: 0.5811824798583984
斉 and 齋: 0.573908805847168

同じ読みをする漢字が結構出てきます。
「状態-状況」「女性-男性」など意味に引っ張られていそうなペアがあります。
「辺-邉」は名字に由来していそうです。
「謗-誹」が近いのはちょっと面白そうです。おそらくそれぞれが「誹謗」という単語でしか使われずほぼ確実に共起するので、類似度が高くなっていそうです。

１文字ペアだと感覚と合うかわかりにくいので、２文字以上でフィルターします。２桁の数字がでてきてしまうので数字もフィルターします。

# 結果を出力
count = 0
for word, id, word2, id2, cos_sim in similar_pairs:
    if count > 20:
        break
    flag = (
        not word.startswith("▁") and 
        not word2.startswith("▁") and 
        is_all_nouns(word) and 
        is_all_nouns(word2) and 
        len(word) > 1 and 
        len(word2) > 1 and
        not word.isdigit() and 
        not word2.isdigit()
    )
    if flag:
        print(f"{word} and {word2}: {cos_sim}")
        count += 1

今日 and きょう: 0.7511346340179443
状況 and 状態: 0.6744210720062256
女性 and 男性: 0.6007189750671387
事件 and 事故: 0.5230157971382141
最後 and 最初: 0.5192393660545349
中国 and 注目: 0.5156936049461365
影響 and 映像: 0.4912232756614685
今回 and 今後: 0.4902392029762268
起き and 落ち: 0.4852031469345093
映像 and 様子: 0.4760962128639221
ック and ップ: 0.45926761627197266
バー and パー: 0.4554426074028015
女性 and 状態: 0.4507766664028168
警察 and 生活: 0.4462202191352844
今回 and 現在: 0.4384438693523407
こと and 今年: 0.4370863139629364
感染 and 男性: 0.4363980293273926
最後 and 最高: 0.4320104122161865
予想 and 様子: 0.42967063188552856
ティ and ディ: 0.4291040301322937
最後 and サイ: 0.4248266816139221

「起き-落ち」「影響-映像」「女性-状態」「警察-生活」「予想-様子」などは音韻の近さが優先されている気がします。ほかは意味が優先されたペアもありそうです。

内積

上記では類似度の指標をcos類似度にしていましたが、一応、内積でも見てみます。
embeddingの大きさをみると1に正規化されているわけではないので、内積にすると結果が変わりそうです。

# 各embeddingのサイズを計算。最初の１０個だけ
for i in range(10):
    embedding_size = np.linalg.norm(embedding_matrix[i])
    print(f"Embedding {i} size: {embedding_size}")

Embedding 0 size: 0.8634994626045227
Embedding 1 size: 1.7790786027908325
Embedding 2 size: 2.722233772277832
Embedding 3 size: 1.2126015424728394
Embedding 4 size: 1.309340238571167
Embedding 5 size: 1.3232029676437378
Embedding 6 size: 1.3354939222335815
Embedding 7 size: 1.322817325592041
Embedding 8 size: 1.3636163473129272
Embedding 9 size: 1.3340636491775513

上位類似度ペア

cos類似度行列を計算したところを内積行列に変えて、スコアの高いペアを眺めます。

import numpy as np

# 内積行列を計算
dot_product_matrix = np.dot(embedding_matrix, embedding_matrix.T)

# 類似度の高いペアtop100を取得
dot_product_similar_pairs = []
vocab_size = embedding_matrix.shape[0]

for word, id in vocab.items():
    for word2, id2 in vocab.items():
        if id < id2:
            dot_product_similar_pairs.append((word, id, word2,id2, dot_product_matrix[id, id2]))

# 類似度の降順にソートしてtop100を取得
dot_product_similar_pairs.sort(key=lambda x: x[-1], reverse=True)

# 結果を出力
for word, id, word2, id2, sim in dot_product_similar_pairs[:10]:
    print(f"{word} and {word2}: {sim}")

▁今 and ▁さあ: 7.183957576751709
▁もう and ▁これは: 7.179233074188232
▁ and ▁そんな: 7.177446365356445
▁今 and ▁もう: 7.172869682312012
▁さあ and ▁えっ: 7.171679496765137
▁もう and ▁そんな: 7.171298503875732
▁そして and ▁さあ: 7.168511390686035
▁そして and ▁今: 7.168503761291504
▁そして and ▁そう: 7.168153762817383
▁お and ▁さあ: 7.168123245239258

結果はcos類似度の場合と異なっていますが、サブワードが上位に来ることは同じようです。
cos類似度とのときと同様にサブワード以外のペアの上位を眺めてみます。

# 結果を出力
count = 0
for word, id, word2, id2, sim in dot_product_similar_pairs:
    if count > 20:
        break
    if not word.startswith("▁") and not word2.startswith("▁"):
        print(f"{word} and {word2}: {sim}")
        count += 1

<s> and </s>: 7.122707843780518
<s> and <pad>: 7.104273796081543
</s> and <pad>: 7.092230796813965
逮 and <s>: 5.964789390563965
逮 and <pad>: 5.8954267501831055
逮 and </s>: 5.7955403327941895
麟 and 麒: 5.258467674255371
紡 and 紬: 5.221790313720703
謗 and 誹: 5.159886360168457
による and によりますと: 5.084996223449707
迭 and </s>: 4.786386489868164
疆 and 錮: 4.781865119934082
迭 and <pad>: 4.775331020355225
疆 and <s>: 4.758637428283691
疆 and <pad>: 4.7537102699279785
迭 and <s>: 4.748174667358398
峰 and 峯: 4.725854873657227
疆 and </s>: 4.687068462371826
諮 and 逮: 4.685540199279785
されている and されています: 4.682553291320801
濫 and 氾: 4.671710968017578

漢字1文字のトークンがたくさん出てきました。
一部ひらがなのもありますが、音韻だけでなく意味の影響もありそうな気がします。

念の為、2文字以上のトークンに限定して結果を見てみます。

# 結果を出力
count = 0
for word, id, word2, id2, sim in dot_product_similar_pairs:
    if count > 20:
        break
    if not word.startswith("▁") and not word2.startswith("▁") and len(word)>1 and len(word2)>1:
        print(f"{word} and {word2}: {sim}")
        count += 1

<s> and </s>: 7.122707843780518
<s> and <pad>: 7.104273796081543
</s> and <pad>: 7.092230796813965
による and によりますと: 5.084996223449707
されている and されています: 4.682553291320801
になります and になりました: 4.175802707672119
されて and されています: 3.935055732727051
されて and されている: 3.860311985015869
んですけど and んですけれども: 3.818894624710083
されました and になりました: 3.7648262977600098
されました and されています: 3.715080738067627
というのは and っていうのは: 3.7116668224334717
ください and てください: 3.681180000305176
される and され: 3.5976014137268066
的に and 的な: 3.562657594680786
みたいな and みたい: 3.561537265777588
される and されている: 3.5403590202331543
くなって and くなる: 3.5101776123046875
を受け and 受け: 3.4987425804138184
12 and 11: 3.4608535766601562
しています and されています: 3.436283588409424

cos類似度と似たような感じでした。

名詞

名詞に限定して結果を見てみます。
2文字以上、数字以外のフィルターもかけます。

# 結果を出力
count = 0
for word, id, word2, id2, sim in dot_product_similar_pairs:
    if count > 20:
        break
    flag = (
        not word.startswith("▁") and 
        not word2.startswith("▁") and 
        is_all_nouns(word) and 
        is_all_nouns(word2) and 
        len(word) > 1 and 
        len(word2) > 1 and
        not word.isdigit() and 
        not word2.isdigit()
    )
    if flag:
        print(f"{word} and {word2}: {sim}")
        count += 1

状況 and 状態: 3.4362664222717285
女性 and 男性: 3.4362664222717285
中国 and 注目: 3.4362664222717285
午後 and 午前: 3.4362664222717285
今日 and きょう: 3.4362664222717285
事件 and 事故: 3.4362664222717285
影響 and 映像: 3.4362664222717285
起き and 落ち: 3.4362664222717285
映像 and 様子: 3.4362664222717285
警察 and 生活: 3.4362664222717285
容疑者 and 会社: 3.4362664222717285
容疑者 and オリンピック: 3.4362664222717285
予想 and 様子: 3.4362664222717285
今回 and 今後: 3.4362664222717285
ラー and ダー: 3.4362664222717285
最後 and 最初: 3.4362664222717285
女性 and 状態: 3.4362664222717285
中国 and 全国: 3.4362664222717285
チーム and 注目: 3.4362664222717285
政府 and 逮捕: 3.4362664222717285
感染 and 男性: 3.4362664222717285

細かい中身は変わりますが、概ねcos類似度と同じ感じで、音韻が優先されていそうなものもあれば、意味に引きづられていそうなものもある、という感じです。
「映像-様子」「容疑者-オリンピック」などは全くとは言いませんが他のに比べると音韻も意味も似てなさそうです。

おわりに

音韻的な特徴を反映した単語のembeddingを取得するために日本語で学習済みのword2vec2のCTCモデルの最終層からトークンに対応するベクトルを取得し、類似度の高いトークンペアを眺めてみました。
結果としては、音韻と意味の近さの両方が考慮されたようなembeddingになっている印象でした。
通常のBERTなどよりはマシな可能性が高そうですが、純粋に音韻の類似度を知りたいという目的に対しては、意味の影響が邪魔をしているようにも思われました。

類似度指標はcos類似度でも内積でも似たような結果でしたが、内積は少し結果が荒れているようにも見えたので、cos類似度が無難な気がしました。

意味に引きづられる問題については、今回のモデルではトークンが単語に近く意味を持ちすぎているので、音素列でCTC学習されたモデルを使うと、もう少し音韻の影響が強いembeddingが得られる可能性があると思いました。ただ、その場合、音素の系列であるところの単語のembeddingをどのように得るかが課題になります（日本語トークンも辞書サイズは小さいので同じ問題はあります）。ナイーブには平均を取ってみてもよいですが、語順の影響がなくなるので結構もったいないです。ただ、CTCモデルの最終層に基づく限りは語順を無視してしまうことは避けられないので、音素列で事前学習されたBERT（PhonemeBERTなど）を使うなど、別のアプローチを考える必要があるかもしれません。

総じて、すぐに使えそう、というほどの結果ではありませんでしたが、とはいえ、比較的手軽に音響特徴を捉えたトークンの表現が得られるのはとてもおもしろいなと思いました。音響特徴の表現を得るためのアプローチはいろいろありそうなので、機会があればまた何か試してみたいと思います。