该模型于 2023-05-09 贡献至 Hugging Face Transformers。

RWKV

概述

RWKV 模型（版本 4）在此仓库中提出

它建议对传统 Transformer 注意力机制进行微调，使其变为线性。这样一来，该模型就可以作为循环网络使用：同时传入时间戳 0 和时间戳 1 的输入，与先传入时间戳 0 的输入、再传入时间戳 1 的输入以及时间戳 0 的状态是等价的（参见下方示例）。

这比常规 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.Tensor, **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()])

备注

• 在批量生成时使用左侧填充。请参阅 填充侧。

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 注意力机制与循环公式

在传统的自回归 Transformer 中，注意力机制表示为 $$O = \hbox{softmax}(QK^{T} / \sqrt{d}) V$$

其中 $Q$、 $K$ 和 $V$ 是形状为 seq_len x hidden_size 的矩阵，分别称为查询（query）、键（key）和值（value）（它们实际上是包含批次维度和注意力头维度的更大矩阵，但我们只关注最后两个维度，因为矩阵乘法是在这两个维度上进行的，所以为了简洁起见，我们只考虑这两个维度）。乘积 $QK^{T}$ 的形状为 seq_len x seq_len ，我们可以将其与 $V$ 进行矩阵乘法，得到形状与其他矩阵相同的输出 $O$ 。

用 softmax 的值替换 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 为止），因为注意力机制不允许查看未来的 token（只能查看过去的 token）。

相比之下，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$ 为 sigmoid 函数）。 $W$ 是一个表示 token 位置的新向量，其定义为 $$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}}$$