LLM(大規模言語モデル)を教師モデルにして、小規模言語モデルに蒸留した話 実験版

DeepSeekの登場以来、「蒸留」という言葉が結構流行ってきましたね。

蒸留とは

蒸留(distillation) とは、
「大きくて高性能な教師モデルの知識を、小さくて軽量な生徒モデルに効率的に移す技術」
のことを指します。

元々は自然言語処理などの分野で、巨大モデルをそのままデバイスに載せるのは難しいため、
・教師モデルが出す出力(予測確率など)
・中間層の特徴量
を利用して、生徒モデルに教師モデルの「ふるまい」を学ばせる仕組みとして考えられました。

蒸留のポイントは、
・小さいモデル(生徒モデル)が教師の出力分布を真似る
・教師の持つパターン認識や知識を圧縮して伝える
・学習データが少なくても、教師の知識を最大限引き継げる
といったところです。

蒸留に使うモデルはこれだ

今までは高額なお金をかけて(GPUなど)学習していくと言う感じでしたが、この蒸留と言う手法を使えば、比較的手軽に専門領域に特化した小規模言語モデルを構築することができます

今回は、この蒸留の手法を使って、Googleコラボ上(有料版)で全て完結できるように、小さいスペックのモデルを使って実験してみたいと思います

まず、今回選ばれたモデルは
SakanaAIさんが公開している言語モデル TinySwallow-1.5B-Instructです。

通常の7B以上とかのモデルだと大きすぎてコラボにうまく載せられません。
そこで性能のいい大規模言語モデルのTinySwallow-1.5Bを選択しました。

このモデルを教師モデルにして、いわゆる生徒モデルもTinySwallow-1.5Bを使って、「空の生徒モデル」を構築し、そこに知識を蒸留(distillation)していくプロセスを試しました。

教師モデルの準備

まずは教師役のモデル、
SakanaAI/TinySwallow-1.5B-Instruct
を使います。

パラメータ数は約 1.54B(15億パラメータ) で、推論能力に優れたインストラクション・フォロー型のLLMです。

このモデルはHugging Face上に公開されており、まずColabへダウンロードしました。

モデル本体(model.safetensors)、トークナイザーの設定ファイル、その他メタデータまで合わせて14ファイル。
容量としてはおおむね3GB超ですが、2回目以降のキャッシュが効くため実質的には快適に動きます。

生徒モデルの構築

次に、生徒モデルをTinySwallow-1.5Bで「空の状態」から作ります。

ここでいう「空」とは、重みをほとんど初期化状態にして、教師の動きを真似させるところから始めることです。

構築した生徒モデルは およそ108Mパラメータ。
これは教師の約1/14ほどのサイズで、とても軽量です。

蒸留の流れとLossの推移

ここから教師モデルの出力をガイドにして、生徒モデルに学習させます。

やり方としては教師モデルの出力ロジットを教師信号として、損失関数にKLダイバージェンスなどを適用して生徒モデルを収束させていく形を取ります。

いわゆる Knowledge Distillation の定番手法です。

今回は以下の条件で蒸留を走らせました。

学習回数:10エポック

Optimizer:Adam

学習率:調整済

損失関数:教師出力との距離を最小化

実際のLossの推移はこんな感じです。

Epoch 1/10 – Loss: 13.1150
Epoch 2/10 – Loss: 12.8502

Epoch 10/10 – Loss: 11.9068

Lossはゆるやかに減少しており、少なくとも教師のふるまいを表面的には学習できるところまで進んでいるように見えました。

簡易テキスト生成のテスト

学習が終わったところで、生徒モデルのテキスト生成能力をテストしました。
(最終エポックでロスが11以上あるので、それなりの回答を想定)

プロンプトとして
「昔々あるところに、」
を入力してみたところ、生成結果は以下の通り。

昔々あるところに、thenReturn matière’),
sent_AP-directory Caucasian UserData\Active نوف UserData剑 nrparer買う-directoryFinoxious lanes.a比利 apenas要坚持_APB foo딛食べて剑 grants临近小额贷款 UFO apenas-directoryDiese depparer atIndexparermov SeiteDis fooموظ

正直なところ、単語が多言語混在で意味不明気味です。

しかしこれは想定内で、まずまずオッケーと言ったところです。

理由としては、学習データが非常に少量だったことが主因で、文章の一貫性を持たせるほどの蒸留は進んでいないです。

とはいえパラメータ108Mという極小モデルに10エポックだけでここまで出力させるのは、それなりに面白い成果ともいえます。

まとめと今後の課題

今回の取り組みは

・教師モデルに小さめの大規模言語モデル(TinySwallow-1.5B)を使用
・そこからさらに桁違いに小さい生徒モデル(108M)を作る
・生徒は初期化状態から蒸留で構築する

という実験でした。

結果としては、プロンプトに対して「なんらかの言葉を返せる」レベルには到達しました。

より本格的にかつ専門領域に特化した蒸留を行うなら

・教師モデルの出力を使って大量のペアデータを生成
・なおかつ専門性の高い、質の良いデータを揃える
・学習エポックを増やす

などの工夫が不可欠です。

もっと大きなモデルでやりたければ、AWSなどを使って、中堅の大規模言語モデルを教師モデルとし、小規模言語モデルへ蒸留していけば、結構いい感じの小型モデルの構築が可能です。

さて今回は「とにかく空のモデルからコラボでサクッと蒸留する」というチャレンジでしたが、技術的におもしろい試みでした。

今後はもう少しデータを増やして、専門特化した小型の言語モデルを作ってみたいと思います。

(とはいえ、数十万件のデータそろえるのはちょっとなあ…。1万件くらいならいけるかなあ。
ゴルフ専用モデル、アート専用モデルとか作りたいよなあ)

****************

最近のデジタルアート作品を掲載!

X 旧ツイッターもやってます。
https://x.com/ison1232

インスタグラムはこちら
https://www.instagram.com/nanahati555/

****************

Follow me!

PAGE TOP