見出し画像

Chat VectorならぬMath Vectorは作れるのか

はち

はじめに

この記事は以下記事の続きになります。

Chat Vectorと呼ばれる、重みの足し引きでFine TuningなしにChat能力を事前学習モデルに付与できるという技術あります。
この発想から、Chat能力以外にも能力の切り貼りはできるのかという検証が本記事の趣旨となります。
今回は以下の能力について試したいと思います。

  • 数学的推論能力

結論だけ書くとある程度うまくいきました。検証記録とコード、モデルリンクをまとめていきます。


1. Skill Build

今回は以下のモデルを使っていきます。

・英語Base Model：Mistral-7B-v0.1
・Skilled Model：OpenMath-Mistral-7B-v0.1-hf
・日本語Base Model：Swallow-MS-7b-v0.1

OpenMath-Mistral-7B-v0.1-hfはNvidiaが出したモデルです。Mistral-7Bに対してmath instruction tuningを行うことで数学的推論能力を上げたモデルになります。
ここからMath Reasoning能力を抽出し、Swallow-MS-7b-v1.0に装備されます。
基本的な流れは前の記事と同様です。

1-1. モデルの取得

モデルを取得します。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

base_model_name = "mistralai/Mistral-7B-v0.1"
skilled_model_name = "nvidia/OpenMath-Mistral-7B-v0.1-hf"
jp_model_name = "tokyotech-llm/Swallow-MS-7b-v0.1"

# 英語ベースモデル
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)
# Skilledモデル
skilled_model = AutoModelForCausalLM.from_pretrained(
    skilled_model_name,
    torch_dtype=torch.bfloat16,
    device_map="cpu",
)
# 日本語ベースモデル
jp_model = AutoModelForCausalLM.from_pretrained(
    jp_model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

1-2. SkillTreeの作成

毎回Chat Vectorならぬ〇〇Vectorと呼ぶのは長いので、ゲームのスキルツリーのようにモデルに能力を与えるという意味込めてSkillTreeと呼ぼうと思います。（一般的な名前が定まるまで。）
基本的にはlayernorm等、直接能力に効かなそうなところを除外して重みの差分を取ります。

# 除外対象
skip_layers = ["model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"]

for k, v in base_model.state_dict().items():
    # layernormも除外
    if (k in skip_layers) or ("layernorm" in k):
        continue
    chat_vector = skilled_model.state_dict()[k] - base_model.state_dict()[k]
    new_v = chat_vector.to(v.device)
    v.copy_(new_v)

Skillを抽出したらSkillTreeと名付けて保存します。この重みに数学的推論能力が秘められているはずです。

skill_tree_name = "HachiML/SkillTree-Math-OpenMath-Mistral-7B-v0.1"
base_model.save_pretrained(f"./models/{skill_tree_name}", repo_id=skill_tree_name, push_to_hub=True)

SkillTree-Math（Math Vector）

1-3. Skill Enhansed Modelの作成

作ったSkillTreeを使って、日本語Base Model（Swallow-MS-7b-v0.1）に数学的推論能力を付与します。

def apply_skill(model, skill_tree):
    # excluded object
    skip_layers = ["model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"]
    # apply skill
    for k, v in model.state_dict().items():
        # layernorm is also excluded
        if (k in skip_layers) or ("layernorm" in k):
            continue
        vector = skill_tree.state_dict()[k]
        new_v = v + vector.to(v.device)
        v.copy_(new_v)
    return model

model = apply_skill(jp_model, skill_tree)

tokenizerを取得して、一緒に保存します。

# 日本語ベースモデルのトークナイザ
jp_tokenizer = AutoTokenizer.from_pretrained(jp_model_name)

# 保存
model_name = "HachiML/Swallow-MS-7b-v0.1-MathSkill-OpenMath"
jp_tokenizer.save_pretrained(f"./models/{model_name}", repo_id=model_name, push_to_hub=True)
model.save_pretrained(f"./models/{model_name}", repo_id=model_name, push_to_hub=True)

Math強化モデル

2. 検証

2-1. Skill Enhansed Modelの能力検証

作ったモデルを読み込み、まずは元の日本語能力が落ちていないかを確認します。
東京工業大学の主なキャンパスは、田町キャンパスと大岡山キャンパスの2つらしいです。この2つを主なと呼んでいいかは私は知りませんが、検索すると少なくともこの2つのキャンパスは実在する様です。
普通の日本語の生成という意味では能力は残り続けている様です。

model_name = "HachiML/Swallow-MS-7b-v0.1-MathSkill-OpenMath"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")

# 日本語問題1
prompt = "東京工業大学の主なキャンパスは、"
input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
tokens = model.generate(input_ids.to(device=model.device), max_new_tokens=256, temperature=0.99, top_p=0.95, do_sample=True)
out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

東京工業大学の主なキャンパスは、東京都大田区にある東京工業大学田町キャンパスと、東京都港区にある東京工業大学大岡山キャンパスの2つがあります。 田町キャンパスは東京工業大学の旧称である東京工業学校時代からのキャンパスで、東京工業大学田町キャンパスには工学系の学部と、大学院工学系研究科、および附置研究所が置かれています。田町キャンパスは地下鉄都営浅草線の三田駅に隣接しており、アクセスが便利なのが特徴です。 大岡山キャンパスは192年に開校した東京工業大学の2番目のキャンパスで、大学院情報理工学研究科や、大学院生命理工学研究科、また総合理工系の学部が置かれています。 地下鉄目黒線大岡山駅に隣接したキャンパスで、田町キャンパスに比べて緑が多く、閑静な環境であるのが特徴です。 学校の特長 東京工業大学は、明治時代に日本で最初に設立された技術者養成機関である、東京工業学校を前身とします。 東京工業学校を前身とする大学は、東工

日本語問題1の解答

次に、本番の数学的推論能力を見ていきます。英語と日本語で数学の文章題を出します。日本語の方は英語のものを和訳しただけのもので、どちらも正解は48+24で72です。
英語も日本語も余計なテキストは生成しているものの、最終的には72と答えを出しています。どうやら数学的推論能力を手に入れている様です！

# 数学問題1
prompt = "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?\n"
input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
tokens = model.generate(input_ids.to(device=model.device), max_new_tokens=256, temperature=0.99, top_p=0.95, do_sample=True)
out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

# 数学問題2
prompt = "ナタリアは4月に48人の友人にクリップを売り、5月にはその半分の数のクリップを売った。ナタリアが4月と5月に売ったクリップの数は？:\n"
input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
tokens = model.generate(input_ids.to(device=model.device), max_new_tokens=256, temperature=0.99, top_p=0.95, do_sample=True)
out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?
Let's solve this problem using Python code.
<llm-code>
number_of_friends_april = 48
number_of_friends_may = number_of_friends_april / 2
number_of_clips_april = number_of_friends_april * 1
number_of_clips_may = number_of_friends_may * 1
number_of_clips_total = number_of_clips_april + number_of_clips_may
number_of_clips_total
</llm-code>
<llm-code-output> 72.0 </llm-code-output>
Thus Natalia sold \boxed{72} clips altogether.

数学問題1の解答

ナタリアは4月に48人の友人にクリップを売り、5月にはその半分の数のクリップを売った。ナタリアが4月と5月に売ったクリップの数は？:
ナタリアは4月に48人の友人にクリップを売り、5月にはその半分の数のクリップを売った。ナタリアが4月と5月に売ったクリップの数は? ナタリアは4月に48人の友人にクリップを売り、5月にはその半分の数のクリップを売った。ナタリアが4月と5月に売ったクリップの数は?
Let's solve this problem using Python code.
<llm-code>
clip_cost = 20
clip_sold_in_april = 48
clip_sold_in_may = clip_sold_in_april / 2
clip_sold_in_both_months = clip_sold_in_april + clip_sold_in_may
clip_sold_in_both_months
</llm-code>
<llm-code-output> 72.0 </llm-code-output>

数学問題2の解答

2-2. （参照）日本語Base Modelの能力確認

最後に、数学的推論能力が得られたものなのか元から持っていたものなのかを確認します。
同じ問題をベースとなった日本語Base Model（Swallow-MS-7b-v0.1）に与えてみます。

少なくとも英語の方は何となく答えようと頑張っているものの、答えることができませんでした。つまり、Skill Buildによって新たに能力を獲得したということになります！

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "tokyotech-llm/Swallow-MS-7b-v0.1"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")

# 数学問題1
prompt = "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?\n"
input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
tokens = model.generate(input_ids.to(device=model.device), max_new_tokens=256, temperature=0.99, top_p=0.95, do_sample=True)
out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

# 数学問題2
prompt = "ナタリアは4月に48人の友人にクリップを売り、5月にはその半分の数のクリップを売った。ナタリアが4月と5月に売ったクリップの数は？:\n"
input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
tokens = model.generate(input_ids.to(device=model.device), max_new_tokens=256, temperature=0.99, top_p=0.95, do_sample=True)
out = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(out)

Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?
A: 43 in April and 20 in May
Q: How many clips did Natalia sell in April?
A: 43
Q: How many clips did Natalia sell in May?
A: 20
Q: How many friends did Natalia sell clips to in April?
A: 48
Q: How many friends did Natalia sell clips to in May?
A: 20
Q: How many days did Natalia have in April and May?
A: 31 for each month.
Q: How many clips does Natalia have in each month?
A: 126 at the beginning of April and 120 at the beginning of May.
Q: How many clips did Natalia sell altogether in April and May?
A: 43 in April and 20 in May.
Q: How many clips did Natalia sell to her friends in May?
A: 10
Q: How many friends did Natalia have at the beginning of May?
A: 48
Q: How many friends does Natalia have after the month is over?
A:

数学問題1の解答

ナタリアは4月に48人の友人にクリップを売り、5月にはその半分の数のクリップを売った。ナタリアが4月と5月に売ったクリップの数は？:
問題 7 :次の文は、「次の数列を求める」に関する問題である。(1)30から65までの66個の数の和は49316である。(2)30から80までの81個の数の和は56550である。65と80の間の偶数は2ずつ増える。(3)65から80までの81個の数の和は、前問の問(1)の数よりも5650(0は0を表す)多くなる。65から80までのこの数列の合計の差は何であるか。
問題 8 :次の文は、「第2次大戦中、フランスに駐屯していたドイツ兵の兵営で見た奇妙なこと」に関する問題である。(1)兵営の入口には、“Wenn man eine Maus sieht, dann ist diese ja nicht so toll”ということを意味する看板が立てられていた。(2)ある日、兵営の中庭に小さな猫がやってきた。その

数学問題2の解

3. まとめ

現状の結果をまとめます。
⭕️　Chat
× 　Code
⭕️ 　Math Reasoning

Codeだけ追加の事前学習を含んでいるので、やはりChatVectorの理論が通るのはTuningだけなのかもしれません。事前学習をやってしまうとモデルの中身が大きく変わってしまうことが原因ではないでしょうか。
また別の能力で試してみたいと思います。

続き

参照


この記事が気に入ったらサポートをしてみませんか？