このモデルは2023年5月9日に Hugging Face Transformers に貢献されました。

RWKV

概要

RWKVモデル（バージョン4）は、このリポジトリで提案されました。

従来の Transformer の注意機構を調整して線形化することを提案しています。これにより、モデルを再帰型ネットワークとして使用できます。タイムスタンプ 0 とタイムスタンプ 1 の入力を同時に渡すことは、まずタイムスタンプ 0 の入力を渡し、次にタイムスタンプ 0 の状態とともにタイムスタンプ 1 の入力を渡すことと同等です（以下の例を参照）。

これは通常の Transformer よりも効率的であり、任意の長さの文を処理できます（モデルが学習時に固定のコンテキスト長を使用する場合でも）。

このモデルは sgugger によって提供されました。 元のコードは こちら で確認できます。

使用例

import torch

from transformers import AutoTokenizer, RwkvModel


model = RwkvModel.from_pretrained("sgugger/rwkv-430M-pile", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("sgugger/rwkv-430M-pile")

inputs = tokenizer("This is an example.", return_tensors="pt").to(model.device)
# Feed everything to the model
outputs = model(inputs["input_ids"])
output_whole = outputs.last_hidden_state

outputs = model(inputs["input_ids"][:, :2])
output_one = outputs.last_hidden_state

# Using the state computed on the first inputs, we will get the same output
outputs = model(inputs["input_ids"][:, 2:], state=outputs.state)
output_two = outputs.last_hidden_state

torch.allclose(torch.cat([output_one, output_two], dim=1), output_whole, atol=1e-5)

'\n\n' が検出された時点でモデルの生成を停止させたい場合は、以下の停止条件の使用を推奨します。

from transformers import StoppingCriteria


class RwkvStoppingCriteria(StoppingCriteria):
    def __init__(self, eos_sequence = [187,187], eos_token_id = 537):
        self.eos_sequence = eos_sequence
        self.eos_token_id = eos_token_id

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        last_2_ids = input_ids[:,-2:].tolist()
        return self.eos_sequence in last_2_ids


output = model.generate(inputs["input_ids"], max_new_tokens=64, stopping_criteria = [RwkvStoppingCriteria()])

注記

• バッチ生成には左パディングを使用します。Padding sideを参照してください。

RwkvConfig

>>> from transformers import RwkvConfig, RwkvModel

>>> # Initializing a Rwkv configuration
>>> configuration = RwkvConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = RwkvModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

RwkvModel

RwkvLMHeadModel

Rwkvのattentionと再帰式

従来の自己回帰型 Transformer では、アテンションは以下のように表されます。 $$O = \hbox{softmax}(QK^{T} / \sqrt{d}) V$$

と $Q$ は $K$ 、 $V$ は形状が seq_len x hidden_size の行列で、それぞれクエリ、キー、バリューと呼ばれます（実際にはバッチ次元とアテンションヘッド次元を持つより大きな行列ですが、行列積が計算される最後の2次元にのみ注目するため、簡略化のためにこれら2次元のみを考慮します）。積 $QK^{T}$ の形状は seq_len x seq_len となり、 $V$ との行列積を計算することで、出力 $O$ が得られ、他の行列と同じ形状になります。

softmax をその値で置き換えると、次のようになります: $$O_{i} = \frac{\sum_{j=1}^{i} e^{Q_{i} K_{j}^{T} / \sqrt{d}} V_{j}}{\sum_{j=1}^{i} e^{Q_{i} K_{j}^{T} / \sqrt{d}}}$$

の項目において、 $QK^{T}$ に対応する $j > i$ の要素はマスクされています（合計は j で打ち切られます）。これは、アテンションが将来のトークンではなく、過去のトークンのみに参照を許可されるためです。

これに対し、RWKV のアテンションは次のように与えられます。 $$O_{i} = \sigma(R_{i}) \frac{\sum_{j=1}^{i} e^{W_{i-j} + K_{j}} V_{j}}{\sum_{j=1}^{i} e^{W_{i-j} + K_{j}}}$$

ここで $R$ は著者が受容率 (receptance) と呼ぶ新しい行列、 $K$ および $V$ は依然としてキーとバリューです（ここでは $\sigma$ はシグモイド関数です）。 $W$ はトークンの位置を表す新しいベクトルであり、以下のように与えられます $$W_{0} = u \hbox{ and } W_{k} = (k-1)w \hbox{ for } k \geq 1$$

の $u$ および $w$ は、コード内でそれぞれ time_first および time_decay と呼ばれる学習可能なパラメータです。分子と分母は、どちらも再帰的に表すことができます。これらを $N_{i}$ および $D_{i}$ と名付けると、次のようになります。 $$N_{i} = e^{u + K_{i}} V_{i} + \hat{N}_{i} \hbox{ where } \hat{N}_{i} = e^{K_{i-1}} V_{i-1} + e^{w + K_{i-2}} V_{i-2} \cdots + e^{(i-2)w + K_{1}} V_{1}$$

したがって $\hat{N}_{i}$ (（コードでは numerator_state と呼ばれます）は、 $$\hat{N}_{0} = 0 \hbox{ and } \hat{N}_{j+1} = e^{K_{j}} V_{j} + e^{w} \hat{N}_{j}$$

および $$D_{i} = e^{u + K_{i}} + \hat{D}_{i} \hbox{ where } \hat{D}_{i} = e^{K_{i-1}} + e^{w + K_{i-2}} \cdots + e^{(i-2)w + K_{1}}$$

したがって $\hat{D}_{i}$ (（コードでは denominator_state と呼ばれます）は $$\hat{D}_{0} = 0 \hbox{ and } \hat{D}_{j+1} = e^{K_{j}} + e^{w} \hat{D}_{j}$$

実際に使用される再帰式は、数値の安定性を保つために大きな数の指数を計算したくないため、わずかに複雑になっています。通常、softmax はそのまま計算されるのではなく、最大項の指数を分子と分母で割ります: $$\frac{e^{x_{i}}}{\sum_{j=1}^{n} e^{x_{j}}} = \frac{e^{x_{i} - M}}{\sum_{j=1}^{n} e^{x_{j} - M}}$$

最大値 $M$ すべての $x_{j}$ の最大値です。つまり、ここでは分子の状態 ($\hat{N}$) と分母の状態 ($\hat{D}$) を保存するだけでなく、指数関数で遭遇したすべての項の最大値も追跡します。したがって、実際には $$\tilde{N}_{i} = e^{-M_{i}} \hat{N}_{i} \hbox{ and } \tilde{D}_{i} = e^{-M_{i}} \hat{D}_{i}$$

以下の再帰式で定義されます: $$\tilde{N}_{0} = 0 \hbox{ and } \tilde{N}_{j+1} = e^{K_{j} - q} V_{j} + e^{w + M_{j} - q} \tilde{N}_{j} \hbox{ where } q = \max(K_{j}, w + M_{j})$$

および $$\tilde{D}_{0} = 0 \hbox{ and } \tilde{D}_{j+1} = e^{K_{j} - q} + e^{w + M_{j} - q} \tilde{D}_{j} \hbox{ where } q = \max(K_{j}, w + M_{j})$$

および $M_{j+1} = q$。これらを用いて、次を計算できます $$N_{i} = e^{u + K_{i} - q} V_{i} + e^{M_{i}} \tilde{N}_{i} \hbox{ where } q = \max(u + K_{i}, M_{i})$$

および $$D_{i} = e^{u + K_{i} - q} + e^{M_{i}} \tilde{D}_{i} \hbox{ where } q = \max(u + K_{i}, M_{i})$$

これにより最終的に得られるのは $$O_{i} = \sigma(R_{i}) \frac{N_{i}}{D_{i}}$$