이 모델은 2023-05-09에 Hugging Face Transformers에 기여되었습니다.

RWKV

개요

RWKV 모델(4버전)은 이 저장소에서 제안되었습니다.

전통적인 Transformer attention을 선형으로 조정하는 방식을 제안합니다. 이렇게 하면 모델을 recurrent network로 사용할 수 있습니다. timestamp 0과 timestamp 1의 입력을 함께 전달하는 것은, timestamp 0의 입력을 전달한 후 timestamp 0의 상태와 함께 timestamp 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.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()])

The input is a single word "Notes". 참고

• 배치 생성에는 왼쪽 패딩을 사용합니다. 패딩 방향을 참조하세요.

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에서는 attention을 다음과 같이 표현합니다. $$O = \hbox{softmax}(QK^{T} / \sqrt{d}) V$$

여기서 $Q$는 $K$ 이고 $V$ 는 각각 seq_len x hidden_size 크기의 행렬로, query, key, value라고 불립니다(실제로는 batch 차원과 attention head 차원이 포함된 더 큰 행렬이지만, 행렬 곱이 수행되는 마지막 두 차원만 관심 대상이므로 단순화를 위해 이 두 차원만 고려합니다). 곱 $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에서 중단됨). 이는 attention이 미래 토큰을 참조할 수 없고(과거 토큰만 참조 가능) 과거 토큰만 참조할 수 있기 때문입니다.

비교하자면, RWKV의 attention은 다음과 같이 주어집니다. $$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$ 은 여전히 key와 value입니다(여기서 $\sigma$는 sigmoid 함수입니다). $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}}$$

with $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}}$$