Mô hình này đã được đóng góp cho Hugging Face Transformers vào ngày 2023-05-09.

RWKV

Tổng quan

Mô hình RWKV (phiên bản 4) được đề xuất trong repo này

Đề xuất một điều chỉnh trong cơ chế chú ý Transformer truyền thống để biến nó thành tuyến tính. Nhờ đó, mô hình có thể được sử dụng như một mạng nơ-ron tuần tự: việc truyền các đầu vào cho mốc thời gian 0 và mốc thời gian 1 cùng lúc là tương đương với việc truyền đầu vào tại mốc thời gian 0, sau đó truyền đầu vào tại mốc thời gian 1 cùng với trạng thái tại mốc thời gian 0 (xem ví dụ bên dưới).

Điều này có thể hiệu quả hơn so với Transformer thông thường và có thể xử lý câu có độ dài bất kỳ (ngay cả khi mô hình sử dụng độ dài ngữ cảnh cố định cho huấn luyện).

Mô hình này được đóng góp bởi sgugger. Mã nguồn gốc có thể được tìm thấy tại đây.

Ví dụ sử dụng

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ếu bạn muốn đảm bảo mô hình dừng tạo khi phát hiện '\n\n', chúng tôi khuyên bạn nên sử dụng các tiêu chí dừng sau:

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()])

Ghi chú

• Sử dụng left padding cho việc tạo batched. Xem 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

Cơ chế chú ý Rwkv và các công thức đệ quy

Trong một Transformer tự hồi quy truyền thống, attention được viết dưới dạng $$O = \hbox{softmax}(QK^{T} / \sqrt{d}) V$$

với $Q$, $K$ và $V$ là các ma trận có hình dạng seq_len x hidden_size được gọi là query, key và value (thực tế chúng là các ma trận lớn hơn có thêm chiều batch và chiều attention head, nhưng chúng ta chỉ quan tâm đến hai chiều cuối cùng, nơi phép nhân ma trận được thực hiện, nên vì sự đơn giản, chúng ta chỉ xét hai chiều đó). Tích $QK^{T}$ sau đó có hình dạng seq_len x seq_len và chúng ta có thể thực hiện phép nhân ma trận với $V$ để lấy đầu ra $O$ có cùng hình dạng với các ma trận khác.

Thay softmax bằng giá trị của nó, ta có: $$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}}}$$

Lưu ý rằng các phần tử trong $QK^{T}$ tương ứng với $j > i$ bị che (tổng dừng lại ở j) vì attention không được phép nhìn vào các token tương lai (chỉ nhìn vào các token quá khứ).

Ngược lại, cơ chế chú ý của RWKV được cho bởi $$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}}}$$

trong đó $R$ là một ma trận mới được tác giả gọi là receptance, $K$ và $V$ vẫn là key và value ($\sigma$ ở đây là hàm sigmoid). $W$ là một vector mới đại diện cho vị trí của token và được cho bởi $$W_{0} = u \hbox{ and } W_{k} = (k-1)w \hbox{ for } k \geq 1$$

với $u$ và $w$ các tham số có thể học được gọi là time_first và time_decay tương ứng. Tử số và mẫu số đều có thể được biểu diễn đệ quy. Đặt tên cho chúng là $N_{i}$ và $D_{i}$ ta có: $$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}$$

sao $\hat{N}_{i}$ (được gọi là numerator_state trong mã) thỏa mãn $$\hat{N}_{0} = 0 \hbox{ and } \hat{N}_{j+1} = e^{K_{j}} V_{j} + e^{w} \hat{N}_{j}$$

và $$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}}$$

nên $\hat{D}_{i}$ (được gọi là denominator_state trong mã) thỏa mãn $$\hat{D}_{0} = 0 \hbox{ and } \hat{D}_{j+1} = e^{K_{j}} + e^{w} \hat{D}_{j}$$

Công thức đệ quy thực tế được sử dụng phức tạp hơn một chút, vì để đảm bảo ổn định số, chúng ta không muốn tính lũy thừa của các số lớn. Thông thường, softmax không được tính trực tiếp, mà lũy thừa của số hạng lớn nhất được chia cho cả tử và mẫu: $$\frac{e^{x_{i}}}{\sum_{j=1}^{n} e^{x_{j}}} = \frac{e^{x_{i} - M}}{\sum_{j=1}^{n} e^{x_{j} - M}}$$

với $M$ là giá trị lớn nhất của tất cả $x_{j}$. Vì vậy, ở đây ngoài việc lưu trạng tử tử số ($\hat{N}$) và trạng tử mẫu số ($\hat{D}$), chúng ta còn theo dõi giá trị lớn nhất của tất cả các số hạng gặp trong các hàm mũ. Do đó, chúng ta thực sự sử dụng $$\tilde{N}_{i} = e^{-M_{i}} \hat{N}_{i} \hbox{ and } \tilde{D}_{i} = e^{-M_{i}} \hat{D}_{i}$$

được định nghĩa bởi các công thức đệ quy sau: $$\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})$$

và $$\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})$$

và $M_{j+1} = q$. Với những giá trị này, chúng ta có thể tính toán $$N_{i} = e^{u + K_{i} - q} V_{i} + e^{M_{i}} \tilde{N}_{i} \hbox{ where } q = \max(u + K_{i}, M_{i})$$

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

cuối cùng cho chúng ta $$O_{i} = \sigma(R_{i}) \frac{N_{i}}{D_{i}}$$