RWKV가 하려는 일

이중 표현을 다른 방향에서 풀었어

RWKV는 Receptance Weighted Key Value의 약자고 RwaKuv라고 읽어. 처음부터 목표가 분명한 linear attention과 RNN의 하이브리드야. Transformer처럼 병렬로 학습하고, 추론할 때는 RNN으로 실행해. SSM에도 비슷한 이중성이 있지만 RWKV는 구조상 어텐션에 더 가까워. R, W, K, V라는 매개변수 이름도 우연이 아니야. 어텐션의 QKV를 일부러 비튼 이름이지.

네 매개변수가 역할을 나눠. R은 이 토큰이 누적 요약을 얼마나 받아들일지 정하는 수용 게이트야. W는 오래된 토큰을 얼마나 빨리 흐리게 할지 정하는 지수 감쇠 가중치고, K와 V는 key와 value야. WKV는 시간이 흐를수록 감쇠를 적용하면서 가중 key-value 합을 갱신해. T×T 행렬을 만들 일은 없어.

SSM과 같은 장점을 다른 말로 설명해

SSM은 제어 이론의 말로 효율을 설명해. 과거를 상태로 압축한다는 식이지. RWKV는 같은 이득을 어텐션에 가까운 개념으로 설명해. 지수 감쇠를 적용하는 순환 상태로 어텐션을 근사한다고 보는 거야. SSD가 보여 주듯 수학은 이웃해 있지만 공학적인 감촉과 사용 방식은 달라. RWKV 모델은 Transformer 형태의 체크포인트처럼 보이고 llama.cpp, MLX, ggml 계열 기술 스택에 비교적 쉽게 옮길 수 있어. 상태 안정성을 따로 다루지 않고 RNN처럼 실행할 수 있다는 점도 실용적이야.

토큰마다 드는 시간이 일정해

RWKV가 아주 깔끔하게 보장하는 이득이 있어. 추론 비용이 토큰마다 일정해. 고정 크기 상태를 가진 RNN과 증명 가능한 방식으로 같기 때문에 생성 지연 시간은 문맥이 길어져도 평평해. 기기 내 서비스에서 이 성질은 강력해. 첫 번째 토큰 뒤를 만들든 100,000번째 토큰 뒤를 만들든 다음 토큰의 비용이 같거든.

Code

RWKV의 WKV recurrence (단순화)·python

# 매 timestep t, input x_t 주어졌을 때:
#   k_t, v_t = K(x_t), V(x_t)
#   r_t = sigmoid(R(x_t))
# 두 running quantity 유지 ('state'):
#   num   <- num * exp(-w) + exp(k_t) * v_t
#   denom <- denom * exp(-w) + exp(k_t)
# Output:
#   wkv_t = num / denom
#   y_t   = r_t * wkv_t
# 'num' 과 'denom' 이 전체 state. t 무관 constant size.

Exercise

Hugging Face에서 작은 RWKV 모델을 받아 글을 생성해 봐. RWKV/rwkv-4-169m-pile이면 충분해. 프롬프트 길이를 256, 8K, 64K 토큰으로 늘리면서 새 토큰 하나의 지연 시간을 재. WKV 상태 크기가 같으니 지연 시간은 거의 평평해야 해. 프리필은 제곱으로 느려지고 KV 캐시 메모리는 선형으로 커지는 Transformer와 비교해 봐.