日本語事前学習済みALBERTモデルを公開します

2022/04/21 追記

本モデルのスピンオフ的な、トークナイザーを差し替えたものを新たに公開したのでお好みでどうぞ
https://zenn.dev/ken_11/articles/0e2e231f321c5d

本題

どうもこんばんは。
今回は掲題の通り、日本語事前学習済みALBERTモデルを公開したので、その過程やらなにやらを紹介します。(ほぼポエム)

albert-base-japanese-v1

https://huggingface.co/ken11/albert-base-japanese-v1
こちらがそのモデルです。
よければ使ってみてください。
ここから先はわりとどうでもいい話です。

ALBERTって？

詳しい話は論文なり解説記事なり読んでください。
大切なのはこれが「A Lite BERT」のことで、すごく雑に言えば「軽量化されたBERT」ということです。

なぜ事前学習済みモデルを作ったのか

結局のところ「自分がちょうど欲しいくらいの事前学習済みモデルがなかった」から作ったというDIY精神にほかなりません。
今回だと前提として「BERTはいいけどモデルサイズが大きくて取り回しが効かないシーンがある」という課題を感じていて「ALBERTなら精度を大きく落とさずに軽量化できるのでは」と思ったので作成しました。

やったこと

学習データにはWikipedia全文とlivedoorニュースコーパスを使いました。
定番なので、この辺の入手方法や使い方はここでは紹介しません。

ALBERTにはbase/largeといったサイズの違いと、v1とv2がありますが、アーキテクチャ的にはbase-v1です。
dropoutが隠し味程度にちょっぴり入ってます←
ちなみにv2の違いは公式には以下のように説明されています。

In this version, we apply 'no dropout', 'additional training data' and 'long training time' strategies to all models.

僕の雑な理解ではgeluがgelu_newになってdropoutがなくなった、くらいの感覚です。

バッチサイズ128で200万ステップほど学習させました。
これはベンチしていたBERTの東北大モデルが256で100万ステップだったので、とりあえずそれくらいを目標にした次第です。(もちろんBERTとALBERTでは全然異なるので、この意思決定自体はとても無意味で適当なものです)

困った点

実はhuggingfaceのTokenizerは、Sentencepieceを使ったときにspecial tokenのあとに余計なトークンが出現するという問題に気づいてしまいました。
なにが困るって、これケアしないとまともに推論できないので、pipelineとかで雑に確認するとおかしな結果になってしまいます。
公開したモデルカードでも同じことを書いてますが、以下のようにケアしてあげないと正しい推論結果が得られません。
なので、モデルカードの横についているウィジェットもカッコいい推論結果が出ないんです。

from transformers import (
    AlbertForMaskedLM, AlbertTokenizerFast
)
import torch


tokenizer = AlbertTokenizerFast.from_pretrained("ken11/albert-base-japanese-v1")
model = AlbertForMaskedLM.from_pretrained("ken11/albert-base-japanese-v1")

text = "大学で[MASK]の研究をしています"
tokenized_text = tokenizer.tokenize(text)
del tokenized_text[tokenized_text.index(tokenizer.mask_token) + 1]

input_ids = [tokenizer.cls_token_id]
input_ids.extend(tokenizer.convert_tokens_to_ids(tokenized_text))
input_ids.append(tokenizer.sep_token_id)

inputs = {"input_ids": [input_ids], "token_type_ids": [[0]*len(input_ids)], "attention_mask": [[1]*len(input_ids)]}
batch = {k: torch.tensor(v, dtype=torch.int64) for k, v in inputs.items()}
output = model(**batch)[0]
_, result = output[0, input_ids.index(tokenizer.mask_token_id)].topk(5)

print(tokenizer.convert_ids_to_tokens(result.tolist()))
# ['英語', '心理学', '数学', '医学', '日本語']

学習してみて「いまどんな感じだろう」なんて思ってちょろっとpipelineに突っ込んで試すみたいなのを最初の頃ずっとやってたんですが、「全然いい感じの結果にならねえじゃん」ってなってました。全部これのせいです。
今となっては「こんなに学習しなくても十分いいモデルになってたのでは？」という疑念すら沸いてきます。もっと早く気づくべきだった :kusa:

ちなみに、これ同様の問題がis_split_into_words=Trueオプションにも隠れていて、これやると漏れなくトークンのあとにmeta symbolが出現するという魅惑の機能になってしまっています。
例: ['▁', 'こんにち', '▁', 'は', '▁', '[MASK]', '▁', 'の人々']
なので、NERタスクの学習とかマジで気をつけないとうまくできないです。
全く精度出なくてしばらく悩んでましたがこれのせいでした。泣きたい :ksua:

精度について

日本語はGLUE的なのないので、精度を簡単に比較する方法ってないんですよね。。
参考程度に、最近よくやってるお手軽NERで東北大BERTと比較してみました。

ルール
いずれも同じデータセットを使って10エポック程度学習(ファインチューニング)し、その後同じテストデータを使って結果を比較
データはストックマーク株式会社が公開しているner-wikipedia-datasetを使用
テストには全体の8%のデータを使用
結果はseqevalを使ってmode=defaultのscheme=IOB2で評価
Tokenizerがそれぞれ異なるので、support数は若干ズレるが大きな問題ではないので無視する

今回作成したALBERTのモデルをファインチューニングしたもの

東北大BERTをファインチューニングしたもの

悪くなさそうです。
あまりにも局所的なタスクなのでこの精度差に意味はあまりないと思いますが、事前学習済みモデルとしては十分役に立つレベルになってそうというのはわかりました。
これならある程度実用できそうです。

まとめ

というわけで、日本語事前学習済みのALBERTモデルを作成し、huggingface上で公開しました。
恐らく僕の2021年最後のアウトプットかなと。
ALBERTはホントに軽量モデルなので、このモデルをベースにファインチューニングすることで、様々なタスクをエッジ含め様々なシーンで実行できるようになると思います。
精度等どれくらい求めるかにもよってくると思いますが、モデルが軽くなればコストも抑えやすくなりますし、今まで躊躇していた領域でも活かせるのではないかなと期待しています。

拙作でもお役に立てそうであればぜひともご活用ください。