fastText で量子化し、実用的な 1.7MB のテキスト分類器を作る
英語記事をAI関連 or AI関連でない、というテキスト分類器を fastText + 量子化で作ったら、ファイルサイズ1.7MBで実用的なモデルができてしまって驚き、というお話をメモ記事に。1.7MB ですよ!!
AI Newsでは、AI関連 or AI関連でないのテキスト分類にAI News の公開と、裏側の OpenAI の活用話で書いた通り、OpenAIのtext-embedding-ada-002で1536次元のベクトルにしたものをlightGBMで学習させたものを使っている。この方法の問題は、すべての記事の判定に必ず OpenAI の API を通す必要があり、長文記事が沢山あると日によっては一日あたり数十円がかかってしまっている。月にすると500~1000円はかかってそうで、チリツモヤマトテナデコである。
というわけで、そろそろデータも溜まってきたしと、OpenAIのAPIにコストを掛けない方法で記事分類をしたいと結果を見る。なおデータ元は、英語約1100記事タイトルと本文(AI記事2:それ以外8)を7:2:1(train/valid/test)に分けて学習。また、「AI記事をAI記事でない」という誤判定をできる限りしてほしくないため、正解率accに加え、recallの値も見ている。
今の分類器: OpenAI embeddings + lightGBM で学習
正解率 acc は 0.9636、recall は 0.777。fine-tuneなしのembeddingsなのに、かなり高い。
transformer (deberta-v3-xsmall)
正解率 acc は 0.9636、recall は 0.888。ちゃんと transformer (deberta-v3-xsmall) でfine-tuneしたもの。テストデータ件数が少ない(110件ぐらい)ので何とも言えないが、今今の分類機とほぼ同性能。まぁfine-tuneしたしね…。
ちなみに、deverta-v3-large でもやってみたが、かえってスコアが落ちる感じ。学習データが少ないから、大きなモデルのfine-tuneではうまくfitしなかったのかも。
fastText (cc.en.300)
正解率 acc は 0.9454、recall は 1.0。正解率は下がったが、recall は 1.0。テストデータが少ないからrecall性能めちゃ良い！と強く言えるというわけではないが、バランス良さそうな仕上がり。
ちなみにこの時の学習後のモデルサイズは 4.6GB である。もともとの cc.en.300 が大きいからね。
fastText (ag news)
正解率 acc は 0.9454、recall は 1.0。cc.en.300 と変わらない結果になった。ちなみに学習後のモデルサイズは 88MB とこの時点でだいぶ小さくなった。
なお ag news はSupervised modelsから持ってきたもの。cc.en.300 が Common Crawl と Wikipedia から学習させた300次元のデータで、ag newsはニュース記事のタイトルと説明のコーパスから学習した10次元のデータ。今回の用途では結構マッチしているデータセットなので、相性も良かったと思える。
と、この時点でもそれなりに小さくて実用的なモデル担ったのだけど、Supervised modelsを見ると超軽量なモデルも同時に提供していて、量子化したものとな。スコアも例えば ag news なら 0.924 -> 0.92 とそんなに悪くなっていない。かつモデルサイズは 387MB → 1.6MB とな。
というわけで、量子化したものでも性能を見てみよう。
fastText (ag news) + 量子化
正解率 acc は 0.9363、recall は 1.0。正解率が若干下がったが、モデルサイズは 88MB → 1.7MB と大幅な減量。
ちなみに量子化は、fasttext コマンドならこんな感じでサクッとできる。また量子化したモデルを使って推論するときも、とくに何もせずに利用ができた。
fasttext quantize -output ./trained.ag_news -input ./trained.ag_news.bin -qnorm -retrain -cutoff 100000
おまけ: fastText (cc.en.300) + 量子化
正解率 acc は 0.9181、recall は 0.714。おお、めちゃスコアが下がったぞ。モデルサイズは 4.6G → 16MB と大幅に下がったけど。cc.en.300 は300次元と次元数が多いのが災いしているのかもしれない。
性能まとめ
今まで試した分類器の性能は以下の感じ。OpenAI の embeddings 性能の良さがうかがえる。お金を気にしなければこれでいいじゃんという感がしてしまう。deberta-v3 は動かすマシンリソースがあるなら良い。ただAI Newsのデータ処理には、メモリ1GBのVPSで動かしているのでオーバースペックだ。
というわけで、量子化されて省メモリ、かつ性能も実用的な量子化したfastText(ag news)を1st stageの分類機として使い、2nd stage ではOpenAI embeddings + lightGBMで、という二段構えのテキスト分類とした。1st stage で8割の記事はふるい落とされるので、2nd stage で OpenAI を叩く回数自体はガクッと減りそうだ。
- OpenAI embeddings + lightGBM
- acc 0.9636, recall: 0.777
- deberta-v3-xsmall
- acc 0.9636, recall: 0.8888
- fasttext(cc.en.300): fine-tune後: 4.6G
- acc 0.9454, recall: 1.0
- fasttext(cc.en.300): fine-tune + 量子化後: 16MB
- acc 0.9181, recall: 0.7142
- fasttext(ag news): fine-tune後: 88M
- acc 0.9454, recall 1.0
- fasttext(ag news): fine-tune + 量子化後: 1.7M
- acc 0.9363, recall 1.0
fastText + 量子化、の選択肢
今まで fastText のモデルを量子化するとこれだけモデルサイズが減るのだ、というのを知らずに過ごしてきたので、今後はスペック低くテキスト分類推論を動かしたいときには fastText + 量子化も強力な選択肢の一つとして考えていきたい。
おまけメモ bin -> vec
fasttext print-word-vectors ag_news.bin
で、単語ベクトルのみの .vec ファイルを出力できる。pretrainモデルにはこの .vec ファイルを指定する。
