RWKV(Receptance Weighted Key Value)をつかってみた

本記事は、当社オウンドメディア「Doors」に移転しました。

約5秒後に自動的にリダイレクトします。


このたびブレインパッドは、LLM/Generative AIに関する研究プロジェクトを立ち上げ、この「Platinum Data Blog」を通じてLLM/Generative AIに関するさまざまな情報を発信をしています。
この記事では、Google Colaboratory を使って、チャットAI「ChatRWKV」の利用方法、ファインチューニング手法を紹介する他、ファインチューニング済みモデルとベースモデルの挙動の比較結果もご紹介します。

はじめに

こんにちは。ブレインパッドの丸山です。
最近、GPTなどのTransformerベースの生成系AIが注目を集めていますが、より高速な推論が可能なモデルとして、RWKV(Receptance Weighted Key Value)という新たなアーキテクチャが提案されました。RWKVはApache License 2.0の下で提供されており、商用利用が可能です。RWKVの登場により、自身の環境にチャットAIを導入し、利用目的に合わせてカスタマイズすることができる可能性が広がりました。
この記事では、Google Colaboratoryを使用して、RWKVをベースにしたチャットAIであるChatRWKVの利用方法と、ファインチューニング手法も紹介し、ファインチューニング済みモデルとベースモデルの挙動を確認していきます。

ChatRWKVの活用手順

Google Colaboratoryを使用して以下の手順で実行します。
なお、こちらの記事を参考に実行しています。

# Googleドライブのマウント
from google.colab import drive
drive.mount('/content/drive')

# 作業フォルダへの移動
import os
os.makedirs("/content/drive/My Drive/rwkv", exist_ok=True)
%cd '/content/drive/My Drive/rwkv'

# 環境構築
!git clone https://github.com/BlinkDL/ChatRWKV
%cd ChatRWKV
!pip install -r requirements.txt
!pip install transformers
!pip install ninja

モデルは事前にhttps://huggingface.co/BlinkDLからダウンロードして、事前にディレクトリに配置する必要があります。今回は「”RWKV-4-Pile-7B”」を使用します。
v2/chat.py内でモデルのpathを指定して、v2/chat.pyを実行するとチャットが可能になります。

まずは、英語で以下の2つの質問を投げかけてみます。

user: What is artificial intelligence?
bot: Artificial intelligence is the study of intelligent machines and other forms of artificial intelligence. AI has been the focus of many research and development in science, technology and philosophy. It has always been the subject of interest for computer scientists, mathematicians, engineers and scientists. The advent of AI has become a turning point in the history of science and technology.

user: What is the most famous novel?
bot: The most famous novel is "Robinson Crusoe". In fact, this novel was written by Daniel Defoe, and it was published in 1719. It tells a story about the adventures of a shipwrecked sailor.

次に、日本語の入力に対しての回答を見てみます。

user: 人工知能とは何ですか?
bot: 人工知能は、人間の知性をコンピューターに取り込むことです。人間の知性は、自分の思考や感情をコンピューターに取り込むことでできます。

user: 最も有名な小説は何ですか?
bot: 「ハリー・ポッター」です。ハリー・ポッターは、世界中の人々に愛されています。

日本語にもある程度対応しているようですが、回答の精度が英語に比べると下がるようです。

RWKVのファインチューニング

ここからはファインチューニングを行っていきたいと思います。
今回は以下の2つの記事を参考にニュース記事をインプットして、ニュースのタイトルを生成するタスクに適応させてみたいと思います。

今回学習用データセットには、日本のさまざまなニュース記事から構成されるデータセットであり、多様なトピックや文体を含んでいるlivedoor newsコーパスを用います。
ファインチューニングの手順は以下の通りです。まず、livedoor newsコーパスをダウンロードし、以下の形式に置き換え、Googleドライブにアップロードします。

#学習データセットのサンプル例(筆者が作成した内容で実際のデータセットには含まれていません。)
# Instruction: 
以下は、あなたへの依頼事項です。次のニュースのタイトルを考えて,
# Input: 
【ロサンゼルス】近日公開予定のSFアクション映画『ABCD』が、観客を魅了する特殊効果と壮大なストーリー展開で話題を集めています。本作は、未来の地球を舞台にした壮大な冒険を描いた作品であり、映画ファンからの期待が高まって~(以下続く))
# Response:
新作SF映画『ABCD』、観客を魅了する特殊効果と壮大なストーリー展開

次にこちらのレポジトリ内にある以下のnotebookを使用してファインチューニングを実行します。
「Base Model Options」のセルでbase modelを指定します。今回は"RWKV-4-Pile-430M"をベースモデルとして使用しています。(筆者はcolabo pro環境下で実行しましたが、"RWKV-4-Pile-430M"より大きなモデルでは学習できませんでした。。)
「Training Data Options」で先ほどアップロードした学習用データセットのパスを指定してtrain.pyを実行します。1エポックあたり、2分30秒程度で学習することができました。今回は100epoch実行しています。

ファインチューニング済みのモデルの実行

では、実際に学習したモデルとベースモデルの出力を比較してみたいと思います。
タスクは以下の3つで、このような形式でプロンプトとして入力しています。

タスク①
# Instruction:
以下は、あなたへの依頼事項です。次のニュースのタイトルを考えて
# Input:
株式会社ブレインパッドは、顧客一人ひとりに自然なアクションを促す、Web/アプリのコンテンツ最適化プラットフォーム「Rtoaster action+(アールトースター アクション・プラス)」のバージョンアップを実施し、施策の高度化・自動化を強化するレコメンデーション表示制御機能の提供を開始したことを発表いたします。 ■今回の主な新機能、バージョンアップの内容  ●レコメンデーション表示制御機能の強化  ●ABテスト機能の強化  企業のマーケティングにおいてパーソナライズされたコミュニケーションの重要性が増す中、デジタルチャネルにおいても顧客の状況・行動・属性に応じたきめ細やかな接客を実行することに対する期待が、さらに高まっています。

タスク②
# Instruction:
以下は、あなたへの依頼事項です。次のニュースのタイトルを考えて
# Input:
2022年4月1日(木)、ブレインパッドに、新卒社員数としては過去最多となる41名の仲間が入社しました! 実は、この数年の新卒入社者数は毎年「過去最多」を更新しているので、この冒頭の紹介文も定型になりつつあるのが嬉しい限りです。 ブレインパッドが新卒採用を始めた当初はデータサイエンティストだけだった募集コースも、今ではデータサイエンティスト、ビジネス、エンジニアの3コースに広がり、さまざまなバックグラウンドを持つ方々が入社してくれるようになりました。 それでは、今年入社した素敵な仲間達の様子を、入社式の写真と共にご紹介していきます。

タスク③
# Instruction:
以下は、あなたへの依頼事項です。次のニュースのタイトルを考えて
# Input:
野球チーム「ブレインパッド」のマルヤマ投手が「9番・投手」で先発出場し、投手としては5回を投げて被安打3、6奪三振・6四死球で3失点し3勝目を逃した。打者では3安打という活躍をみせ、チームの5連勝に大きく貢献した。投手としては粘りが光った。立ち上がりに制球を乱し、先頭打者に四球を与えると4番のサトウに2ランを被弾。5回に、四球と死球でピンチを招くと、3番フランスに同点タイムリーを許した。その後のピンチを併殺で切り抜け、6奪三振でリリーフ陣に後を託した。


各モデルの回答を以下に示します。

RWKV-4-Pile-430M(ファインチューニングのベースモデル)の回答

タスク①
# Response: ----
レコメンデーション表示の強化、ブラウザー効率化の強化、そしてWeb/アプリのテストの提供が進められています。今後とも、次の取り組みを展開していきます。
# 関連情報
# 動画配信:http://www.younow.com/product/487813
# データベース内容:http://www.younow.com/product/487814
# ニュースサイト:http://www.younow.com/news/143084
# プラットフォーム情報:https://marketplace.bazaarvoice.com/item/4eae2b0c-5d6a-4f5c-b1d8-7bfb8d5dcf8f
# ステップインの新機能・動画配信:https://web.younow.com/action+plus/487814
# プ

タスク②
# Response: ----
本日の新卒採用者は、現在進学を控えています。彼らは、学生にも選ばれており、その場合は、新卒入社理由の過去について、考えたいことがあります。 1. 今年1月31日(木)から30日(日)に、現在の仲間過去最多の40名が新卒入社式で行われました。 今年2月に入社した新卒採用者のうち、3人以上は、自分たちが嫌な人だったことを明かしました。 2. 新卒採用者の数は、新卒採用者数の多くを通して、全体的に合計20人以上です。 この30名が「初めて」に入社した時には、このページで「初めての仲間過去最多」と記載されていることをご覧ください。

タスク③
# Response: ----
打者の指先に向かって、先制5球を振り回し、2球を投げて3球を倒し、5球を失った。5球を投げて、4球を投げて、4球を投げて3球を投げて5球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて3球を投げて三ロームで一致する。5回に加え、4回にもピンチ


ファインチューニング済みモデルの回答

タスク①
# Response: ----
パーソナルコミュニケーションの流行に共通する理由【デジ通】

タスク②
# Response: ----
今年の入社式をチェックしよう! 写真は過去1人撮影の様子!

タスク③
# Response: ----
【Sports Watch】マルセル・フランス球界に聞いてみた


ベースモデルの回答とファインチューニング済みのモデルの回答を比較してみると、以下の観点で、ファインチューニングによって特定のタスクに適応している痕跡があるように感じました。

  • 「次のニュースのタイトルを考えて」というタスクに対して、ユーザーのニーズに合わせた明確で適切な回答をしていました。ニュースのタイトルとして文長や表現も適切であるように見えます。
  • スポーツ関連のニュースであることを判断して、【Sports Watch】という見出しをつけています。

一方で、「マルセル・フランス球界」など入力文章から読み取れない表現が出力されており改善の余地があることも感じました。

chat.pyのなかにFor better chat & QA quality: reduce temp, reduce top-p, increase repetition penaltiesという記載があったことから、これらのパラメーターを調整することで出力が改善される可能性があります。(なお、今回は以下のパラメーターで実行しています。)

GEN_TEMP = 1.1 # It could be a good idea to increase temp when top_p is low
GEN_TOP_P = 0.5 # Reduce top_p (to 0.5, 0.2, 0.1 etc.) for better Q&A accuracy (and less diversity)
GEN_alpha_presence = 0.2 # Presence Penalty
GEN_alpha_frequency = 0.2 # Frequency Penalty

また現在、Hugging FaceのスペースにはRWKV-4-Pile 1.5/3/7/14BモデルをベースにAlpacaデータセットなどでファインチューニングしたRavenモデルが続々と作成されています。日本語を学習データに使ったものも公開されており、これらのモデルをベースモデルとして使用することで上記の問題も解決するかもしれません。(これらのモデルはApache License 2.0となっていますが、Alpacaデータセットを使用しているため商用利用には注意が必要な認識です。)

終わりに

本記事では、RWKVをベースにしたチャットAIであるChatRWKVをGoogle Colaboratoryで実行する手順や、livedoor newsコーパスを活用してRWKVモデルをファインチューニングする方法について紹介しました。簡便にモデルを実行できることやファインチューニングが行えることが、ご確認いただけたと思います。この枠組みを利用して、特定のタスクに特化したモデルを作成することができると、ビジネス適用の可能性も広がりそうです。

参考文献

本記事は、以下を参考に調査・執筆を行いました。
huggingface.co
huggingface.co
github.com
github.com
www.rondhuit.com
note.com
note.com
qiita.com