はじめに

こんにちは、RetailAIの@long10langです。さて、advent calendar今年は何書こう？というわけで、あれこれ考えた結果、あまり理解が進んでいなかったRWKVについてちゃんと調べてみようと思い立ちまして、それなら単にRWKVの説明をしたところで面白くないので、juliaでいっちょやってみっかということで、やって行きたいと思います。

RWKVについて

やはり、今年はなんと言ってもChatGPT元年ということで、いろんなことがありました。もうすっかり市民権を獲得したと言っても過言ではないtransformerですが、その根幹をなすSelf-Attentionという仕組みは、もちろん画期的な発明ではありましたが、学習量がすんごいことになるため、一般人がそう簡単に手出しできるような代物ではありませんよね。

とはいえ、RNNベースのモデルだと、メモリと計算要件の面で線形にスケールするものの、並列化と拡張性の制限からtransformerと同等の性能を達成できるかといえばそれもまた困難です。

そこで，transformerの効率的な並列学習と、RNNの効率的な推論の両方を兼ね備えたモデルとしてRWKV（Receptance Weighted Key Value）という新たなモデルアーキテクチャーが提案されました。

詳しくは論文を参照して頂ければと思うのですが、ざっくり言ってしまえば、このモデルは、数百億のパラメータまでスケールする初の非Transformerアーキテクチャでありながら、同じサイズのTransformerと同等の性能を発揮することが論文内で示されています。

そんなわけで、現在GPTをはじめとしたtransformerベースのモデルよりも、高速に推論可能（？）なモデルとして注目されている、というわけです。

左がtransformerで、右がRWKVです。Channel MixingとTime Mixingが大きく異なる点ですが、そこに登場するモデル要素にちなんで、RWKVと呼ばれています。

R：過去の情報の受容度を表現するReceptanceベクトル。
W：位置の重み減衰ベクトル。訓練可能なモデルパラメータ。
K：一般的な注意機構におけるK（Key）に類似のベクトル。
V：一般的な注意機構におけるV（Value）に類似のベクトル。

めちゃくちゃざっくりと言ってしまえば、tranformerもRWKVも、どちらも計算上で、文章などの情報を理解する方法という点では一緒なのですが、tranformerはSelf-Attention、つまり文章のすべての部分に注意を払い、複雑な内容も理解できる反面たくさんの計算が必要になる。一方、RWKVはよりシンプルで、少ない計算で済むためサクサク動作します。というわけで、長い文章を扱うときRWKVの方が効率的ですが、transformerの方がより深く内容を理解できる、乱暴ですがそんな理解です。

実装あれこれ

それでは、四の五の言わず、juliaで実装を見ていきましょう。まずは、ネットワークから。

１．Layer Normalization (LN) 構造体と関数：層正規化を行うカスタム構造体とその動作を定義しています。

struct LN
    γ
    β
    ϵ
    dims
end
LN(γ, β) = LN(γ, β, 1e-5, 1)
LN(n::Integer) = LN(ones(Float32, n), zeros(Float32, n))
@Flux.functor LN
Flux.trainable(m::LN) = (γ=m.γ, β=m.β)
function (m::LN)(x::AbstractArray{T,N}) where {T,N}
    μ = mean(x, dims=m.dims)
    σ² = var(x, dims=m.dims)
    (T).(m.γ .* (x .- μ) ./ (σ² .+ m.ϵ).^0.5 .+ m.β)
end

2.次に、ユーティリティ関数：時間の混合、二乗ReLU、指数混合などの関数を定義していきます。

time_mix(x, x_prev, mix) = @. x * mix + x_prev * (1 - mix)
square_relu(x::T) where T = max(zero(T), x)^2
function exp_mix(v1::AbstractArray{T}, v2::AbstractArray{T}, p1::AbstractArray{T}, p2::AbstractArray{T}) where T
    p = max.(p1, p2)
    (@. exp(p1 - p) * v1 + exp(p2 - p) * v2, p)
end

function exp_selfadj(v1, v2, p)
    p_new = @. max(min(abs(asinh(v1)), p), min(abs(asinh(v2)), p))
    factor = @. exp(p - p_new)
    @. (v1*factor, v2*factor, p_new)
end

3.State構造体：ネットワークの状態を保持するための構造体を定義します。

mutable struct State
    x_tm  # token mixing
    x_cm  # channel mixing
    a # numerator
    b # denominator
    p # largest exponent seen
end

State(n_embed::Integer, n_layer::Integer) = begin
    dim = (n_embed, n_layer)
    State(zeros(Float32, dim), zeros(Float32, dim), zeros(Float32, dim), zeros(Float32, dim), zeros(Float32, dim))
end
@Flux.functor State

4.TokenMixing と ChannelMixing 構造体：トークンとチャネルの混合を行うための層を定義しています。

function recur_step(left::Vector, right::Vector; w)
    a_prev, b_prev, p_prev = left
    expkv, expk, p = right
    a_new, p_new = exp_mix(a_prev, expkv, p_prev .+ w, p)
    b_new, _ = exp_mix(b_prev, expk, p_prev .+ w, p)
    [a_new, b_new, p_new]
end

struct TokenMixing{T}
    Tₖ::AbstractArray{T, 1}
    Tᵥ::AbstractArray{T, 1}
    Tᵣ::AbstractArray{T, 1}
    r_proj
    k_proj
    v_proj
    out_proj 
    time_first::AbstractArray{T, 1}
    time_decay::AbstractArray{T, 1}  # <-- w
end

@Flux.functor TokenMixing

TokenMixing(n_embed::Integer) = TokenMixing(
    zeros(Float32, n_embed), # Tₖ
    zeros(Float32, n_embed), # Tᵥ
    zeros(Float32, n_embed), # Tᵣ
    Dense(n_embed, n_embed, bias=false), # r_proj
    Dense(n_embed, n_embed, bias=false), # k_proj
    Dense(n_embed, n_embed, bias=false), # v_proj
    Dense(n_embed, n_embed, bias=false), # out_proj
    zeros(Float32, n_embed), # time first
    ones(Float32, n_embed),  # time_decay
)

function (m::TokenMixing)(x::AbstractArray{T,2}, state::State; i) where T
    n_embed, n_seq = size(x)

    x_prev = hcat(@views(state.x_tm[:, i]), @views(x[:, 1:end-1]))

    xₖ = time_mix(x, x_prev, m.Tₖ)
    xᵥ = time_mix(x, x_prev, m.Tᵥ)
    xᵣ = time_mix(x, x_prev, m.Tᵣ)

    r = m.r_proj(xᵣ) .|> sigmoid
    k = m.k_proj(xₖ)
    v = m.v_proj(xᵥ)

    p = k
    expk = k.*0 .+ 1
    expkv = v

    step_f = (a, b) -> recur_step(a, b; w=m.time_decay)
    a_prev, b_prev, p_prev = @views(state.a[:, i]), @views(state.b[:, i]), @views(state.p[:, i])
    substrate = [[@views(expkv[:,i]), @views(expk[:,i]), @views(p[:,i])] for i = 1:n_seq]
    abp = accumulate(step_f, substrate; init=[a_prev, b_prev, p_prev])

    a_prev = batch([a_prev, [abp[i][1] for i = 1:n_seq-1]...])
    b_prev = batch([b_prev, [abp[i][2] for i = 1:n_seq-1]...])
    p_prev = batch([p_prev, [abp[i][3] for i = 1:n_seq-1]...])

    c, _ = exp_mix(a_prev, expkv, p_prev, p .+ m.time_first)
    d, _ = exp_mix(b_prev, expk,  p_prev, p .+ m.time_first)
    rwkv = @. r * c / d

    # update state
    @views state.x_tm[:, i] .= x[:, end]
    @views state.a[:, i] .= abp[end][1]
    @views state.b[:, i] .= abp[end][2]
    @views state.p[:, i] .= abp[end][3]

    m.out_proj(rwkv), state
end
    
# stateless for training
function (m::TokenMixing)(x::AbstractArray{T,3}) where T
    n_embed, n_batch, n_seq = size(x)

    @views x_prev = pad_zeros(x, (0,0,0,0,1,0))[:, :, 1:end-1]

    xₖ = time_mix(x, x_prev, m.Tₖ)
    xᵥ = time_mix(x, x_prev, m.Tᵥ)
    xᵣ = time_mix(x, x_prev, m.Tᵣ)

    r = m.r_proj(xᵣ) .|> sigmoid
    k = m.k_proj(xₖ)
    v = m.v_proj(xᵥ)

    p = k
    expk = zeros_like(k) .+ 1
    expkv = v

    step_f = (a, b) -> recur_step(a, b; w=m.time_decay)
    a_prev = b_prev = p_prev = zeros_like(k, eltype(v), (n_embed, n_batch))
    substrate = @views [[expkv[:,:,i], expk[:,:,i], p[:,:,i]] for i = 1:n_seq]
    abp = accumulate(step_f, substrate; init=[a_prev, b_prev, p_prev])

    a_prev = batch([a_prev, [abp[i][1] for i = 1:n_seq-1]...])
    b_prev = batch([b_prev, [abp[i][2] for i = 1:n_seq-1]...])
    p_prev = batch([p_prev, [abp[i][3] for i = 1:n_seq-1]...])

    c, _ = exp_mix(a_prev, expkv, p_prev, p .+ m.time_first)
    d, _ = exp_mix(b_prev, expk,  p_prev, p .+ m.time_first)
    rwkv = @. r * c / d

    m.out_proj(rwkv)
end

struct ChannelMixing{T}
    Tₖ::AbstractArray{T, 1}  # will be taken out in the future
    Tᵣ::AbstractArray{T, 1}  # will be taken out in the future
    r_proj
    k_proj
    v_proj
end

@Flux.functor ChannelMixing

ChannelMixing(n_embed::Integer) = ChannelMixing(
    zeros(Float32, n_embed), # Tₖ
    zeros(Float32, n_embed), # Tᵣ
    Dense(n_embed, n_embed, bias=false), # r_proj
    Dense(n_embed, n_embed*4, bias=false), # k_proj
    Dense(n_embed*4, n_embed, bias=false), # v_proj
)

function (m::ChannelMixing)(x::AbstractArray{T, 2}, state::State; i) where T
    n_embed, n_seq = size(x)

    x_prev = @views(state.x_cm[:, i])
    if size(x, 2) > 1
        x_prev = hcat(x_prev, @views(x[:, 1:end-1]))
    end
    xₖ = time_mix(x, x_prev, m.Tₖ)
    xᵣ = time_mix(x, x_prev, m.Tᵣ)

    r = m.r_proj(xᵣ) .|> sigmoid
    k = m.k_proj(xₖ) .|> square_relu

    # update state
    @views state.x_cm[:, i] .= x[:, end]

    r .* (m.v_proj(k)), state
end

function (m::ChannelMixing)(x::AbstractArray{T, 3}) where T
    # n_embed, n_batch, n_seq = size(x)
    @views x_prev = pad_zeros(x, (0,0,0,0,1,0))[:, :, 1:end-1]

    xₖ = time_mix(x, x_prev, m.Tₖ)
    xᵣ = time_mix(x, x_prev, m.Tᵣ)

    r = m.r_proj(xᵣ) .|> sigmoid
    k = m.k_proj(xₖ) .|> square_relu

    r .* (m.v_proj(k))
end

5.Block 構造体：LN、TokenMixing、ChannelMixingを組み合わせたブロック構造体を定義します。

struct Block
    ln1
    token_mixing
    ln2
    channel_mixing
end

@Flux.functor Block

Block(n_embed::Integer) = Block(
    LN(n_embed),
    TokenMixing(n_embed),
    LN(n_embed),
    ChannelMixing(n_embed),
)

function (m::Block)(x, state::State; i)
    xp, state = m.token_mixing(m.ln1(x), state; i=i)
    x = x + xp
    xp, state = m.channel_mixing(m.ln2(x), state; i=i)
    x = x + xp 
    x, state
end

function (m::Block)(x)
    xp = m.token_mixing(m.ln1(x))
    x = x + xp
    xp = m.channel_mixing(m.ln2(x))
    x = x + xp
    x
end

6.RWKVモデル：言語モデルのメイン構造体で、埋め込み層、LN層、複数のブロック、最終LN層、言語モデルのヘッドを含んでいます。

struct RWKV
    ln_init
    embedding
    blocks
    ln_final
    lm_head
end

@Flux.functor RWKV

RWKV(n_embed::Integer, n_blocks::Integer, n_vocab::Integer) = RWKV(
    Embedding(n_vocab, n_embed),
    LN(n_embed),
    [Block(n_embed) for _ in 1:n_blocks],
    LN(n_embed),
    Embedding(n_embed, n_vocab)
)

7.モデルのフォワードパス：状態を持つか持たないかに基づいて、モデルのフォワードパスを定義しています。

(m::RWKV)(x::AbstractArray{T, 1}, state::State) where T = begin
    x = m.embedding(x)
    x = m.ln_init(x)
    for i in 1:length(m.blocks)
        x, state = m.blocks[i](x, state; i=i)
    end
    x = m.ln_final(x)

    # x: [n_embed, n_seq]
    x = m.lm_head.weight' * x

    x, state
end
(m::RWKV)(x::Integer, state::State) = begin
    out, state = m([x], state)
    out[:, end], state
end

(m::RWKV)(x::AbstractArray{T, 2}) where T = begin
    x = m.embedding(x)
    # better performance
    x = permutedims(x, (1, 3, 2))
    x = m.ln_init(x)
    for i in 1:length(m.blocks)
        x = m.blocks[i](x)
    end
    x = m.ln_final(x)
    # better performance
    x = permutedims(x, (1, 3, 2))
    # x: [n_embed, n_seq, n_batch]
    batched_mul(m.lm_head.weight', x)
end

以上で、RWKVのネットワーク実装が終わりました。では、早速このネットワークを用いてトレーニングをしていきましょう。

include("data.jl") # <- データ
include("rwkv.jl") # <- 上で定義したネットワーク
include("utils.jl") # <- 下で定義するユーティリティ関数

using Flux.Losses
using OneHotArrays
using Dates, Logging
using CUDA
using BSON: @save
using Optimisers

function save_model(model, path)
    model = model |> cpu
    @save path model
end

cfg = (
    learning_rate = 3e-4,
    epochs = 5,
    dataset = "sample.jsonl",
    use_cuda = true,
    ctx_len = 256,
    batch_size = 8,
    log_per_nbatch = 5,
    save_per_nbatch = 1000,
    n_vocab = 50277,
    n_layer = 12,
    n_head = 12
)

data_file = cfg.dataset
tokenizer = get_tokenizer()
ts = TextSplitter(cfg.ctx_len+1, 16, tokenizer)

function loss_func(y_pred, y; n_vocab=cfg.n_vocab)
    return logitcrossentropy(y_pred, onehotbatch(y, 1:n_vocab))
end

if cfg.use_cuda && CUDA.functional()
    CUDA.allowscalar(false)
    @info "Using CUDA"
    device = gpu
else
    @info "Using CPU"
    device = identity
end

Flux.trainable(m::RWKV) = (blocks=m.blocks,)
Flux.trainable(m::Block) = (token_mixing=m.token_mixing,)

# device = identity  # debug
model = rwkv_from_pth("RWKV-4-Pile-169M-20220807-8023.pth"; cfg.n_layer) |> f32
model = model |> device

# setup optimizer
opt = Optimisers.setup(Optimisers.Adam(cfg.learning_rate), model)
i_batch = 1
for epoch in 1:cfg.epochs
    global i_batch
    batches = jsonl_reader(data_file) |> ch->batch_sampler(ch, ts; batch_size=cfg.batch_size)
    for (x, y) in batches
        x = x |> device
        y = y |> device
        val, grads = Flux.withgradient(model) do m
            y_pred = m(x)
            loss_func(y_pred, y)
        end
        Flux.update!(opt, model, grads[1])
        if i_batch % cfg.log_per_nbatch == 0
            @info "metrics" loss=val
            println("Ep $(epoch) Batch $(i_batch) Loss $(val)")
        end
        if i_batch % cfg.save_per_nbatch == 0
            @info "saving model"
            save_model(model, "rwkv-169m-$(now()).bson")
        end
        i_batch += 1
    end
end

ユーティリティ関数（utils.jl）がこちら↓

using PyCall
using Flux
using StatsBase

include("rwkv.jl")
torch = pyimport("torch")

function rwkv_from_pth(pth_path="RWKV-4-Pile-169M-20220807-8023.pth"; n_layer=12)
    data = torch.load(pth_path, "cpu")

    ln_init = LN(
        data["blocks.0.ln0.weight"].float().numpy(),
        data["blocks.0.ln0.bias"].float().numpy(),
    )
    ln_final = LN(
        data["ln_out.weight"].float().numpy(),
        data["ln_out.bias"].float().numpy()
    )
    embedding = Embedding(
        data["emb.weight"].float().numpy()'
    )
    lm_head = Embedding(
        data["head.weight"].float().numpy()'
    )
    blocks = []
    for i = 0:n_layer-1
        ln1 = LN(
            data["blocks.$i.ln1.weight"].float().numpy(),
            data["blocks.$i.ln1.bias"].float().numpy(),
        )
        ln2 = LN(
            data["blocks.$i.ln2.weight"].float().numpy(),
            data["blocks.$i.ln2.bias"].float().numpy(),
        )
        time_first = data["blocks.$i.att.time_first"].float().numpy()
        time_decay = -exp.(data["blocks.$i.att.time_decay"].float().numpy())
        token_mixing = TokenMixing(
            dropdims(data["blocks.$i.att.time_mix_k"].float().numpy(), dims=(1,2)),
            dropdims(data["blocks.$i.att.time_mix_v"].float().numpy(), dims=(1,2)),
            dropdims(data["blocks.$i.att.time_mix_r"].float().numpy(), dims=(1,2)),
            Dense(
                data["blocks.$i.att.receptance.weight"].float().numpy(),
                false
            ), # r_proj
            Dense(
                data["blocks.$i.att.key.weight"].float().numpy(),
                false
            ), # k_proj
            Dense(
                data["blocks.$i.att.value.weight"].float().numpy(),
                false
            ), # v_proj
            Dense(
                data["blocks.$i.att.output.weight"].float().numpy(),
                false
            ), # out_proj
            time_first,
            time_decay,
        )
        channel_mixing = ChannelMixing(
            dropdims(data["blocks.$i.ffn.time_mix_k"].float().numpy(), dims=(1,2)),
            dropdims(data["blocks.$i.ffn.time_mix_r"].float().numpy(), dims=(1,2)),
            Dense(
                data["blocks.$i.ffn.receptance.weight"].float().numpy(),
                false
            ), # r_proj
            Dense(
                data["blocks.$i.ffn.key.weight"].float().numpy(),
                false
            ), # k_proj
            Dense(
                data["blocks.$i.ffn.value.weight"].float().numpy(),
                false
            ), # v_proj
        )
        push!(blocks, Block(
            ln1,
            token_mixing,
            ln2,
            channel_mixing
        ))
    end
    RWKV(
        ln_init,
        embedding,
        blocks,
        ln_final,
        lm_head
    )
end

function get_tokenizer()
    py"""
    from tokenizers import Tokenizer

    tokenizer = Tokenizer.from_file('20B_tokenizer.json')
    """
    py"tokenizer"
end

function sample_logits(logits; temperature=1.0, top_p=0.9, use_argmax=false)
    if use_argmax
        return argmax(logits)
    end
    probs = softmax(logits; dims=1)
    sorted_probs = sort(probs; rev=true)
    cum_probs = cumsum(sorted_probs)
    cutoff = sorted_probs[argmax(cum_probs .> top_p)]
    probs = (probs .> cutoff) .* probs
    if temperature != 1.0
        probs .^= 1/temperature
    end
    probs = probs / sum(probs)

    sample(collect(1:length(probs)), ProbabilityWeights(probs)) |> Int
end

上記をトレーニングをして、推論を実行すると、結果が表示されます。

include("rwkv.jl")
include("utils.jl")

device = gpu
tokenizer = get_tokenizer()
model = rwkv_from_pth("RWKV-4-Pile-430M-20220808-8066.pth"; n_layer=24) |> device

function generate(model, prompt, n_tokens=50; top_p=0.99, temperature=1.0, use_argmax=false, device=identity)
    input_ids = tokenizer.encode(prompt).ids .+ 1

    state = State(size(model.embedding.weight, 1), length(model.blocks)) |> device

    out, state = model(input_ids[1:end-1], state)

    println("-------------------------")
    println(prompt)
    println("-------------------------")

    input_id = input_ids[end]
    for i = 1:n_tokens
        out, state = model(input_id, state)
        out_id = sample_logits(out |> cpu; top_p=top_p, temperature=temperature, use_argmax=use_argmax)
        print(tokenizer.decode([out_id-1]))
        input_id = out_id
    end
    println()
    state
end

prompt = "Japan is an economic powerhouse, but its growth rate is slowing. In the days ahead, Japan will";
generate(model, prompt, 100; top_p=0.999, device=device);

プロンプトとして、「Japan is an economic powerhouse, but its growth rate is slowing. In the days ahead, Japan will（日本は経済大国だが成長率は鈍化している。この先日本は・・・）」と渡してみたところ、返答として「 be dragged down by the decrease in exports and by the fall in the official growth rate to 6% in FY 2011. If the Japanese government continues its policies in this area, growth rates will drop below the previous rate in all of FY 2011（輸出の減少と、2011年度の公式成長率が6%に低下したことが足を引っ張った。日本政府がこの分野の政策を続ければ、2011年度の成長率はすべて以前の成長率を下回るだろう。）」と、正しいのかよくわかりませんがそれっぽいのが返ってきました。

苦労話

RWKVのネットワーク実装については、RWKV-LMをほぼパクってますので、実態としてはpyCallしまくってます（笑）juliaを使うメリットは、ちょっとjuliaだけでは実装が厳しいとなっても、すぐPythonでなんとかできる点が良い点です。(いざとなればCが呼べるのもの嬉しいですね。)

ただ、PyCallの呼び出しで、実行させたいPythonのバージョンを指定するのに、以下のようにREPLで環境変数を指定してあげた後に、　(@hoge) pkg> build PyCall といった感じでビルドしてあげる必要があります。

               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.9.4 (2023-11-14)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org/ release
|__/                   |

julia> ENV["PYTHON"]="/home/long/.local/share/virtualenvs/XXXX/bin/python"

個人的には、juliaは、言語としてPythonよりも構文がはっきりと構造的だし、パイプ演算子があることで、状態変化の見通しがつきやすい点で、処理速度が速いというだけでなく、読みやすさという点でも、Pythonより上回ってると感じてます。どんどん使っていきたい。

まとめ

いかがでしたでしょうか？RWKVをより実装レベルで理解するために、Juliaで書き下してみることで、様々な気づきがありました。生成AI関連は来年もきっとすごいことになって行きそうですので、こんな感じで好きなプログラミング言語をあえて使いながら、面白がってキャッチアップしていきたいと思います。

さて明日は、サーバサイドエンジニアの@mizoguchi_ryosukeさんが、「初めてのpython cvxpy」というタイトルで発表します！ぜひ、お楽しみください〜

RWKVをjuliaで試してみる

はじめに

目次

RWKVについて

実装あれこれ

苦労話

まとめ

参考文献