ChatGPTに100ページの論文を要約させたい。リアルタイムでストリーミング音声を処理したい。しかしTransformerのSelf-Attentionは系列長 $n$ に対して $O(n^2)$ の計算量を要求し、系列長を倍にするたびにコストは4倍に膨れ上がります。一方で、RNNは $O(n)$ の計算量で系列を処理できますが、訓練時に並列化できず、GPUの性能を引き出せません。
「RNNの推論効率」と「Transformerの訓練並列性」を同時に手に入れることはできないのでしょうか？ この問いに対して、正面から解を与えたのがRWKV（Receptance Weighted Key Value）です。RWKVは、Attentionを使わずにTransformerと同等の性能を達成しながら、推論時は $O(1)$ メモリ（1ステップあたり）、訓練時はTransformerと同様に全トークンを並列計算できるアーキテクチャです。
RWKVの理論を理解すると、以下の分野に直接応用できます。
- エッジデバイスでの推論: スマートフォンや組み込みデバイスのような限られたメモリ環境で、長いコンテキストを保持しながらリアルタイム推論を実現する
- 超長系列モデリング: 数十万トークンの文書、ゲノム配列、長時間の時系列データを、メモリを爆発させずに処理する
- ストリーミング処理: RNNモードで1トークンずつ逐次処理できるため、音声認識やチャットボットでのレイテンシ削減に直結する
- 新しいアーキテクチャ設計のヒント: Attentionの代替手法を理解することで、MambaやRetNetなど次世代モデルの設計思想も見えてくる
この図がRWKVの立ち位置を一枚で表しています。横軸は「訓練の並列性」、縦軸は「推論メモリ効率」です。RNN（左上）は推論はO(1)で軽いが訓練を並列化できず、Transformer（右下）は訓練は並列だが推論のKVキャッシュが系列長に比例して膨らみます。RWKV（右上）は両者の良いところだけを継承し、訓練は並列・推論はO(1)メモリという2つの長所を同時に手に入れます。本記事はこの「両取り」がどんな仕組みで実現されるのかを、数式とコードで一段ずつ解きほぐしていきます。
本記事の内容
- TransformerとRNNのトレードオフ — なぜ両方の長所が必要なのか
- RWKVの名前の由来 — R（Receptance）、W（Weight）、K（Key）、V（Value）の役割
- WKV機構の数学的定式化 — 時間混合（Time Mixing）の全導出
- チャネル混合（Channel Mixing）の定式化
- Token Shift操作 — 現在と過去の情報を混ぜる仕組み
- 訓練時の並列化 — RNN形式とTransformer形式の双対性
- 推論時の $O(1)$ メモリ — 再帰的な状態更新
- RWKV v4 / v5 / v6 の進化
- Pythonでの簡易RWKV実装（WKV機構）
- Transformer / Mamba / RWKV の比較
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
TransformerとRNNのトレードオフ
並列性と逐次性のジレンマ
RWKVの設計を理解するために、まずTransformerとRNNの根本的なトレードオフを整理しましょう。
TransformerのSelf-Attentionは、全トークンペアの関係を一度に計算します。系列 $(x_1, x_2, \dots, x_n)$ が与えられたとき、$n \times n$ の注意行列を構成するため、計算量は $O(n^2 d)$（$d$ は特徴次元）です。しかし、この行列演算はGPU上で高度に並列化できるため、訓練時の効率は非常に高いです。問題は推論時に現れます。新しいトークンを1つ生成するたびに、過去の全トークンとの注意を再計算するか、KVキャッシュを保持する必要があります。KVキャッシュのメモリ量は系列長に比例して増え続け、長い対話では数十GBに達することもあります。
RNN（LSTM、GRU含む）は、隠れ状態 $\bm{h}_t$ を逐次的に更新する再帰モデルです。
$$ \bm{h}_t = f(\bm{h}_{t-1}, x_t) $$
1ステップあたりの計算は $O(d^2)$ で系列長に依存せず、メモリも隠れ状態のサイズ $O(d)$ で固定です。推論時の効率は理想的です。しかし、訓練時には各ステップが前のステップの結果に依存するため、系列を順番に処理するしかなく、GPUの並列性能を活かせません。さらに、長い系列では勾配消失・爆発が起こり、遠くの依存関係を学習しにくいという問題もあります。
理想のモデル
この状況を整理すると、私たちが求めるモデルの条件は次のようになります。
| 要件 | Transformer | RNN | 理想のモデル |
|---|---|---|---|
| 訓練の並列化 | $\checkmark$ | $\times$ | $\checkmark$ |
| 推論メモリ | $O(n)$（KVキャッシュ） | $O(1)$ | $O(1)$ |
| 推論の1ステップ計算量 | $O(n)$ | $O(1)$ | $O(1)$ |
| 長距離依存性 | $\checkmark$（全トークン参照） | 困難 | $\checkmark$ |
Transformerの左列とRNNの中列を見ると、それぞれ正反対の長所を持っていることがわかります。RWKVは、この両方の長所を同時に実現することを目指して設計されました。
では、RWKVがこの理想をどのように実現するのか、アーキテクチャの核心に入っていきましょう。
RWKVの名前の由来 — 4つの構成要素
Receptance, Weight, Key, Value
RWKVという名前は、このモデルの4つの中核的な概念の頭文字から取られています。TransformerのQ（Query）、K（Key）、V（Value）と対比すると、RWKVの設計思想がよく見えてきます。
Transformerでは「Query がKeyに問い合わせ、Valueを引き出す」という検索メカニズムがAttentionの本質でした。RWKVではQueryをなくし、代わりにReceptance（受容度）という新しい概念を導入します。
- R（Receptance）: 現在のトークンが過去の情報をどの程度「受け入れる」かを制御するゲート。値域は $[0, 1]$（sigmoid関数を通すため）。TransformerのQueryに相当する「検索」の役割を担いますが、softmaxベースの注意ではなく、要素ごとのゲーティングとして機能します
- W（Weight Decay）: 時間方向の減衰係数。過去のトークンの影響が時間とともにどれだけ薄れるかを制御します。チャネルごとに学習される実数値パラメータで、$e^{-w}$（$w > 0$）の形で指数減衰として働きます。RNNの「忘却」に相当します
- K（Key）: 各トークンの「特徴ラベル」。Transformerと同じ名前ですが、全トークンペアの類似度を計算するのではなく、WKV機構の中で時間減衰と組み合わせて使われます
- V（Value）: 各トークンが持つ「情報の中身」。最終的にKとWによる重み付けで集約され、Rでゲーティングされて出力になります
4つの要素の役割をカードに整理しました。R（Receptance）はsigmoidで$[0,1]$のゲートを作り「過去情報をどれだけ受け入れるか」を決めます。W（Weight）は$e^{-w}$の指数減衰で「過去がどれだけ薄れるか」を、K（Key）は$e^{k}$で各トークンの重要度を、V（Value）は実際に運ばれる情報の中身を担います。TransformerのQueryに当たる「検索」の役割をRが引き受け、全トークンペアの内積計算を一切しない点が最大の違いです。
直感的に言えば、Transformerが「誰に聞くべきか（Query→Key）」を全トークンに同時に問い合わせるのに対し、RWKVは「過去の情報を指数的に忘れながら蓄積し（W, K, V）、現在のトークンがそれをどれだけ受け入れるか（R）で出力を決める」という仕組みです。図書館のアナロジーで言えば、Transformerは全ての本の棚を同時にスキャンしますが、RWKVは最近見た本ほど鮮明に覚えている司書が、あなたのリクエスト（R）に応じて記憶から情報を引き出す、というイメージです。
この4つの要素がどのように数式で組み合わされるのか、まずはRWKVブロックの全体構造を見てから、核心であるWKV機構の定式化に入りましょう。
RWKVブロックの全体構造
2つのサブブロック
RWKVの各層は、TransformerのEncoder層と似た構造を持っています。Transformerの各層が「Self-Attention + FFN」で構成されるように、RWKVの各層は以下の2つのサブブロックで構成されます。
- Time Mixing（時間混合）— Transformerの Self-Attention に相当。トークン間の情報交換を行う
- Channel Mixing（チャネル混合）— TransformerのFFN（Feed-Forward Network）に相当。チャネル方向の非線形変換を行う
各サブブロックの前にはLayerNormが適用され、残差接続（Residual Connection）で入力が加算されます。数式で書くと、$l$ 番目の層の計算は次のようになります。
$$ \bm{x}_{l}’ = \bm{x}_{l-1} + \text{TimeMixing}(\text{LayerNorm}(\bm{x}_{l-1})) $$
$$ \bm{x}_{l} = \bm{x}_{l}’ + \text{ChannelMixing}(\text{LayerNorm}(\bm{x}_{l}’)) $$
1層の中身を図にすると、入力から出力までの流れが見えてきます。LayerNorm → Time Mixing → 残差接続、続いてLayerNorm → Channel Mixing → 残差接続という2段構成です。Time MixingがSelf-Attentionに、Channel MixingがFFNに対応し、どちらもPre-Norm型で残差接続（破線の迂回路）が付いています。TransformerのEncoder層と骨格がそっくりで、違いは「Attentionの中身がWKV機構に置き換わっている」点だけだと分かります。
この構造はTransformerの「Pre-Norm」型と同じです。残差接続のおかげで勾配が層をまたいで直接伝播し、深いモデルでも安定して学習できます。
重要なのは、Time MixingとChannel Mixingのどちらも、現在のトークンと1つ前のトークンだけを参照するという点です。「え、たった1つ前しか見ないのに長距離依存性を学習できるの？」と疑問に思うかもしれませんが、これは各層が前の層の出力を入力とする深いネットワーク構造によって、間接的に遠くの情報を伝播させています。また、WKV機構の中で指数減衰による過去の情報の蓄積が行われるため、明示的に全トークンを参照しなくても長い依存関係を捉えることができます。
それでは、RWKVの心臓部であるTime Mixing（WKV機構）の数学的な定式化に入りましょう。
Token Shift — 現在と過去を混ぜる操作
なぜToken Shiftが必要なのか
Time MixingとChannel Mixingの定式化に入る前に、両方で共通して使われるToken Shift（トークンシフト）という操作を理解しておく必要があります。
Transformerでは、Self-Attentionにおいて全トークンの情報が注意行列を通じて混ざり合います。そのため、各トークンの入力ベクトルがそのままQ, K, Vの射影に使われても問題ありません。しかしRWKVでは、R, W, K, Vの計算は各トークン位置で独立に（Attentionなしで）行われます。もし入力ベクトル $\bm{x}_t$ だけを使ってR, K, Vを計算すると、現在のトークンの情報しか反映されず、文脈をまったく考慮できないことになってしまいます。
そこでRWKVは、現在のトークンの入力と1つ前のトークンの入力を線形補間することで、最小限の文脈情報を注入します。これがToken Shiftです。
数式の定義
時刻 $t$ のトークンの入力ベクトルを $\bm{x}_t \in \mathbb{R}^d$ とします。Token Shiftは、学習可能な混合係数 $\bm{\mu} \in \mathbb{R}^d$（チャネルごとに異なる値）を使って、次のように定義されます。
$$ \bar{\bm{x}}_t = \bm{\mu} \odot \bm{x}_t + (1 – \bm{\mu}) \odot \bm{x}_{t-1} $$
ここで $\odot$ は要素ごとの積（Hadamard積）です。混合係数 $\bm{\mu}$ は各チャネルで独立に学習され、$[0, 1]$ の範囲に制約されます。$\mu_i = 1$ のチャネルでは現在のトークンの情報だけが使われ、$\mu_i = 0$ のチャネルでは1つ前のトークンの情報だけが使われます。中間の値では両方の情報がブレンドされます。
実際には、R, K, V の計算にそれぞれ異なる混合係数 $\bm{\mu}_r, \bm{\mu}_k, \bm{\mu}_v$ が使われます。これにより、「受容度（R）の計算では現在のトークンを重視するが、Key（K）の計算では前のトークンの情報を多く取り入れる」といった柔軟な制御が可能になります。
$$ \bar{\bm{x}}_t^{(r)} = \bm{\mu}_r \odot \bm{x}_t + (1 – \bm{\mu}_r) \odot \bm{x}_{t-1} $$
$$ \bar{\bm{x}}_t^{(k)} = \bm{\mu}_k \odot \bm{x}_t + (1 – \bm{\mu}_k) \odot \bm{x}_{t-1} $$
$$ \bar{\bm{x}}_t^{(v)} = \bm{\mu}_v \odot \bm{x}_t + (1 – \bm{\mu}_v) \odot \bm{x}_{t-1} $$
左の模式図がToken Shiftの仕組みです。現トークン$x_t$と前トークン$x_{t-1}$を混合係数$\mu$で線形補間し、その結果$\bar{x}_t$をR, K, Vの計算に渡します。右の積み上げ棒グラフは、$\mu$の値によって現在と過去の混合比がどう変わるかを示しています。$\mu=1$なら現トークンのみ、$\mu=0$なら前トークンのみ、中間では両者がブレンドされます。$\mu$はチャネルごとに独立に学習されるため、あるチャネルは現在を重視し、別のチャネルは過去を多く取り込むという使い分けが自然に生まれます。
この操作は計算量が $O(d)$ で極めて軽量です。1つ前のトークンの入力を保持するだけでよいため、メモリ増加もごくわずかです。それでいて、「現在のトークンだけ見る」のと「前のトークンの文脈も考慮する」のとでは、モデルの表現力に大きな差が生まれます。
Token Shiftの仕組みが分かったところで、いよいよRWKVの核心であるWKV機構（Time Mixing）の数式を見ていきましょう。
WKV機構 — Time Mixingの数学的定式化
直感的なイメージ
WKV機構は、Transformerの Self-Attention を「指数減衰付きの重み付け平均」で置き換えたものです。
Self-Attentionでは、現在のトークンのQueryと全トークンのKeyの内積をsoftmaxで正規化し、Valueの重み付け平均を取りました。これにより「どのトークンに注目するか」が動的に決まりますが、全ペアの計算が必要なため $O(n^2)$ のコストがかかりました。
WKV機構のアイデアはシンプルです。「注目するかどうか」を全トークンペアについて計算する代わりに、過去のトークンの寄与を時間とともに指数的に減衰させることで、自然と最近のトークンに重みが集中するようにします。これは「最近の出来事ほどよく覚えている」という人間の記憶の仕組みにも似ています。
最後のトークンから過去を振り返ったときの重み分布を、減衰率$w$を変えて描いたものです。$w=1.5$のように大きいと直近の数トークンだけに鋭く集中し、$w=0.1$のように小さいと遠い過去まで緩やかに広がります。Transformerの注意がデータに応じて任意の形を取れるのに対し、RWKVの重みは「指数減衰」という決まった形に縛られている点が対照的です。$w$をチャネルごとに学習させることで、短期記憶のチャネルと長期記憶のチャネルが自然に分業します。
数式の定式化（RWKV-4）
Token Shift後の入力 $\bar{\bm{x}}_t^{(r)}, \bar{\bm{x}}_t^{(k)}, \bar{\bm{x}}_t^{(v)}$ から、まず線形射影で R, K, V を計算します。
$$ \bm{r}_t = \bm{W}_r \bar{\bm{x}}_t^{(r)}, \quad \bm{k}_t = \bm{W}_k \bar{\bm{x}}_t^{(k)}, \quad \bm{v}_t = \bm{W}_v \bar{\bm{x}}_t^{(v)} $$
ここで $\bm{W}_r, \bm{W}_k, \bm{W}_v \in \mathbb{R}^{d \times d}$ は学習可能な重み行列です。
次に、WKV機構のコアとなる演算を定義します。時刻 $t$ におけるチャネル $i$ のWKV出力は、次のように計算されます。
$$ \text{wkv}_{t,i} = \frac{e^{u_i + k_{t,i}} v_{t,i} + \sum_{j=1}^{t-1} e^{-(t-1-j)w_i + k_{j,i}} v_{j,i}}{e^{u_i + k_{t,i}} + \sum_{j=1}^{t-1} e^{-(t-1-j)w_i + k_{j,i}}} $$
この式を直感的に読み解きましょう。
分子は、Valueの重み付き和です。各トークン $j$ のValue $v_{j,i}$ に重み $e^{-(t-1-j)w_i + k_{j,i}}$ が掛けられています。この重みは2つの要素で決まります。
- 指数減衰 $e^{-(t-1-j)w_i}$: トークン $j$ が現在時刻 $t$ からどれだけ離れているかに応じて、指数的に減衰します。$w_i > 0$ のとき、遠い過去のトークンほど重みが小さくなります
- キーの寄与 $e^{k_{j,i}}$: トークン $j$ のKeyの値が大きいほど、そのトークンの重要度が高くなります
分母は正規化係数で、重みの総和です。softmaxと同じ役割を果たし、出力がValueの加重平均になることを保証します。
現在のトークンの特別扱い: 現在のトークン $t$ には減衰 $w_i$ ではなく、別のパラメータ $u_i$（ボーナス項）が使われています。これは「現在のトークン自身の情報は減衰させるべきではない」という設計意図を反映しています。$u_i$ は学習可能なパラメータで、現在のトークンの重要度をモデルが調整できるようにしています。
WKVの重みを行列として可視化すると、その構造がはっきり見えます。行が現在時刻$t$、列が過去のKey位置$j$で、上三角がゼロ（未来は見ない）の因果的な下三角行列になっています。対角要素（現在のトークン自身）が最も明るく、左へ進むほど指数的に暗くなる — これがTransformerのMasked Attentionと同じ下三角構造を持ちつつ、中身が指数減衰で固定されているというRWKVの本質です。後で見るように、この下三角構造こそが訓練時の並列化を可能にします。
再帰形式への書き換え
上の式をナイーブに計算すると、各時刻 $t$ で過去の全トークンを走査する必要があり、全体で $O(n^2)$ になってしまいます。しかし、指数減衰の構造を利用すると、再帰的に $O(1)$ で1ステップを計算できます。
2つの累積変数を導入します。
$$ a_{t,i} = e^{-w_i} a_{t-1,i} + e^{k_{t,i}} v_{t,i} $$
$$ b_{t,i} = e^{-w_i} b_{t-1,i} + e^{k_{t,i}} $$
$a_{t,i}$ は重み付きValueの累積和、$b_{t,i}$ は重みの累積和です。前のステップの値に減衰 $e^{-w_i}$ を掛けてから、現在のステップの寄与を加えるという、単純な更新式になっています。
この再帰更新を時間軸に沿って図にしました。状態$(a_t, b_t)$は時刻を進むたびに「前の状態に$e^{-w}$を掛けて、現ステップのKey・Valueの寄与を足す」だけで更新されます。過去の全トークンを走査する必要はなく、保持するのは$d$次元のベクトル2本だけ。だから1ステップあたりのメモリは系列長$n$に一切依存せず$O(d)$で済みます。RNNの「隠れ状態を1つだけ持ち回る」効率を、指数減衰の構造で実現しているわけです。
この累積変数を使うと、次の時刻のWKV出力は次のように書けます。
$$ \text{wkv}_{t+1,i} = \frac{e^{u_i + k_{t+1,i}} v_{t+1,i} + a_{t,i}}{e^{u_i + k_{t+1,i}} + b_{t,i}} $$
現在のトークンのボーナス付き寄与 $e^{u_i + k_{t+1,i}} v_{t+1,i}$ と、過去のトークンの累積寄与 $a_{t,i}$ を足して正規化するだけです。累積変数 $a_t, b_t$ はそれぞれ $d$ 次元のベクトルなので、1ステップあたりのメモリは $O(d)$ で済みます。系列長 $n$ に依存しないのがポイントです。
数値的安定性
実装上の注意点があります。$e^{k_{j,i}}$ は $k$ の値が大きいとオーバーフローを起こします。これを避けるために、最大値を引いてからexpを計算するlog-sum-expトリックを使います。
累積変数を対数領域で保持する方法を見てみましょう。$p_{t,i} = \max(k_{t,i}, q_{t-1,i})$ として、
$$ a_{t,i} = e^{q_{t-1,i} – p_{t,i}} \cdot a’_{t-1,i} + e^{k_{t,i} – p_{t,i}} \cdot v_{t,i} $$
ここで $q_{t,i}$ は累積指数の最大値を追跡する変数です。このテクニックにより、数値が常に $[0, 1]$ 近辺に保たれ、float32でもfloat16でも安定した計算が可能になります。
Time Mixingの出力
WKV機構の出力にReceptance（受容度）のゲーティングを適用して、最終的なTime Mixingの出力を得ます。
$$ \bm{o}_t = \sigma(\bm{r}_t) \odot \text{wkv}_t $$
ここで $\sigma$ はsigmoid関数です。$\sigma(\bm{r}_t)$ は各チャネルで $[0, 1]$ の値を取り、WKV出力をどれだけ「受け入れるか」を制御します。Receptanceが0に近いチャネルではWKVの情報がブロックされ、1に近いチャネルではそのまま通過します。これはLSTMの出力ゲートに似た役割を果たしています。
最終的に、線形射影で出力次元に変換します。
$$ \text{TimeMixing}(\bm{x}_t) = \bm{W}_o \bm{o}_t $$
Time Mixingの仕組みが理解できたところで、次にチャネル方向の情報変換を担うChannel Mixingの定式化を見ていきましょう。
Channel Mixing — チャネル方向の非線形変換
TransformerのFFNとの対応
TransformerのFeed-Forward Network（FFN）は、Attention層の後にチャネル方向（特徴次元方向）の非線形変換を行い、各トークン位置の表現を豊かにする役割を果たします。RWKVのChannel Mixingはこれに対応するコンポーネントですが、Token ShiftとReceptanceゲーティングが組み込まれている点が異なります。
数式の定義
Channel Mixingでも、まずToken Shiftで時間方向の情報を混ぜます。ただし、Time Mixingとは別の混合係数 $\bm{\mu}_r’, \bm{\mu}_k’$ を使います。
$$ \bar{\bm{x}}_t^{(r’)} = \bm{\mu}_r’ \odot \bm{x}_t + (1 – \bm{\mu}_r’) \odot \bm{x}_{t-1} $$
$$ \bar{\bm{x}}_t^{(k’)} = \bm{\mu}_k’ \odot \bm{x}_t + (1 – \bm{\mu}_k’) \odot \bm{x}_{t-1} $$
次に、Key射影にsquared ReLU活性化関数を適用して中間表現を計算し、Receptanceでゲーティングします。
$$ \bm{r}_t’ = \sigma(\bm{W}_{r’} \bar{\bm{x}}_t^{(r’)}) $$
$$ \bm{k}_t’ = \text{sqReLU}(\bm{W}_{k’} \bar{\bm{x}}_t^{(k’)}) = \left(\max(0, \bm{W}_{k’} \bar{\bm{x}}_t^{(k’)})\right)^2 $$
$$ \text{ChannelMixing}(\bm{x}_t) = \bm{W}_{v’} (\bm{r}_t’ \odot \bm{k}_t’) $$
ここで $\bm{W}_{r’} \in \mathbb{R}^{d \times d}$、$\bm{W}_{k’} \in \mathbb{R}^{d_{\text{ff}} \times d}$、$\bm{W}_{v’} \in \mathbb{R}^{d \times d_{\text{ff}}}$ です。$d_{\text{ff}}$ は中間次元で、Transformerと同様に $d_{\text{ff}} = 4d$ が典型的です。
squared ReLU $(\max(0, x))^2$ は通常のReLUより疎な（多くの要素がゼロになる）活性化を生み出し、かつ微分がゼロ点で連続であるため学習が安定します。Transformerの FFN で使われるGELUやSwiGLUの代替として機能します。
Channel Mixingの構造を整理すると、次のようなデータフローになります。
- Token Shiftで前のトークンの情報を混入
- Key射影 + squared ReLUで中間表現を生成（非線形変換）
- Receptance射影 + sigmoidでゲートを生成
- ゲートと中間表現の要素積を取り、Value射影で出力次元に戻す
TransformerのFFNとの最大の違いは、Token ShiftとReceptanceゲーティングが加わっている点です。Token Shiftにより、チャネル混合の段階でも隣接トークンの文脈が考慮されます。Receptanceゲーティングにより、各チャネルの情報を選択的に通過させるフィルタリングが働きます。
これでRWKVブロックの2つのサブブロック — Time MixingとChannel Mixing — の定式化が完了しました。次に、RWKVの最も巧妙な性質である「訓練時の並列化」について詳しく見ていきましょう。
訓練時の並列化 — RNN形式とTransformer形式の双対性
2つの顔を持つモデル
RWKVの最も革新的な性質は、同じモデルをRNN形式でもTransformer形式でも計算できるという双対性です。これは「数学的に等価な2つの計算方法がある」ということであり、状況に応じて使い分けることで、訓練と推論の両方で最適な効率を実現します。
前節で見た再帰形式をもう一度書きます。
$$ a_{t,i} = e^{-w_i} a_{t-1,i} + e^{k_{t,i}} v_{t,i} $$
$$ b_{t,i} = e^{-w_i} b_{t-1,i} + e^{k_{t,i}} $$
この再帰式を展開すると、$a_{t,i}$ は次のように書けます。
$$ a_{t,i} = \sum_{j=1}^{t} e^{-(t-j)w_i + k_{j,i}} v_{j,i} $$
ここで指数の和の操作 $e^{-(t-j)w_i}$ は $j$ ごとに独立な係数であることに注目してください。つまり、$e^{k_{j,i}} v_{j,i}$ と $e^{-(t-j)w_i}$ の組み合わせは、全ての $(t, j)$ ペアについて独立に計算可能です。
並列計算の実現
全時刻の WKV 出力を一度に計算するには、次の行列演算として定式化できます。
重み行列 $\bm{W}_{\text{decay}} \in \mathbb{R}^{n \times n}$ を定義します（チャネル $i$ について）。
$$ [\bm{W}_{\text{decay}}]_{t,j} = \begin{cases} e^{-(t-1-j)w_i} & \text{if } j < t \\ e^{u_i} & \text{if } j = t \\ 0 & \text{if } j > t \end{cases} $$
この行列は下三角行列（因果マスク付き）であり、TransformerのMasked Attentionと同じ構造を持っています。$e^{k_{j,i}}$ を対角行列 $\text{diag}(e^{\bm{k}_i})$ として掛けると、全時刻のWKV出力を次の行列積で一括計算できます。
$$ \text{wkv}_i = \text{rowwise\_normalize}\left(\bm{W}_{\text{decay}} \cdot \text{diag}(e^{\bm{k}_i}) \cdot \bm{V}_i\right) $$
ここで rowwise_normalize は各行を対応する重みの和で割る操作です。この行列積はGPU上で高度に並列化でき、cuBLASやFlash Attentionと同様の最適化が適用できます。
計算量の比較
| モード | 1ステップ計算量 | 全系列計算量 | メモリ |
|---|---|---|---|
| RNN形式（再帰） | $O(d)$ | $O(nd)$ | $O(d)$ |
| Transformer形式（並列） | — | $O(nd)$（行列演算） | $O(n + d)$ |
注目すべきは、どちらの形式でも全系列の計算量が $O(nd)$ で線形であるということです。Transformerの $O(n^2 d)$ と比較すると、系列長 $n$ に関して1桁少ないオーダーです。
ただし、Transformer形式での実装では $n \times n$ の減衰行列を明示的に構成する必要があるため、メモリが $O(n^2)$ になるように見えます。実際の実装では、CUDAカスタムカーネルを使って減衰行列を陽に構成せずに計算を行い、$O(n)$ のメモリで並列計算を実現しています。
この双対性を1枚にまとめたのが上の図です。左の訓練時は、下三角の減衰行列を一括の行列演算で処理するTransformer形式 — 全時刻を同時に計算できるのでGPUの並列性能をフルに使えます。右の推論時は、状態を1トークンずつ更新していくRNN形式 — メモリは$O(1)$で固定です。重要なのは、この2つがまったく同じパラメータで数学的に等価だということ。学習は左、生成は右と、用途に応じて計算形式を切り替えるだけでよいのです。
この双対性のおかげで、RWKVは訓練時にはGPUの並列性能をフル活用しながら、推論時にはRNNとして $O(1)$ のメモリで効率的に動作できます。同じ学習済みパラメータを使って、モードを切り替えるだけでよいのです。
では、推論時のRNNモードがどのように動作するのか、具体的な状態更新の流れを見ていきましょう。
推論時の $O(1)$ メモリ — 再帰的な状態更新
Transformerの推論問題
Transformerで自己回帰的にテキストを生成するとき、各ステップで過去の全トークンのKey・Valueを保持する必要があります。これがKVキャッシュです。系列長が $n$ トークンに達すると、KVキャッシュのメモリは各層で $O(nd)$ となり、長い対話や文書生成では数十GBに膨れ上がります。
さらに、新しいトークンを生成するたびに、そのトークンのQueryと過去の全KeyのAttentionを計算する必要があるため、1ステップあたりの計算量も $O(nd)$ です。
RWKVのRNNモード推論
RWKVの推論は根本的に異なります。各層で保持すべき状態は、累積変数 $a_{t,i}$ と $b_{t,i}$（チャネルごとにスカラー2つ）だけです。
新しいトークン $x_{t+1}$ が入力されたとき、以下の手順で出力を計算します。
ステップ1: Token Shift
前のトークンの入力 $\bm{x}_t$ を保持しておき、混合します。
$$ \bar{\bm{x}}_{t+1} = \bm{\mu} \odot \bm{x}_{t+1} + (1 – \bm{\mu}) \odot \bm{x}_t $$
ステップ2: R, K, V の計算
$$ \bm{r}_{t+1} = \bm{W}_r \bar{\bm{x}}_{t+1}^{(r)}, \quad \bm{k}_{t+1} = \bm{W}_k \bar{\bm{x}}_{t+1}^{(k)}, \quad \bm{v}_{t+1} = \bm{W}_v \bar{\bm{x}}_{t+1}^{(v)} $$
ステップ3: WKV出力の計算と状態更新
各チャネル $i$ について、
$$ \text{wkv}_{t+1,i} = \frac{e^{u_i + k_{t+1,i}} v_{t+1,i} + a_{t,i}}{e^{u_i + k_{t+1,i}} + b_{t,i}} $$
状態を更新します。
$$ a_{t+1,i} = e^{-w_i} a_{t,i} + e^{k_{t+1,i}} v_{t+1,i} $$
$$ b_{t+1,i} = e^{-w_i} b_{t,i} + e^{k_{t+1,i}} $$
ステップ4: Receptanceゲーティングと出力
$$ \bm{o}_{t+1} = \sigma(\bm{r}_{t+1}) \odot \text{wkv}_{t+1} $$
全ての操作はベクトル演算であり、系列長 $n$ に依存しません。1ステップあたりの計算量は $O(d^2)$（行列-ベクトル積が支配的）、保持すべき状態のメモリは $O(Ld)$（$L$ は層数）です。
Transformerとのメモリ比較
具体的な数値で比較してみましょう。$d = 4096$、$L = 32$ のモデルで $n = 100{,}000$ トークンを生成する場合を考えます。
Transformer（KVキャッシュ）:
KVキャッシュのメモリは各層でKey・Valueそれぞれ $n \times d$ の行列を保持するため、
$$ \text{メモリ} = 2 \times L \times n \times d \times 2 \text{bytes} \approx 2 \times 32 \times 100{,}000 \times 4{,}096 \times 2 \approx 50 \text{ GB} $$
（float16を仮定）
RWKV（RNNモード）:
各層で保持するのは $a_t \in \mathbb{R}^d$ と $b_t \in \mathbb{R}^d$ と $\bm{x}_t \in \mathbb{R}^d$（Token Shift用）なので、
$$ \text{メモリ} = 3 \times L \times d \times 2 \text{bytes} \approx 3 \times 32 \times 4{,}096 \times 2 \approx 0.75 \text{ MB} $$
実に約67,000倍のメモリ差です。これがRWKVの推論効率の根本的な優位性であり、エッジデバイスでの長文生成や、メモリが限られた環境での大規模モデル運用を可能にします。
両者の推論メモリを対数軸の棒グラフで並べました。Transformerは数十GB、RWKVは1MB未満と、桁が4つ以上違うことが一目で分かります（縦軸が対数なので見た目の差以上に開いています）。Transformer側は系列長$n$に比例して伸び続けますが、RWKV側は系列長によらず一定です。10万トークンの対話を1台のGPUで回せるか、それとも複数台を要するか — この差が実運用での決定的な分かれ目になります。
この圧倒的なメモリ効率は実用上非常に大きなインパクトを持ちますが、当然ながら精度との兼ね合いが気になります。RWKVの各バージョンがどのように性能を向上させてきたのか、進化の歴史を見ていきましょう。
RWKV v4 / v5 / v6 の進化
RWKV-4: 基盤アーキテクチャの確立
2023年に発表されたRWKV-4（論文: “RWKV: Reinventing RNNs for the Transformer Era”）は、ここまで解説してきた基本アーキテクチャを確立しました。主な特徴は次のとおりです。
- チャネルごとの固定減衰 $w_i$: 学習可能だが、訓練後は時間方向に一定
- スカラー状態: 各チャネルの状態は $(a_{t,i}, b_{t,i})$ のスカラーペア
- Token Shift: 現在と1つ前のトークンの線形補間
- モデルサイズ: 169M〜14Bパラメータの系列で公開
RWKV-4は、同規模のTransformerモデルと比較して、多くの自然言語処理ベンチマークで同等の性能を示しました。特に長系列タスクではメモリ効率の面で圧倒的な優位性を持ちます。
しかし、いくつかの限界も明らかになりました。固定の減衰率 $w_i$ では、入力に応じて「今は過去の情報をたくさん保持すべき」「ここでは過去を忘れるべき」という動的な切り替えができません。また、スカラー状態は情報の表現容量に限界があります。
RWKV-5 (Eagle): 行列値状態への拡張
RWKV-5では、状態空間モデル（SSM）の研究から着想を得て、以下の拡張が行われました。
マルチヘッドのベクトル値状態: RWKV-4ではチャネルごとにスカラーの状態 $(a_{t,i}, b_{t,i})$ を持っていましたが、RWKV-5では行列値の状態 $\bm{S}_t \in \mathbb{R}^{h \times (d/h) \times (d/h)}$ を導入しました。ここで $h$ はヘッド数です。TransformerのMulti-Head Attentionと同様に、特徴空間を複数のヘッドに分割し、各ヘッドが独立に状態を更新します。
状態更新式は次のようになります（ヘッド $m$ について）。
$$ \bm{S}_t^{(m)} = \text{diag}(\bm{\alpha}_t^{(m)}) \bm{S}_{t-1}^{(m)} + (\bm{k}_t^{(m)})^\top \bm{v}_t^{(m)} $$
ここで $\bm{\alpha}_t^{(m)} \in \mathbb{R}^{d/h}$ は減衰ベクトルで、Key次元ごとに異なる減衰率を持ちます。また、$(\bm{k}_t^{(m)})^\top \bm{v}_t^{(m)}$ は外積であり、Key-Valueペアの情報を行列として状態に書き込みます。
この行列値状態は、TransformerのKVキャッシュにおける「Key-Valueペアの蓄積」を、固定サイズの行列で近似していると解釈できます。行列のランクが情報容量の上限となりますが、系列長に依存しないメモリで情報を圧縮・保持できるという利点があります。
Group Normalization の導入: 各ヘッドの出力にGroup Normalization（LayerNormの一種）を適用し、訓練の安定性を向上させました。
RWKV-6 (Finch): 入力依存の動的パラメータ
RWKV-6は2024年に発表された最新版で、MambaのSelective Mechanism（入力依存のパラメータ）に近いアイデアを取り入れました。
時間依存の減衰（Data-Dependent Decay）: RWKV-4/5では減衰率 $w_i$ が入力に依存しない固定値でしたが、RWKV-6では入力に応じて動的に変化します。
$$ \bm{w}_t = \bm{w}_0 + \text{tanh}(\bm{W}_w \bar{\bm{x}}_t^{(w)}) \odot \bm{c}_w $$
ここで $\bm{w}_0$ はベースの減衰率、$\bm{W}_w$ は射影行列、$\bm{c}_w$ はスケーリングベクトルです。入力 $\bar{\bm{x}}_t^{(w)}$ に依存するため、「この入力は重要だから過去の記憶を強く保持しよう」「この入力はノイズだから過去を忘れよう」という動的な判断が可能になります。
改良されたToken Shift（LoRA-enhanced）: Token Shiftの混合係数もデータ依存にし、さらにLoRA（Low-Rank Adaptation）的な低ランク構造で効率的にパラメータ化しました。
$$ \bm{\mu}_t = \bm{\mu}_0 + \text{tanh}(\bm{x}_t \bm{A}) \bm{B} $$
ここで $\bm{A} \in \mathbb{R}^{d \times r}$、$\bm{B} \in \mathbb{R}^{r \times d}$（$r \ll d$）は低ランク行列です。
バージョン間の比較
| 特徴 | RWKV-4 | RWKV-5 (Eagle) | RWKV-6 (Finch) |
|---|---|---|---|
| 状態の型 | スカラー | 行列 | 行列 |
| 減衰率 | 固定（学習済み） | 固定（学習済み） | 入力依存（動的） |
| Token Shift | 固定混合係数 | 固定混合係数 | 入力依存（LoRA） |
| ヘッド構造 | なし | マルチヘッド | マルチヘッド |
| 正規化 | LayerNorm | GroupNorm | GroupNorm |
3つのバージョンの進化を1枚にまとめました。RWKV-4はスカラー状態と固定減衰で基盤を確立し、RWKV-5（Eagle）は状態を行列値・マルチヘッド化して表現容量を拡張、RWKV-6（Finch）は減衰やToken Shiftを入力依存の動的パラメータに変えました。矢印が示すとおり、進化の方向は一貫して「固定パラメータ → 入力に応じて動く動的パラメータ」です。これはMambaがS4からS6でSelective Mechanismを導入した流れと完全に並行しており、両陣営が独立に同じ結論にたどり着いたことを物語っています。
バージョンが進むにつれて、「固定パラメータ」が「入力依存の動的パラメータ」に置き換えられてきたことがわかります。この進化の方向性は、Mamba（S4 → S6でSelective Mechanismを導入）と平行しており、「入力に応じた動的なフィルタリングが長距離依存性の学習に不可欠」という共通の知見を反映しています。
理論的な構造が理解できたところで、実際にPythonでWKV機構を実装して、動作を確認してみましょう。
Pythonでの簡易RWKV実装
WKV機構の実装
まず、RWKV-4のWKV機構を NumPy で実装します。再帰形式（RNNモード）と直接計算形式の両方を実装し、数値的に一致することを確認します。
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(42)
def wkv_naive(w, u, k, v):
"""
WKV機構のナイーブ実装（直接計算、O(n^2)）
w: (d,) 減衰パラメータ（正の値）
u: (d,) 現在トークンのボーナスパラメータ
k: (n, d) Key
v: (n, d) Value
戻り値: (n, d) WKV出力
"""
n, d = k.shape
output = np.zeros((n, d))
for t in range(n):
for i in range(d):
# 分子: 現在トークンのボーナス付き寄与 + 過去トークンの減衰付き寄与
numerator = np.exp(u[i] + k[t, i]) * v[t, i]
denominator = np.exp(u[i] + k[t, i])
for j in range(t):
weight = np.exp(-(t - 1 - j) * w[i] + k[j, i])
numerator += weight * v[j, i]
denominator += weight
output[t, i] = numerator / denominator
return output
def wkv_recurrent(w, u, k, v):
"""
WKV機構の再帰実装（RNNモード、O(n)）
数値的安定性のためlog-sum-expトリックを使用
"""
n, d = k.shape
output = np.zeros((n, d))
# 累積変数の初期化
a = np.zeros(d) # 重み付きValueの累積
b = np.zeros(d) # 重みの累積
# log領域での最大値追跡
max_prev = np.full(d, -np.inf)
for t in range(n):
# 現在トークンのボーナス付き寄与
current_exp = u + k[t]
if t == 0:
# 最初のトークンは過去がないので現在トークンのみ
output[t] = v[t]
else:
# 数値安定性のために最大値を取る
max_val = np.maximum(current_exp, max_prev)
e_curr = np.exp(current_exp - max_val)
e_prev = np.exp(max_prev - max_val)
output[t] = (e_curr * v[t] + e_prev * a) / (e_curr + e_prev * b)
# 状態更新（log-sum-expトリック）
if t == 0:
a = np.exp(k[t]) * v[t]
b = np.exp(k[t])
max_prev = k[t].copy()
else:
new_max = np.maximum(max_prev - w, k[t])
a = np.exp(max_prev - w - new_max) * a + np.exp(k[t] - new_max) * v[t]
b = np.exp(max_prev - w - new_max) * b + np.exp(k[t] - new_max)
max_prev = new_max.copy()
return output
# テストデータ
n, d = 20, 8 # 系列長20、特徴次元8
w = np.abs(np.random.randn(d)) * 0.5 # 正の減衰パラメータ
u = np.random.randn(d) * 0.1 # ボーナスパラメータ
k = np.random.randn(n, d) * 0.5 # Key
v = np.random.randn(n, d) # Value
# 両方の実装で計算
out_naive = wkv_naive(w, u, k, v)
out_recurrent = wkv_recurrent(w, u, k, v)
# 数値的一致の確認
max_diff = np.max(np.abs(out_naive - out_recurrent))
print(f"ナイーブ実装と再帰実装の最大差: {max_diff:.2e}")
print(f"一致判定 (< 1e-10): {max_diff < 1e-10}")
このコードでは、WKV機構の2つの等価な実装を検証しています。wkv_naive は定義式どおりに全トークンペアを走査する $O(n^2)$ の実装であり、wkv_recurrent は累積変数を使った $O(n)$ の再帰実装です。数値的に一致することで、再帰形式への変換が正しいことを確認できます。再帰実装の最大差は浮動小数点の丸め誤差程度（$10^{-12}$ オーダー以下）に収まるはずです。
指数減衰の可視化
次に、WKV機構における重みの減衰パターンを可視化して、Transformerの注意重みとの違いを直感的に理解しましょう。
import numpy as np
import matplotlib.pyplot as plt
# 異なる減衰率での重みパターン
n = 50
decay_rates = [0.1, 0.3, 0.7, 1.5]
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# 左: 各減衰率での重みの時間変化
ax = axes[0]
for w_val in decay_rates:
t_query = n - 1 # 最後のトークンから見た重み
positions = np.arange(n)
weights = np.exp(-(t_query - positions) * w_val)
weights[t_query] = 0 # 現在トークンは別扱い（ボーナス項）
weights = weights / weights.sum() # 正規化
ax.plot(positions, weights, label=f'w = {w_val}', linewidth=2)
ax.set_xlabel('Token position', fontsize=12)
ax.set_ylabel('Attention weight (normalized)', fontsize=12)
ax.set_title('WKV Exponential Decay Weights\n(from last token perspective)', fontsize=13)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
# 右: RWKVの減衰行列 vs Transformerの注意行列（ヒートマップ）
ax = axes[1]
n_small = 16
w_val = 0.3
decay_matrix = np.zeros((n_small, n_small))
for t in range(n_small):
for j in range(t + 1):
if j < t:
decay_matrix[t, j] = np.exp(-(t - 1 - j) * w_val)
else:
decay_matrix[t, j] = 1.5 # 現在トークンのボーナス
# 各行を正規化
for t in range(n_small):
row_sum = decay_matrix[t, :t+1].sum()
if row_sum > 0:
decay_matrix[t, :t+1] /= row_sum
im = ax.imshow(decay_matrix, cmap='viridis', aspect='auto', origin='upper')
ax.set_xlabel('Key position (j)', fontsize=12)
ax.set_ylabel('Query position (t)', fontsize=12)
ax.set_title(f'WKV Weight Matrix (w={w_val})\n(Causal, Exponential Decay)', fontsize=13)
plt.colorbar(im, ax=ax, label='Weight')
plt.tight_layout()
plt.savefig('rwkv_decay_pattern.png', dpi=150, bbox_inches='tight')
plt.show()
左のグラフから、減衰率 $w$ が大きいほど「直近のトークンだけに集中する鋭い重み分布」になり、$w$ が小さいほど「過去のトークンにも広く注意を向ける緩やかな分布」になることが読み取れます。これはTransformerの注意重みがデータ依存で任意のパターンを取れるのと対照的です。RWKVでは減衰率 $w$ がチャネルごとに学習されるため、あるチャネルは直近の文脈に特化し、別のチャネルは長距離の文脈を保持するという役割分担が自然に生まれます。
右のヒートマップは、WKVの重み行列が因果的（下三角）かつ指数減衰のパターンを持つことを示しています。対角要素（現在のトークン自身）が最も明るく、左に行くほど暗くなる — これがRWKVの「最近の情報を優先し、遠い過去を忘れていく」メカニズムの可視化です。
Time Mixing ブロック全体の実装
WKV機構を含むTime Mixingブロック全体をクラスとして実装します。
import numpy as np
class RWKVTimeMixing:
"""RWKV-4 Time Mixing ブロックの簡易実装"""
def __init__(self, d_model, seed=42):
rng = np.random.RandomState(seed)
self.d = d_model
# 学習可能パラメータ（本来は訓練で最適化される）
scale = 0.1
self.W_r = rng.randn(d_model, d_model) * scale
self.W_k = rng.randn(d_model, d_model) * scale
self.W_v = rng.randn(d_model, d_model) * scale
self.W_o = rng.randn(d_model, d_model) * scale
# Token Shift の混合係数（0〜1）
self.mu_r = np.clip(rng.rand(d_model), 0.3, 0.7)
self.mu_k = np.clip(rng.rand(d_model), 0.3, 0.7)
self.mu_v = np.clip(rng.rand(d_model), 0.3, 0.7)
# WKV パラメータ
self.w = np.abs(rng.randn(d_model)) * 0.5 # 減衰率（正）
self.u = rng.randn(d_model) * 0.1 # 現在トークンボーナス
def sigmoid(self, x):
return 1.0 / (1.0 + np.exp(-np.clip(x, -20, 20)))
def token_shift(self, x_curr, x_prev, mu):
"""Token Shift: 現在と前のトークンの線形補間"""
return mu * x_curr + (1 - mu) * x_prev
def forward_recurrent(self, x_seq):
"""
再帰モード（RNNモード）での順伝播
x_seq: (n, d) 入力系列
戻り値: (n, d) 出力系列
"""
n, d = x_seq.shape
outputs = np.zeros((n, d))
# 状態の初期化
a = np.zeros(d)
b = np.zeros(d)
max_prev = np.full(d, -np.inf)
x_prev = np.zeros(d) # Token Shift用
for t in range(n):
# Token Shift
x_r = self.token_shift(x_seq[t], x_prev, self.mu_r)
x_k = self.token_shift(x_seq[t], x_prev, self.mu_k)
x_v = self.token_shift(x_seq[t], x_prev, self.mu_v)
# R, K, V の計算
r = x_r @ self.W_r.T
k = x_k @ self.W_k.T
v = x_v @ self.W_v.T
# WKV 計算
if t == 0:
wkv = v
else:
current_exp = self.u + k
max_val = np.maximum(current_exp, max_prev)
e_curr = np.exp(current_exp - max_val)
e_prev = np.exp(max_prev - max_val)
wkv = (e_curr * v + e_prev * a) / (e_curr + e_prev * b + 1e-12)
# Receptance ゲーティング
o = self.sigmoid(r) * wkv
# 出力射影
outputs[t] = o @ self.W_o.T
# 状態更新
if t == 0:
a = np.exp(k) * v
b = np.exp(k)
max_prev = k.copy()
else:
new_max = np.maximum(max_prev - self.w, k)
a = np.exp(max_prev - self.w - new_max) * a + np.exp(k - new_max) * v
b = np.exp(max_prev - self.w - new_max) * b + np.exp(k - new_max)
max_prev = new_max.copy()
x_prev = x_seq[t].copy()
return outputs
# 動作確認
d_model = 64
n_tokens = 32
tm = RWKVTimeMixing(d_model)
# ランダム入力で動作テスト
x_input = np.random.randn(n_tokens, d_model) * 0.5
output = tm.forward_recurrent(x_input)
print(f"入力形状: {x_input.shape}")
print(f"出力形状: {output.shape}")
print(f"出力の統計: mean={output.mean():.4f}, std={output.std():.4f}")
print(f"出力にNaNなし: {not np.any(np.isnan(output))}")
このクラスはRWKV-4のTime Mixingブロックを忠実に再現しています。forward_recurrent メソッドがRNNモードの推論を実装しており、各時刻で状態 $(a, b, \text{max\_prev})$ を更新しながら出力を計算します。出力の平均が0に近く、標準偏差が有限の正の値であり、NaNが含まれていないことが確認できれば、実装が正常に動作していると判断できます。
計算量のスケーリング比較
最後に、RWKVの再帰実装、ナイーブ実装（$O(n^2)$）、そしてTransformerのAttention（$O(n^2)$）の計算時間を系列長に対して比較します。
import numpy as np
import matplotlib.pyplot as plt
import time
def benchmark_wkv_recurrent(n, d):
"""RWKVの再帰実装の計算時間を測定"""
w = np.abs(np.random.randn(d)) * 0.5
u = np.random.randn(d) * 0.1
k = np.random.randn(n, d) * 0.5
v = np.random.randn(n, d)
start = time.perf_counter()
_ = wkv_recurrent(w, u, k, v)
return time.perf_counter() - start
def benchmark_attention(n, d):
"""Transformerの Scaled Dot-Product Attention の計算時間を測定"""
Q = np.random.randn(n, d) * 0.5
K = np.random.randn(n, d) * 0.5
V = np.random.randn(n, d)
start = time.perf_counter()
scores = Q @ K.T / np.sqrt(d)
# 因果マスク
mask = np.triu(np.ones((n, n)) * (-1e9), k=1)
scores = scores + mask
# softmax
scores_exp = np.exp(scores - scores.max(axis=-1, keepdims=True))
attn = scores_exp / scores_exp.sum(axis=-1, keepdims=True)
_ = attn @ V
return time.perf_counter() - start
# ベンチマーク
seq_lengths = [50, 100, 200, 500, 1000, 2000, 3000]
d = 32 # 特徴次元（小さめにしてNumPyで実行可能に）
times_rwkv = []
times_attn = []
for n in seq_lengths:
# 各設定で3回測定して中央値
t_rwkv = np.median([benchmark_wkv_recurrent(n, d) for _ in range(3)])
t_attn = np.median([benchmark_attention(n, d) for _ in range(3)])
times_rwkv.append(t_rwkv)
times_attn.append(t_attn)
print(f"n={n:5d}: RWKV={t_rwkv:.4f}s, Attention={t_attn:.4f}s")
# プロット
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# 線形スケール
ax = axes[0]
ax.plot(seq_lengths, times_rwkv, 'o-', label='RWKV (recurrent, O(n))',
linewidth=2, markersize=6, color='#00bcd4')
ax.plot(seq_lengths, times_attn, 's-', label='Attention (O(n²))',
linewidth=2, markersize=6, color='#ff5722')
ax.set_xlabel('Sequence length (n)', fontsize=12)
ax.set_ylabel('Time (seconds)', fontsize=12)
ax.set_title('Computation Time: RWKV vs Attention', fontsize=13)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
# 対数スケール
ax = axes[1]
ax.loglog(seq_lengths, times_rwkv, 'o-', label='RWKV (recurrent)',
linewidth=2, markersize=6, color='#00bcd4')
ax.loglog(seq_lengths, times_attn, 's-', label='Attention',
linewidth=2, markersize=6, color='#ff5722')
# 理論的なスケーリングの参照線
n_ref = np.array(seq_lengths, dtype=float)
ax.loglog(n_ref, n_ref / n_ref[0] * times_rwkv[0], '--',
color='#00bcd4', alpha=0.5, label='O(n) reference')
ax.loglog(n_ref, (n_ref / n_ref[0])**2 * times_attn[0], '--',
color='#ff5722', alpha=0.5, label='O(n²) reference')
ax.set_xlabel('Sequence length (n)', fontsize=12)
ax.set_ylabel('Time (seconds)', fontsize=12)
ax.set_title('Computation Time (log-log scale)', fontsize=13)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('rwkv_vs_attention_benchmark.png', dpi=150, bbox_inches='tight')
plt.show()
左の線形スケールのグラフから、Attentionの計算時間が系列長の増加とともに急激に増大するのに対し、RWKVの再帰実装はほぼ線形に増加していることが明確に読み取れます。系列長3000の時点でAttentionはRWKVの数十倍以上の時間を要しています。
右の対数-対数スケールでは、各手法のスケーリング特性がより明確に見えます。RWKVのプロットは傾き1の参照線（$O(n)$）に沿い、Attentionのプロットは傾き2の参照線（$O(n^2)$）に沿っています。これは理論的な計算量の予測と一致しており、実装の正しさを裏付けています。ただし、このベンチマークはNumPyによるCPU実装であるため、実際のGPU実装ではFlash Attentionなどの最適化によって定数倍の差は縮まる可能性があります。それでも漸近的なスケーリングの差（$O(n)$ vs $O(n^2)$）は変わりません。
Transformer / Mamba / RWKV の比較
アーキテクチャの構造比較
ここまでRWKVの理論と実装を詳しく見てきましたが、同じ「Transformer代替」を目指すMambaとの違いを理解しておくことは、次世代モデルの設計を俯瞰する上で重要です。
| 特徴 | Transformer | Mamba (S6) | RWKV-6 |
|---|---|---|---|
| コア機構 | Softmax Attention | Selective SSM | WKV + Receptance |
| 理論的背景 | 全ペア類似度 | 状態空間方程式 | 指数減衰付き重み平均 |
| 訓練計算量 | $O(n^2 d)$ | $O(n d)$（スキャン） | $O(n d)$（行列演算） |
| 推論メモリ | $O(n d)$（KVキャッシュ） | $O(d N)$（状態） | $O(d^2 / h)$（行列状態、v5以降） |
| 推論1ステップ | $O(n d)$ | $O(d N)$ | $O(d^2)$ |
| 入力依存パラメータ | 全て（Q,K,V） | $\bm{B}, \bm{C}, \Delta$ | $w_t, \mu_t$（v6） |
| 並列訓練 | 行列演算 | 並列スキャン | 行列演算 / CUDAカーネル |
| 位置エンコーディング | 必要（RoPE等） | 不要（連続SSM） | 不要（Token Shift） |
設計思想の違い
Transformerは「全トークンペアの関係を明示的に計算する」アプローチです。表現力が最も高く、任意の長距離依存パターンを学習できます。しかし、その表現力の代償として $O(n^2)$ の計算量がかかります。
Mambaは制御理論の状態空間方程式を出発点としています。連続時間の微分方程式を離散化し、入力依存のパラメータ（Selective Mechanism）で情報のフィルタリングを行います。Mambaの核心は選択的な忘却です — 入力に応じて「何を状態に書き込み、何を忘れるか」を動的に制御します。
RWKVは、AttentionのQuery-Key-Valueフレームワークを出発点としつつ、softmaxを指数減衰で置き換えるというアプローチを取ります。核心は時間減衰付きの重み付け平均 + Receptanceゲーティングです。Attentionの「どこに注目するか」を、「過去の情報をどの速度で忘れるか」と「現在の出力にどれだけ反映するか」の2段階に分解したと解釈できます。
性能比較
公開されているベンチマーク結果に基づくと、おおまかな傾向は次のとおりです。
言語モデリング（Perplexity）: 同じパラメータ数で比較した場合、Transformer（GPT系）が最も低いPerplexityを示し、Mamba → RWKV-6の順で僅差で続きます。ただし差は小さく、モデルサイズが大きくなるほど差は縮まります。
長系列タスク: 系列長が数万トークンを超えるタスク（Long Range Arena等）では、Mamba と RWKV がTransformerを上回る場合があります。これはメモリの制約でTransformerがそもそも長い系列を処理できないケースがあるためです。
推論スループット: 長系列の自己回帰生成では、RWKVとMambaがTransformerを大幅に上回ります。特にバッチサイズ1での推論（チャットボット等）では、KVキャッシュの読み書きがボトルネックになるTransformerに対し、固定サイズの状態で済むRWKV・Mambaの優位性が顕著です。
それぞれの適性
| ユースケース | 最適なアーキテクチャ | 理由 |
|---|---|---|
| 汎用LLM（精度最優先） | Transformer | 表現力が最も高い |
| エッジデバイス推論 | RWKV / Mamba | メモリ消費が極めて少ない |
| 超長系列処理 | RWKV / Mamba | 線形計算量で系列長を拡張可能 |
| リアルタイムストリーミング | RWKV | RNNモードで $O(1)$ ステップ推論 |
| ゲノム・時系列 | Mamba | SSMの構造が連続時間データと親和性が高い |
現時点では「万能のアーキテクチャ」は存在せず、タスクの要件に応じた使い分けが重要です。ただし、RWKVとMambaの急速な進化を考えると、Transformerの $O(n^2)$ Attentionが今後も支配的であり続けるかは不透明です。
まとめ
本記事では、RWKVアーキテクチャの理論と実装を解説しました。
- RWKVはRNNとTransformerの長所を融合したモデルです。訓練時はTransformerのように並列計算でき、推論時はRNNのように $O(1)$ メモリで逐次処理できます
- WKV機構は、Self-Attentionの softmax を指数減衰付きの重み付け平均で置き換えたものです。再帰形式と並列形式の双対性により、同じパラメータで計算モードを切り替えられます
- Token Shiftは、現在と1つ前のトークンの線形補間で最小限の文脈情報を注入する軽量な操作です。R, K, V のそれぞれに異なる混合係数を使うことで柔軟な制御が可能です
- Channel Mixingは、TransformerのFFNに相当し、squared ReLU とReceptanceゲーティングで非線形変換を行います
- Receptance（受容度）は、RWKVの特徴的な概念であり、sigmoid ゲートで過去の情報をどれだけ受け入れるかを制御します
- RWKV v4 → v5 → v6 の進化は、「固定パラメータから入力依存の動的パラメータへ」という方向であり、MambaのSelective Mechanismと共通する知見です
- 推論メモリの差は劇的です。同じモデルサイズで10万トークンを処理する場合、TransformerのKVキャッシュが約50GBを要するのに対し、RWKVの状態は約0.75MBで済みます
RWKVの理解は、線形計算量モデルの全体像を把握するための重要なステップです。次のステップとして、以下の記事も参考にしてください。