はじめに
ACL 2024において、LLMの学習済みパラメータの足し引きを行うことで、事前学習済みモデルに対話能力を与えるChat Vectorが発表されました。
本稿では、Chat Vectorを実装し、評価することでその有効性を理解したいと思います。
Chat Vectorとは
Chat Vectorは、学習済みモデルの演算によって、事前学習モデルにSFTやRLHFを行った場合と同じような能力を装備させる手法です。論文中では、英語の事前学習モデル（LLaMA2）とそのInstruct Tuningモデル（LLaMA2-chat）を用意しています。Instruct Tuningモデルから事前学習モデルのパラメータを引いたベクトル（Chat Vector）は、Instruct Tuningによって得た言語能力を表現しており、これを英語以外の言語の事前学習モデルのパラメータに足し合わせることで、そのモデルも言語能力を獲得できると述べられています。加えて、Chat Vectorによって作られたモデルをその言語データセットで追加学習すると、さらなる性能向上が期待できることも報告されています。
論文より参照
Deep Learningモデルのパラメータを足し引きするという考え方は、突如出てきたものではありません。こちらの論文では、Fine Tuningモデルのパラメータから事前学習モデルのパラメータを引くことで得られるタスクベクトルが、単純な演算で修正可能であり、それによってモデルの振る舞いを制御可能だと示されています。つまり、word2vecで有名な「King - man + woman = Queen」と同様のことが、LLMのタスクベクトルでも成立するということです。Chat Vectorは、大規模な計算資源を必要とせず、Fine-tuningなどの手法に比べてコスト大幅に抑えた手法です。
今回は、Mistral-7B-v0.1を日本語データで継続事前学習したshisa-gamma-7b-v1に数学の能力を付与してみました。
実装
実装はこちらの記事を参考にさせていただきました。コードは以下に公開しています。
使用するモデルと役割をまとめます。
| Model | 役割 |
|---|---|
| mistralai/Mistral-7B-v0.1 | 英語の事前学習モデル |
| WizardLMTeam/WizardMath-7B-V1.1 | 英語の数学問題を解くことに長けたモデル |
| augmxnt/shisa-gamma-7b-v1 | 英語の事前学習モデルを日本語データで継続事前学習したモデル |
1. モデルのロード
まずは、使用するモデルとトークナイザをそれぞれ準備します。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
eng_model = AutoModelForCausalLM.from_pretrained(
"mistralai/Mistral-7B-v0.1",
torch_dtype=torch.bfloat16,
device_map="cpu",
token=HF_TOKEN
)
eng_tokenizer = AutoTokenizer.from_pretrained(
"mistralai/Mistral-7B-v0.1"
)
math_model = AutoModelForCausalLM.from_pretrained(
"WizardLMTeam/WizardMath-7B-V1.1",
torch_dtype=torch.bfloat16,
device_map="cpu",
)
jp_model = AutoModelForCausalLM.from_pretrained(
"augmxnt/shisa-gamma-7b-v1",
torch_dtype=torch.bfloat16,
device_map="cuda",
)
jp_tokenizer = AutoTokenizer.from_pretrained(
'augmxnt/shisa-gamma-7b-v1'
)
※ モデルはすべてCPUに載せて問題ないです。（上記では、メモリ分散のためにGPUを使用しています。）
2. パラメータの演算
モデルを準備したら、Chat Vectorモデルを作成していきます。パラメータの足し引きは、以下のような単純なコードで実装することができます。
※ 英語の事前学習モデルとパラメータを足し合わせるモデルのボキャブラリサイズが異なる場合、そのLayerを演算対象から除外する必要があります。原論文でも、語彙数の異なるモデルに関しては言及されておらず、英語以外の言語モデルへの適用は想定されていないように感じます。
if jp_tokenizer.vocab_size == eng_tokenizer.vocab_size:
# Excluded module
skip_layers = []
else:
exit(1)
for k, v in jp_model.state_dict().items():
if (k in skip_layers) or ("layernorm" in k):
continue
chat_vector = math_model.state_dict()[k] - eng_model.state_dict()[k]
new_v = v + chat_vector.to(v.device)
v.copy_(new_v)
3. Chat Vectorモデルの保存
作成したモデルを別名で保存しておきます。
new_model_name = 'chat_vector_mistral_7B_math'
jp_model.save_pretrained(new_model_name)
精度比較
Chat Vectorによって作成したモデルを、数学ベンチマークのMGSMの日本語版を用いて評価しました。
Evolutionary Model Merge（Sakana AI）によって作られたEvoLLM-JP との比較も行っています。Evolutionary Model Mergeは、Chat Vectorと同様に学習を必要とせずLLMに新たな能力を与えることのできる技術です。
| ID | Model | MGSM-JA(acc↑) |
|---|---|---|
| 1 | WizardMath 7B v1.1 | 35.2 |
| 2 | Shisa Gamma 7B v1 | 6.4 |
| 3 | Chat Vectorモデル | 36.8 |
| 4 | EvoLLM-JP | 41.2 |
- ベースモデル（WizardMath, Shisa Gamma）と比較すると精度は向上
- 特に、Shisa Gamma（Chat Vectorを足す前の日本語モデル）と比較すると大幅に向上
- EvoLLM-JPには、若干劣っている
上記の結果からChat Vectorによって、単純でありながらかなりのスコア向上（6.4 → 36.8）を実現できていることがわかりました。Sakana AIが、Evolutionary Model Mergeのコードを公開していないため、Chat Vectorを適用できるなら代わりに使っても良さそうとだなと思いました。
※ Evolutionary Model Mergeは、mergekitというライブラリを使用すると、実装可能ではあるものの完全再現とはいきません。こちらの記事で実際に実装していますのでご参考までに。
まとめ
本稿では、Chat Vectorによるパラメータの足し引きによって、日本語モデルであるShisa Gammaに数学問題を解く能力を付与してみました。単純な手法ですが、一定の性能向上を確認できました。
実際に実装してみた感想としては、英語に対して日本語のデータの規模が小さいという課題の対策として有効であると感じました。しかし、Chat Vectorはモデル構造がある程度似ていることが前提であり、ChatGPTなどのクローズドモデルにはいつまでも追いつけないのではという懸念は残ります。ただ、Chat Vectorによって作られたモデルを日本語データセットでSFTすることで、さらなる性能向上が期待されるため、リーダーボードなどの今後の動向に注目したいとおもいます。