RWKVアーキテクチャ徹底解説：Transformerに匹敵する線形アテンションRNN

+--------------------------------------------------+
|                RWKV Block (x N)                  |
|                                                  |
|  +--------------------------------------------+  |
|  |           Layer Norm                       |  |
|  +---------------------+----------------------+  |
|                        |                         |
|  +---------------------v----------------------+  |
|  |       Time Mixing (WKV Attention)          |  |
|  |                                            |  |
|  |   x_t --+-- R (Receptance) -- sigmoid(r)  |  |
|  |         +-- K (Key)                        |  |
|  |         +-- V (Value)                      |  |
|  |         +-- W (Time Decay, learned)        |  |
|  |                                            |  |
|  |   wkv_t = weighted_sum(K, V, W, U)        |  |
|  |   out_t = sigmoid(r_t) * wkv_t            |  |
|  +---------------------+----------------------+  |
|                        | + residual              |
|  +---------------------v----------------------+  |
|  |           Layer Norm                       |  |
|  +---------------------+----------------------+  |
|                        |                         |
|  +---------------------v----------------------+  |
|  |       Channel Mixing (FFN)                 |  |
|  |                                            |  |
|  |   x_t --+-- R (Receptance) -- sigmoid(r)  |  |
|  |         +-- K (Key) -- squared_relu(k)     |  |
|  |                                            |  |
|  |   out_t = sigmoid(r_t) * (W_v * relu2(k)) |  |
|  +---------------------+----------------------+  |
|                        | + residual              |
+--------------------------------------------------+

wkv_t = (sum_{i=1}^{t-1} exp(-(t-1-i)*w + k_i) * v_i  +  exp(u + k_t) * v_t)
        / (sum_{i=1}^{t-1} exp(-(t-1-i)*w + k_i)  +  exp(u + k_t))

alpha_t = exp(-w) * alpha_{t-1} + exp(k_t) * v_t
beta_t  = exp(-w) * beta_{t-1}  + exp(k_t)

wkv_t = (exp(-w) * alpha_{t-1} + exp(u + k_t) * v_t)
      / (exp(-w) * beta_{t-1}  + exp(u + k_t))

import torch

def rwkv_wkv_single_step(w, u, k, v, alpha_prev, beta_prev, log_max_prev):
    """
    Single-step WKV computation with numerical stability.

    Args:
        w: time decay (d,) -- positive values
        u: bonus parameter (d,)
        k: key vector (d,)
        v: value vector (d,)
        alpha_prev: running numerator state (d,)
        beta_prev: running denominator state (d,)
        log_max_prev: log of max exponent for stability (d,)

    Returns:
        wkv: output vector (d,)
        alpha_new, beta_new, log_max_new: updated states
    """
    # Compute log-space exponents
    log_ew = -w  # log(e^(-w))
    log_ek = k   # log(e^k)

    # For the past terms: e^(-w) * prev = e^(log_ew + log_max_prev)
    log_past = log_ew + log_max_prev
    # For the current term: e^(u + k)
    log_curr = u + k

    # Numerically stable log-sum-exp
    log_max_new = torch.max(log_past, log_curr)

    past_scale = torch.exp(log_past - log_max_new)
    curr_scale = torch.exp(log_curr - log_max_new)

    wkv = (past_scale * alpha_prev + curr_scale * v) / \
          (past_scale * beta_prev + curr_scale)

    # Update running states
    log_max_alpha = torch.max(log_past, log_ek)
    alpha_new = torch.exp(log_past - log_max_alpha) * alpha_prev + \
                torch.exp(log_ek - log_max_alpha) * v
    beta_new = torch.exp(log_past - log_max_alpha) * beta_prev + \
               torch.exp(log_ek - log_max_alpha)

    return wkv, alpha_new, beta_new, log_max_alpha

def rwkv_wkv_parallel(w, u, k, v):
    """
    Parallel WKV computation for training.

    Args:
        w: time decay (d,)
        u: bonus (d,)
        k: keys (T, d)
        v: values (T, d)

    Returns:
        wkv: output (T, d)
    """
    T, d = k.shape
    ew = torch.exp(-w)  # per-channel decay factor (d,)
    ek = torch.exp(k)   # (T, d)
    ekv = ek * v        # (T, d)

    # Conceptual O(T*d) scan -- in practice uses a custom CUDA kernel
    alpha = torch.zeros(d, device=k.device)
    beta = torch.zeros(d, device=k.device)
    wkv = torch.zeros(T, d, device=k.device)

    for t in range(T):
        # Current token with bonus
        euk = torch.exp(u + k[t])

        wkv[t] = (alpha + euk * v[t]) / (beta + euk)

        # Update running sums (without bonus)
        alpha = ew * alpha + ek[t] * v[t]
        beta = ew * beta + ek[t]

    return wkv

+-----------------------------------------------------------+
|               RWKV Dual Operation Modes                   |
+---------------------------+-------------------------------+
|   Training Mode           |   Inference Mode              |
|   (Transformer-like)      |   (RNN-like)                  |
|                           |                               |
|   +---+---+---+---+      |   +---+                       |
|   |t=1|t=2|t=3|t=4|      |   |t=n|                       |
|   +-+-+-+-+-+-+-+-++      |   +-+-+                       |
|     |   |   |   |         |     |                         |
|     v   v   v   v         |     v                         |
|   +---------------+      |   +-------+   +---------+     |
|   | Parallel Scan |      |   | WKV   |-->| State   |     |
|   | (all at once) |      |   | step  |   | (a, b)  |     |
|   +-------+-------+      |   +---+---+   +---------+     |
|           |               |       |                       |
|           v               |       v                       |
|   +---------------+      |   +-----------+               |
|   |  O(Td) total  |      |   | O(d)/step |               |
|   |  parallelized |      |   | constant  |               |
|   +---------------+      |   +-----------+               |
+---------------------------+-------------------------------+
|  Same weights, same results, different compute pattern    |
+-----------------------------------------------------------+

# Version evolution summary
rwkv_versions = {
    "v4 (Pile)": {
        "year": 2023,
        "max_params": "14B",
        "key_feature": "Original WKV with fixed time decay",
        "state_size": "d per layer (scalar decay)"
    },
    "v5 (Eagle)": {
        "year": 2024,
        "max_params": "7.5B",
        "key_feature": "Multi-headed WKV, matrix-valued states",
        "state_size": "h * (d/h)^2 per layer"
    },
    "v6 (Finch)": {
        "year": 2024,
        "max_params": "7.5B",
        "key_feature": "Data-dependent time decay, LoRA-style mixing",
        "state_size": "h * (d/h)^2 per layer"
    },
    "v7 (Goose)": {
        "year": 2025,
        "max_params": "2.9B (scaling ongoing)",
        "key_feature": "Dynamic state evolution via generalized delta rule",
        "state_size": "h * (d/h)^2 per layer (dynamic)"
    }
}

Tokens/sec at various sequence lengths (3B params, A100 GPU):

Seq Length  | RWKV-6  | Transformer | Mamba-2  | Note
------------|---------|-------------|----------|------------------
512         |  2,400  |    2,200    |  2,500   | All comparable
2,048       |  2,350  |    1,800    |  2,450   | Transformer slowing
8,192       |  2,300  |      950    |  2,400   | KV cache pressure
32,768      |  2,250  |      280    |  2,350   | Transformer struggling
131,072     |  2,200  |      OOM    |  2,300   | Transformer OOM
524,288     |  2,100  |      OOM    |  2,200   | Both linear models OK

Memory Usage (inference state):

Model       | 1K ctx  | 8K ctx  | 128K ctx | 1M ctx
------------|---------|---------|----------|--------
RWKV-6 3B   | 6.2 GB  | 6.2 GB  |  6.2 GB  | 6.2 GB  (constant!)
Transformer | 6.5 GB  | 8.1 GB  | 32.4 GB  |  OOM
Mamba-2 3B  | 6.3 GB  | 6.3 GB  |  6.3 GB  | 6.3 GB  (constant)

Model               | Params | Pile Val PPL | LAMBADA | HellaSwag
--------------------|--------|--------------|---------|----------
RWKV-4              |  7B    |    8.28      |  67.2%  |  52.5%
RWKV-5 Eagle        |  7B    |    8.15      |  68.1%  |  53.2%
RWKV-6 Finch        |  7B    |    8.05      |  69.0%  |  54.1%
RWKV-7 Goose        |  2.9B  |    9.60      |  65.8%  |  51.0%
Pythia (Transformer) |  6.9B |    8.25      |  67.1%  |  52.0%
LLaMA-like          |  7B    |    7.95      |  73.0%  |  56.4%
Mamba               |  2.8B  |    9.80      |  64.9%  |  50.3%

# Clone the RWKV-LM repository
git clone https://github.com/BlinkDL/RWKV-LM.git
cd RWKV-LM/RWKV-v7

# Install dependencies
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install lightning deepspeed wandb ninja

# For custom CUDA kernel compilation
pip install triton

# Verify CUDA kernel builds
python -c "from rwkv.model import RWKV; print('RWKV loaded successfully')"

# train_config.py -- Example training configuration
import lightning as L
from rwkv.model import RWKV

# Model configuration
model_config = {
    "n_layer": 24,         # Number of RWKV blocks
    "n_embd": 2048,        # Embedding dimension
    "vocab_size": 65536,   # Vocabulary size
    "ctx_len": 4096,       # Context length for training
    "head_size": 64,       # Head size for multi-headed WKV (v5+)
}

# Training hyperparameters
train_config = {
    "learning_rate": 6e-4,          # Peak LR
    "lr_schedule": "cosine",        # Cosine decay
    "warmup_steps": 1000,
    "batch_size": 16,
    "accumulate_grad_batches": 4,   # Effective batch = 64
    "max_steps": 100000,
    "precision": "bf16-mixed",      # BF16 mixed precision
    "gradient_clip_val": 1.0,
    "weight_decay": 0.1,
    "beta1": 0.9,
    "beta2": 0.99,
}

# DeepSpeed configuration for multi-GPU
deepspeed_config = {
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"},
    },
    "bf16": {"enabled": True},
}

# Clone RWKV-PEFT
git clone https://github.com/JL-er/RWKV-PEFT.git
cd RWKV-PEFT

# Prepare training data in binidx format
python tools/preprocess_data.py \
    --input your_data.jsonl \
    --output-prefix train_data \
    --tokenizer-type RWKVTokenizer \
    --vocab-size 65536

# Launch LoRA fine-tuning
python train.py \
    --load_model /path/to/rwkv-base-model.pth \
    --proj_dir output/ \
    --data_file train_data \
    --data_type binidx \
    --ctx_len 2048 \
    --n_layer 24 \
    --n_embd 2048 \
    --lora_r 64 \
    --lora_alpha 128 \
    --lora_parts att,ffn \
    --micro_bsz 4 \
    --epoch_steps 1000 \
    --epoch_count 5 \
    --lr_init 2e-4 \
    --lr_final 2e-5 \
    --strategy deepspeed_stage_1 \
    --precision bf16

python merge_lora.py \
    --base_model /path/to/rwkv-base-model.pth \
    --lora_checkpoint output/rwkv-lora-final.pth \
    --output merged_model.pth \
    --lora_r 64 \
    --lora_alpha 128

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load RWKV from HuggingFace Hub
model = AutoModelForCausalLM.from_pretrained(
    "RWKV/rwkv-6-world-7b",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(
    "RWKV/rwkv-6-world-7b",
    trust_remote_code=True
)

# Generate text
prompt = "The RWKV architecture is interesting because"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    temperature=0.8,
    top_p=0.9,
    do_sample=True
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Fine-tuning VRAM requirements (LoRA, ctx_len=2048):

Model Size  | LoRA r=8  | LoRA r=64 | Full Fine-tune
------------|-----------|-----------|----------------
1.5B        |   4 GB    |   6 GB    |   16 GB
3B          |   6 GB    |  10 GB    |   28 GB
7B          |  12 GB    |  18 GB    |   56 GB
14B         |  22 GB    |  34 GB    |  112 GB

Task: "Find the phone number mentioned on page 37 of a 100-page document"

Transformer (GPT-4 class): Correctly retrieves the number (full attention)
RWKV-6 7B:                 Often fails or hallucinates a similar number
Mamba 7B:                  Sometimes succeeds (selective state helps)

Root cause: The exponential decay in WKV means old information is
progressively "forgotten" unless it strongly activates key channels.
The fixed state size (d dimensions) cannot store arbitrary facts.

Prompt: "Alice gave Bob a red ball. Charlie gave Diana a blue cube.
         ... (500 tokens of distraction) ...
         Eve traded her green cone for the object Diana has.
         What color is the object Eve now has?"

Transformer: "Blue" (correct -- attends to both relevant sentences)
RWKV:        "Green" or "Red" (may lose track of multi-hop chain)

# This format works well with RWKV
prompt_good = "User: What is the capital of France?\n\nAssistant:"

# This format may produce worse results
prompt_bad = "capital of france?"

# RWKV is sensitive to newlines, spacing, and role markers.
# Always use consistent chat templates when deploying RWKV models.

class RWKVChatSession:
    """Manage RWKV state across a multi-turn conversation."""

    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.state = None  # Will hold the running RNN state

    def chat(self, user_message):
        # Format the message
        prompt = f"User: {user_message}\n\nAssistant:"
        input_ids = self.tokenizer.encode(prompt)

        # Feed tokens through model, updating state
        output_ids = []
        for token_id in input_ids:
            logits, self.state = self.model.forward(
                token_id, self.state
            )

        # Generate response tokens
        for _ in range(500):
            token_id = self.sample(logits, temperature=0.8)
            if token_id == self.tokenizer.eos_token_id:
                break
            output_ids.append(token_id)
            logits, self.state = self.model.forward(
                token_id, self.state
            )

        return self.tokenizer.decode(output_ids)

    def save_state(self, path):
        """Save conversation state to disk for later resumption."""
        torch.save(self.state, path)

    def load_state(self, path):
        """Resume conversation from saved state."""
        self.state = torch.load(path)

    @staticmethod
    def sample(logits, temperature=1.0, top_p=0.9):
        probs = torch.softmax(logits / temperature, dim=-1)
        sorted_probs, sorted_indices = torch.sort(probs, descending=True)
        cumsum = torch.cumsum(sorted_probs, dim=-1)
        mask = cumsum - sorted_probs > top_p
        sorted_probs[mask] = 0.0
        sorted_probs /= sorted_probs.sum()
        idx = torch.multinomial(sorted_probs, 1)
        return sorted_indices[idx].item()

# Using the rwkv.cpp project for CPU inference
git clone https://github.com/saharNooby/rwkv.cpp.git
cd rwkv.cpp

# Quantize model to INT4
python python/convert_pytorch_to_ggml.py \
    /path/to/rwkv-model.pth \
    /path/to/output.bin \
    FP16

python python/quantize.py \
    /path/to/output.bin \
    /path/to/output-q4_0.bin \
    Q4_0

# Run inference on CPU
python python/chat.py \
    --model /path/to/output-q4_0.bin