新手法「TAID」を用いた小規模日本語言語モデル「TinySwallow-1.5B」の公開


新手法「TAID」によって学習された小規模日本語言語モデル「TinySwallow-1.5B」は、同規模のモデルの中で最高性能を達成しています。さらに、小規模のため、APIなどを介さずお手元のスマートフォンやPCで完結してチャットが可能です。ウェブアプリまたは、GitHubにてお試しいただけます。

本稿は、経済産業省とNEDOが推進するGENIACによって支援されたSakana AIの研究成果を紹介するブログシリーズの第3回目です。


概要

日常的な対話はもとより、数学やコーティングといった複雑なタスクまでも人間と遜色ないレベルでこなせるようになった大規模言語モデル（LLM）は、今後ますます多くの場面での利活用が期待されています。こうした大きな期待の一方で、LLMの開発・活用の現場では、莫大な計算資源の必要性が大きな壁として立ちはだかってます。まずLLMの開発現場では、「大規模データと巨大計算資源によってLLMの性能を向上できる」という指針のもとGPUの購買競争が加熱しており、多くの企業が作成を望んでいるいわゆる「独自LLM」の開発は容易なことではなくなっています。LLMの活用の現場でも、サービスをユーザに届けるためにLLMをスマートフォンなどのエッジデバイスで運用することが期待されていますが、計算リソースの問題でやはり困難が大きい状況です。

Sakana AIはこれまで、Evolutionary Model MergeやThe AI Scientistなど大規模言語モデルを効率良く開発・活用するための研究を進めてきました。また並行して、小規模で高性能な言語モデルの研究・開発にも注力してきました。そしてこの度、小規模ながら高性能な言語モデル（SLM (Small Language Model)）を効率的に構築するための新手法、「TAID (Temporally Adaptive Interpolated Distillation)」を発表します。


「TinySwallow-1.5B」は、を用いたiPhone上でのテキスト生成。動画速度は実際の生成速度です。モデルをGGUFフォーマットに変換したものを、LLMFarmで動作させています。詳しいやり方は、こちらをご覧ください。


TAIDは、LLMの知識をSLMに転移させる「知識蒸留」の新たな手法です。既存の知識蒸留の手法とは異なり、TAIDは、SLMの学習過程に応じて段階的にLLMの知識を転移させます。「常にちょうど良いレベルの先生についてもらう」ことで効率的かつ効果的な知識転移を実現します。更に私たちは、TAIDによって作った小規模モデルを多くの人に使ってもらうため、英語のSLMである「TAID-LLM-1.5B」と日本語のSLMである「TinySwallow-1.5B」を開発しました。TAIDを用いて32BパラメータのLLMから1.5BパラメータのSLMへ知識を転移し、結果として1.5Bパラメータモデルとして最高性能となる日本語の小規模言語モデルを作り出すことに成功しました。

本リリースの要点は以下の通りです。

本研究は、包含さん（京都大学）・横井祥さん（国立国語研究所・東北大学・理研）の両名と共同で実施したもので、論文はarXivで公開されています。日本語のSLM「TinySwallow-1.5B」は、大規模言語モデルSwallowを開発している東京科学大学のメンバー（以下、Swallowチーム）と共同で開発したもので、作り出したモデルはHugging Face Hubで公開されています。ぜひ多くの皆さんにTAIDを活用いただき、効率的な知識蒸留と最先端の小規模言語モデルを楽しんでもらえたらと思います。また我々の知見とモデルが、日本の素晴らしいAIコミュニティのさらなる発展の一助になることを願っています。


ウェブアプリTinySwallow ChatUI。このアプリでは、TinySwallow-1.5Bをブラウザ上で動作させており、外部APIなどを介しておらず、最初のモデルのダウンロードが完了すれば完全オフラインで会話することも可能です。また、完全なローカル実行を希望される方向けに、モデルの重みを含むself-contained版としてTinySwallow ChatUI Localも用意しています。こちらは、Pythonベースのシンプルなインターフェースで、ネットワーク接続を一切必要とせずにローカルマシンからモデルを直接読み込んでチャットができます。


知識蒸留とは

高性能なSLMを学習する手法の一つとして、「知識蒸留」があります。知識蒸留は、高性能なLLM（教師モデル）が自身の知識をSLM（生徒モデル）へ転移させる技術であり、ゼロからSLMを学習するよりも、効率的かつ高性能なSLMを作ることができることで知られています。

知識蒸留の面白い点は、単に「正解」を教えるだけでなく、教師モデルの「考え方」も生徒モデルへ教示できる点です。下図は、言語モデルにおける通常の学習手法（左）と知識蒸留（右）の概念図をそれぞれ示しています。


ここでは例として、「Sakana AI develops efficient methods for ＿＿」という文中の下線部分（＿＿）の単語を予測することを生徒モデルの目標とします。通常の学習（左）では、正解単語である「AI」という情報のみが生徒モデルの学習に利用されます。 一方、知識蒸留（右）では、教師モデルは「AI」が最も適切（35%）だと判断しつつ、「ML」（25%）や「LLM」（15%）も文脈的に自然な選択肢として認識していることを生徒に教示できます。この「確率分布」として表現される判断の仕方そのものが、言語モデルの持つ重要な知識なのです。 このように知識蒸留を行うことで、教師モデルが持つ正解データ以上の豊かな情報を生徒モデルへ転移することが可能です。それでは、より良いSLM（生徒モデル）を学習するには、より大きく高性能なLLMを教師モデルとして選択するのが良いのでしょうか？ 実は、「より大きく高性能なLLMが生徒モデルにとって良いとは限らない」ことが知られており、これは従来手法の知識蒸留における本質的な課題の1つです。生徒モデルの表現力は限定的であるため、教師モデルと生徒モデルの間の能力差によって知識転移が失敗してしまうのです（大学院レベルの専門知識をいきなり小学生に教えようとするようなイメージ）。


TAIDの提案

TAIDは上述の課題を解決する新たな知識蒸留の手法です。TAIDでは、生徒モデルの学習進度に合わせて、教師モデルを段階的に変えることで、効果的な知識転移を実現します。


具体的には、学習段階の生徒モデルの能力に合わせた「中間教師（intermediate teacher）」を導入し学習を行います。中間教師は、生徒モデルにとって理解しやすく、かつ少し高度な知識を持つような、まさに生徒モデルと教師モデルの橋渡しをするような中間的な教師モデルです。生徒の学習が進むにつれて、中間教師は徐々により高度な知識を持つように変化していきます。これは、学校教育で教師が生徒の理解度に合わせて少しずつ難しい内容を教えていくのに似ています。TAIDの考え方は、人間の「良い学習」から着想を得ており、言語モデルの世界でもこのような学習方法が効果的であることを示せました。

下図は、教師モデルの大きさ（横軸：Teacher Size）を変化させたときの、従来手法とTAIDによる知識蒸留を行った生徒モデルの性能（縦軸：Test Accuracy）を表しております。従来手法（KLおよびRKL）は、教師モデルが大きくなるとむしろ性能が低下してしまうことが確認できます。一方、TAIDは、教師モデルが大きさに比例して、生徒モデルの性能は単調増加しています。これは、TAIDが教師と生徒の能力差を効果的に克服できていることを示しています。


異なる教師モデルのサイズに対する生徒モデル（70M）の性能。モデルは英語事前学習モデルPythia Suiteを用いており、SmolLM-Corpusの1Bサブセットで継続事前学習を行いました。Test Accuracyは、lm-eval-harnessを用いたLAMBADAデータセットの正解率を表しています。

TAIDの論文は、機械学習分野におけるトップ国際学会であるICLRに採択されました。論文は、arxivで公開されておりますので、技術的な詳細については論文をご覧ください。


日本語小規模言語モデル「TinySwallow-1.5B」の公開

TAIDのさらなる有効性検証のため、私たちはSwallowチームと共同で日本語SLM「TinySwallow-1.5B」を開発しました。このモデルは、320億パラメータのLLMから約1/20の大きさの15億パラメータのSLMへTAIDによる知識蒸留を行うことで構築しています。 以下の表が示すように、TinySwallow-1.5Bは、日本語での言語モデルベンチマークにおいて、同規模のモデルの中で最高性能を達成しています。

3B未満のSLMに対する、日本語LLM評価ベンチマーク日本語理解・生成タスクの平均スコア

さらに、小規模サイズであるTinySwallow-1.5Bは、PCはもちろん、モバイル端末でのテキスト生成が可能です。概要セクションでお見せしたデモ動画は、TinySwallow-1.5B-InstructをiPhone 14で動作させたときの様子です。ご覧の通り、高速にテキスト生成できていることがわかります。また、PCでの動作を簡単にお試しいただくために、ウェブアプリをご用意しました。


ウェブアプリTinySwallow ChatUI。このアプリでは、TinySwallow-1.5Bをブラウザ上で動作させており、外部APIなどを介しておらず、最初のモデルのダウンロードが完了すれば完全オフラインで会話することも可能です。また、完全なローカル実行を希望される方向けに、モデルの重みを含むself-contained版としてTinySwallow ChatUI Localも用意しています。こちらは、Pythonベースのシンプルなインターフェースで、ネットワーク接続を一切必要とせずにローカルマシンからモデルを直接読み込んでチャットができます。

モデルの重みはHugging Face Hubにて公開しております。日本語事前学習データで継続事前学習を行った「TinySwallow-1.5B」と、さらに指示学習を行った「TinySwallow-1.5B-Instruct」を公開しております。ぜひお試しください。

また、完全なローカル実行を希望される方向けに、モデルの重みを含むself-contained版のウェブアプリも用意しています。

こちらは、Pythonベースのシンプルなインターフェースで、ネットワーク接続を一切必要とせずにローカルマシンからモデルを直接読み込んでチャットができます。ぜひローカル上で小型言語モデルとのチャットを実感してみてください。

小規模なTinySwallowを用いて、以下のような様々な応用が可能です。


動画では、iPhone上でのブログ要約の実際の生成結果を示しています。開発者の皆様のアイデアで、TinySwallowを用いた新しい応用が生まれることを楽しみにしています！


今後の展望

TAIDは、大規模なモデルの知識を効果的に小規模モデルへ転移させる新しい知識蒸留手法として、大きな可能性を秘めています。本稿では、TAIDを用いた日本語SLMの構築について紹介しましたが、TAIDの適用範囲はLLMに限定されません。事実、論文内では小規模な英語視覚言語モデル「TAID-VLM-2B」の構築も行っており、既存手法よりも優れたモデル構築が可能であることを確認しております。今後はさらに様々なモデルや環境条件下での検証を進めていく予定です。

Sakana AIは、これまで自然界から着想を得た様々な研究プロジェクトを発表してきました。TAIDにおいても、人間の学習過程を模すことで、滑らかな知識蒸留を実現しました。

Sakana AIは、高性能なモデルをより身近なものにすることで、AIの恩恵をより多くの人々が享受できる世界の実現を引き続き目指します。




謝辞

本研究は、国立研究開発法人新エネルギー・産業技術総合開発機構（NEDO）および経済産業省が主催する「Generative AI Accelerator Challenge（GENIAC）」のプロジェクト（JPNP20017）として実施されました。プロジェクトへの採択と支援に深く感謝申し上げます。


Sakana AI

日本でのAIの未来を、Sakana AIと一緒に切り拓いてくださる方を募集しています。当社の募集要項をご覧ください。