RWKV（Receptance Weighted Key Value）をつかってみた

はじめに

こんにちは。ブレインパッドの丸山です。
最近、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つの質問を投げかけてみます。

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

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

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

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

• https://qiita.com/toshi_456/items/280efc31950ddb083286
• https://note.com/npaka/n/n91e6dfecd034

今回学習用データセットには、日本のさまざまなニュース記事から構成されるデータセットであり、多様なトピックや文体を含んでいる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つで、このような形式でプロンプトとして入力しています。

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

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

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

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

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