RWKV アーキテクチャの Hugging Face Transformers への統合

RWKV アーキテクチャの Hugging Face Transformers への統合

Hugging Face は、RWKV RNN‑Transformer ハイブリッドモデルを Transformers ライブラリに統合したことを発表しました。これにより、RNN の効率性と Transformer の性能を組み合わせた、オープンソースの長文コンテキスト対応言語モデルが利用可能になります。

RWKV アーキテクチャの Hugging Face Transformers への統合

TL;DR

RWKV は、Transformer のアテンションを模倣した新しい RNN ベースのアーキテクチャで、現在 Hugging Face の transformers ライブラリで正式にサポートされています。これにより、開発者は非常に長いコンテキストを RNN の速度とメモリ効率で処理できるオープンソースの言語モデルにアクセスできます。

RWKV プロジェクトの概要

RWKV プロジェクトは Bo Peng (GitHub: BlinkDL) が主導し、活発な Discord コミュニティによって維持されています。Stability AI がトレーニングに使用された GPU を寄付しました。プロジェクトのロードマップには、パフォーマンス向上（例：RWKV.cpp、量子化）、スケーラビリティ強化（データセット処理）、チャットファインチューニングやマルチモーダルファインチューニングといった研究拡張が含まれます。コミュニティメンバーは公式 Discord チャンネルに参加して貢献できます。

RWKV が RNN と Transformer を橋渡しする方法

RNN の制限と Transformer の利点

  • 従来の RNN は各時間ステップで同じ重みを再利用するため、勾配消失問題や長距離記憶の不足を招きます。LSTM と GRU は部分的にこれを緩和しますが、依然として非常に長いシーケンスに苦労します。
  • Transformer は自己注意機構を用いてすべてのトークンを並列に処理し、クエリ、キー、バリューの射影で注意スコアを計算します。この設計により長距離依存性の問題が解決され、従来の RNN に比べて学習が高速化されます。

RWKV のハイブリッド設計

  • RWKV は Apple の Attention‑Free Transformer に触発され、RNN 互換の形に簡素化されています。
  • Transformer スタイルの埋め込み、レイヤーノーマライゼーション、因果言語モデルヘッドは保持しつつ、注意層を再帰ベースの定式化に置き換え、自己注意と同等の表現力を実現します。
  • TokenShift や SmallInitEmb といった追加のテクニック（公式 GitHub README に記載）を使用することで、モデルは GPT レベルの性能に匹敵します。

RWKV アーキテクチャの技術的ハイライト

長コンテキスト対応能力

  • RWKV は 8 192 トークン（ctx8192）のコンテキストウィンドウを、1 024 トークンモデルと同等の推論速度と RAM 使用量で処理できます。
  • 実証的なロス曲線は、コンテキスト長が大きくなるほどモデルサイズに関係なく言語モデルのロスが改善され、効果的な長距離記憶が実現されていることを示しています。

学習効率

  • 従来の RNN とは異なり、RWKV は「線形化 GPT」方式で学習でき、バッチ間の並列処理と従来の再帰モデルよりも高速な収束を可能にします。
  • 現在の学習パイプラインは 14 B パラメータまでスケールし、RWKV‑4 系列の数値安定性に関する修正が進行中です。

利用可能なモデルチェックポイント

純粋な言語モデル (RWKV‑4)

  • モデルサイズは約 170 M から 14 B パラメータまでです。
  • すべてのモデルは The Pile データセットで事前学習され、最先端ベンチマークと比較して同等の性能を示しています。

指示に基づくファインチューニング済みチャットモデル (RWKV‑4 Raven)

  • Raven シリーズは ALPACA、CodeAlpaca、Guanaco、GPT‑4All、ShareGPT などの指示データセットで RWKV‑4 をファインチューニングします。
  • 言語構成（英語のみ、英語＋中国語＋日本語など）やサイズ（1.5 B、7 B、14 B）に応じたバリエーションがあります。
  • すべてのチェックポイントは Hugging Face Hub の RWKV 組織の下でホストされています。

🤗 Transformers で RWKV を使用する

テキスト生成の例

from transformers import pipeline
model_id = "RWKV/rwkv-4-169m-pile"
pipe = pipeline("text-generation", model=model_id)
print(pipe("In a shocking finding, scientist discovered a herd of dragons...", max_new_tokens=20))

このパイプラインは、Transformer ベースの生成器に匹敵する一貫した続きのテキストを返します。

チャットモデル (Raven) の例

from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "RWKV/rwkv-raven-1b5"
model = AutoModelForCausalLM.from_pretrained(model_id).to(0)
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "### Instruction: Tell me about ravens\n### Response:"
inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=100)
print(tokenizer.decode(output[0], skip_special_tokens=True))

このモデルは Alpaca スタイルの指示フォーマットに従い、詳細な応答を生成します。

元の RWKV 重みを Hugging Face 形式に変換する

convert_rwkv_checkpoint_to_hf.py という変換スクリプトが transformers リポジトリに同梱されています。ユーザーは生のチェックポイントを Hub リポジトリにアップロードし、次のコマンドを実行します：

python convert_rwkv_checkpoint_to_hf.py \
  --repo_id RAW_HUB_REPO \
  --checkpoint_file RAW_FILE \
  --output_dir OUTPUT_DIR

--push_to_hub と --model_name を追加すると、変換されたモデルが直接 Hub にアップロードされます。

今後の方向性

  • Multilingual RWKV – 多言語コーパスとトークナイザーの開発が進行中で、モデルの言語カバレッジが拡大します。
  • Community research – Discord チャンネルでは、新しい学習レシピ、ベンチマーク、アーキテクチャの調整に関するプロジェクトが開催されています。
  • Compression & acceleration – RWKV は行列‑ベクトル演算のみを使用するため、量子化（4‑bit/8‑bit）、ONNX エクスポート、光子アクセラレータなどの実験的ハードウェアに適しています。optimum ライブラリや rwkv.cpp、rwkv-cpp-cuda といったリポジトリとの統合により、推論がさらに高速化されます。

謝辞

Hugging Face チームは Bo Peng、RWKV コミュニティ、そして Johan Wind（RWKV ブログ記事）、ArEnSc（初期 Transformers PR）、Merve Noyan、Maria Khalusova、Pedro Cuenca などの貢献者に、統合のレビューとサポートに感謝します。

引用

研究で RWKV を使用する場合は、RWKV‑LM リポジトリにある CITATION.cff ファイルを用いてプロジェクトを引用してください。

Sources