RWKV - Transformer 시대에 RNN을 재발명하다

Digest: Transformer는 O(T²) 어텐션 때문에 긴 컨텍스트에서 메모리·연산 폭증에 시달리고, 전통적 RNN은 O(Td) 선형 추론은 가능하지만 시간축 병렬 학습이 불가능해 대규모 스케일링에서 밀려났다. RWKV(Receptance-Weighted Key-Value)의 핵심 통찰은 “소프트맥스 어텐션의 지수 가중합 구조를 채널별 시간감쇠(channel-wise time-decay)로 분해된 재귀식으로 재작성하면, 동일 연산이 학습 시엔 시간-병렬 형태(time-parallel mode)로 풀리고 추론 시엔 순수 RNN 상태 업데이트(recurrent mode)로 풀린다”는 이중 형식(duality) 이다. 이를 위해 각 층을 time-mixing 블록(R, K, V + WKV 재귀)과 channel-mixing 블록(R, K, V 게이팅 MLP)으로 구성하고, receptance(수용도) σ(R) 가 과거 정보의 게이트 역할을 한다. 결과적으로 169M ~ 14B 파라미터까지 스케일 업하여 Pythia·GPT-Neo 등 비슷한 연산량의 순수 Transformer와 대등한 zero-shot 언어모델 성능을 보이면서, 추론은 시퀀스 길이에 선형(O(Td)), 메모리는 상수(O(d)) 로 유지된다. 한계로 저자들은 (1) 단일 히든 상태에 과거를 압축하기 때문에 소프트맥스 어텐션 대비 정밀한 장거리 recall 성능 저하 우려, (2) 프롬프트 엔지니어링에 민감한 경향을 인정한다. 열린 질문으로는, 더 큰 재귀 상태(State Space Model, RetNet, Mamba 류)로의 일반화, 표현력의 이론적 상한, multi-modal/긴 reasoning 태스크에서의 한계 등이 남아 있다.

섹션별 요약

Introduction

문제: Transformer는 self-attention의 O(T²d) 비용으로 긴 문맥/실시간 추론에서 비싸고, 전통 RNN은 병렬화가 어려워 LLM 스케일에서 밀림.
기여: (1) RNN처럼 추론하고 Transformer처럼 학습하는 이중 형식 아키텍처 RWKV, (2) 14B까지 공개된 최초의 비-Transformer 대형 LLM, (3) 선형-어텐션/AFT 계열과 RNN의 다리 역할.

Methods

블록 구성: Layer = LayerNorm → time-mixing → LayerNorm → channel-mixing (pre-LN residual).
Token shift: 현재 토큰 $x_{t}$ 와 이전 토큰 $x_{t - 1}$ 을 학습 가능한 비율 $μ$ 로 선형 결합해 R/K/V 입력에 쓴다 → 얕은 1-D conv에 해당.
Time-mixing: $r_{t}, k_{t}, v_{t}$ 를 계산하고 WKV 연산으로 과거를 지수 감쇠로 누적.
Channel-mixing: 위치 고정 GLU-류 MLP. $σ (r) ⊙ (W_{v} \cdot max (0, k)^{2})$ .
Receptance gating: $σ (r_{t})$ 가 “얼마나 받아들일지”를 채널별로 조절.
학습↔추론 이중성: WKV는 분자/분모 각각의 지수 누적합으로 병렬 계산 가능하며, 동시에 상태 $a_{t}, b_{t}$ 로 O(1) 스트리밍 업데이트 가능.

Results

Model	Params	Tokens	LAMBADA (acc, ppl)	PIQA	HellaSwag	WinoGrande
Pythia	1.4B	300B	61.7	71.1	52.1	57.3
RWKV-4	1.5B	332B	56.4	72.4	52.5	54.6
Pythia	6.9B	300B	67.3	75.1	63.9	61.3
RWKV-4	7.5B	332B	67.2	76.1	65.6	61.0
RWKV-4	14B	332B	70.8	77.4	70.4	63.9

(수치는 논문 Table 4/Appendix 참고치 — 동일 연산 예산에서 Transformer와 대등하며 14B에서 더 강해짐)

추론 효율: 토큰당 O(d) 연산 + 상수 메모리. 소프트맥스 어텐션 대비 long-context 생성에서 선형 대 2차 스케일링 이득.
학습: Transformer처럼 시퀀스 병렬, 약 1.2× ~ 비슷한 wall-clock.

Discussion

한계 (저자 인정): (1) 고정 크기 히든 상태 → 세밀한 장거리 look-back 부족, (2) 초기화/프롬프트 민감, (3) 고정 시간감쇠 벡터 $w$ 가 “과거 망각”을 구조적으로 강제.
향후: 상태 크기 확장, 입력-의존 감쇠(→ RetNet/Mamba 방향), encoder/multimodal RWKV.

Insights

주목할 점: 소프트맥스를 “exp-weighted sum”으로 다시 쓰면 선형 어텐션이 되고, 거기에 위치별 지수 감쇠를 주면 곧 RNN 상태 업데이트가 된다 — 즉 “어텐션과 RNN은 같은 식의 두 형태”라는 관점을 14B 스케일로 실증.
연결 고리: Attention Free Transformer(AFT, Zhai 2021)의 position-decay bias를 학습 가능 파라미터와 블록 구조로 일반화.
시사점: 추론 비용 지배의 서비스 환경에서 “Transformer 1등만 바라볼 필요 없다”는 실천적 근거.
비판적 코멘트: 동일 연산 예산에서 소프트맥스 Transformer의 in-context learning 능력을 완전 대체하는지는 더 큰 규모/reasoning에서 추가 검증 필요.

Discussion Points

논쟁점: 고정된 exponential decay가 표현력 병목인가, 아니면 단지 스케일/데이터 부족 문제인가?
검증 필요 가정: $exp (w)$ 감쇠만으로 복잡한 구문/코드 의존성을 포착할 수 있는가.
후속 연구: RetNet, Mamba(SSM), Griffin — 입력 의존 게이팅/상태 확장으로의 자연스러운 일반화.

메타데이터

항목	내용
제목	RWKV: Reinventing RNNs for the Transformer Era
저자	Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, 등 (31명)
소속	EleutherAI / RWKV Foundation / 다수 기관 커뮤니티 협업
연도	2023
발표	arXiv:2305.13048 (EMNLP 2023 Findings)
링크	arXiv, GitHub
키워드	RNN, linear attention, efficient LLM, receptance, time-decay

왜 이 연구를 하는가?

핵심 질문

“Transformer의 병렬 학습 가능성과 RNN의 상수 메모리 추론을 하나의 아키텍처에서 동시에 달성할 수 있는가?”

기존 접근법의 한계

한계	설명
Transformer 추론 비용	토큰당 O(T·d), KV 캐시 O(T·d)로 긴 문맥에서 비용 폭증
전통 RNN 학습 비용	시간축이 순차적이라 BPTT가 병렬화 어려움 → 큰 스케일에서 수렴 느림
선형 어텐션(Linear Attention)	수학적으로 RNN 등가지만 실제 대형 LLM 규모에서 성능이 밀림
AFT(Attention Free Transformer)	위치-의존 bias만 사용 → 성능 및 장기 의존성 제한

핵심 통찰

소프트맥스 어텐션을 “지수 가중 합”으로 재해석 → 채널별 시간감쇠 $w_{c}$ 와 토큰별 보너스 $u$ 로 분리하면 recurrence와 parallel prefix sum 두 방식 모두로 계산 가능.
과거 정보 누적 외에 receptance σ(R) 로 채널별 “읽기 게이트”를 두면 단순 지수 감쇠의 경직성을 완화.
따라서 LM은 학습에선 Transformer-like block(병렬), 배포에선 RNN state(스트리밍) 로 같은 가중치를 양쪽으로 사용.

방법 (Method)

프레임워크 개요

flowchart TB
    subgraph Training["Training Mode — Time-Parallel"]
      X1["x_1..x_T (batch)"] --> TS1["Token Shift μ"]
      TS1 --> RKV1["R, K, V projections"]
      RKV1 --> WKV_P["WKV parallel<br/>(prefix-sum of<br/>exp(w,u)·V and exp·1)"]
      WKV_P --> Gate1["σ(R) ⊙ WKV"]
      Gate1 --> O1["output o_1..o_T"]
    end

    subgraph Inference["Inference Mode — RNN Recurrent"]
      Xt["x_t (streaming)"] --> TS2["Token Shift<br/>(remember x_{t-1})"]
      TS2 --> RKV2["R_t, K_t, V_t"]
      RKV2 --> State["state (a_t, b_t)<br/>O(d) memory"]
      State -->|exp(-w) decay + new k,v| State2["update (a_{t+1}, b_{t+1})"]
      State --> Gate2["σ(R_t) ⊙ (a_t/b_t + u·bonus)"]
      Gate2 --> Ot["output o_t"]
    end

    Training -. 동일 가중치 .- Inference

핵심 구성요소

(1) Token Shift

x_{t}^{shift} = μ ⊙ x_{t} + (1 - μ) ⊙ x_{t - 1}

→ 폭 2의 1-D depthwise conv와 동치. R/K/V 각각에 대해 별도 $μ_{R}, μ_{K}, μ_{V}$ 학습.

(2) Time-mixing (WKV)

채널별 학습 파라미터 시간감쇠 $w \in R^{d}$ (보통 $- exp (w)$ 로 양수 감쇠 부호 고정)와 현재 토큰 보너스 $u$ 로

WKV_{t} = \frac{\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} )}

출력:

o_{t} = W_{o} (σ (r_{t}) ⊙ WKV_{t})

재귀 형태 (추론): $a_{t} = e^{- w} a_{t - 1} + e^{k_{t}} v_{t}$ , $b_{t} = e^{- w} b_{t - 1} + e^{k_{t}}$ . 현재 보너스 $u$ 를 반영해 $WKV_{t} = (a_{t - 1} + e^{u + k_{t}} v_{t}) / (b_{t - 1} + e^{u + k_{t}})$ .
병렬 형태 (학습): 분자/분모 각각 지수 감쇠 prefix-sum → GPU 친화적.

(3) Channel-mixing

r_{t} = W_{r} x_{t}^{shift}, k_{t} = W_{k} x_{t}^{shift}

o_{t}^{ch} = σ (r_{t}) ⊙ (W_{v} \cdot max (0, k_{t})^{2})

→ ReLU² 활성 + receptance 게이트의 GLU-변형 MLP. 위치 독립.

(4) Receptance Gating

$σ (r_{t})$ = 각 채널이 이번 스텝에 과거 집계를 “얼마나 받을지” 결정. LSTM output gate와 유사하나 단일 히든 상태 + 채널별 고정 감쇠라는 점이 다름.

발견 (Findings)

주요 결과

14B까지 안정적 학습, loss curve가 Pythia와 교차 없이 유사 수준 유지.
Zero-shot 상식/LM 벤치마크에서 동급 Transformer와 평균 ±1% 내외.
컨텍스트 길이 증가 시 Transformer는 KV 캐시 선형 증가, RWKV는 상수 — 추론 GPU 메모리 수십 배 이득.

핵심 발견

스케일링 법칙 유사성: RWKV도 Chinchilla-스타일 loss vs compute 직선 관계를 따른다.
긴 컨텍스트 생성 시 실측 지연시간 1/10 이하 (논문 Fig. 7 참고).
고정 감쇠의 표현력 한계: 정밀 “needle-in-haystack” 스타일 recall은 약한 경향 → 후속 RetNet/Mamba의 입력 의존 상태 동기 부여.

이론적 의의

RNN ↔ 선형 어텐션의 통일

소프트맥스를 지수 가중합으로 재해석 + 채널별 감쇠를 도입 = 선형 어텐션 = 지수 감쇠 RNN. RWKV는 이 등가성을 실용 LLM 규모로 실증했다.

학습-추론 이중성(Duality)을 아키텍처 원리로

“같은 수식을 두 계산 그래프로 풀 수 있도록 설계한다”는 원칙은 이후 RetNet의 retention, Mamba의 selective SSM의 parallel scan 등에 계승된다.

재현성 및 신뢰도 평가

항목	등급	비고
코드 공개	✅	BlinkDL/RWKV-LM, HuggingFace, CUDA kernel 공개
데이터 공개	✅	Pile / RedPajama 계열 공개 데이터 사용
하이퍼파라미터	✅	Appendix에 레이어/차원/lr 등 기재
실험 환경	⚠️	GPU 클러스터 구성 일부만 기재
통계적 신뢰도	⚠️	다중 시드 표준편차 보고 부족, 단일 실행 기반 수치
종합 등급	A	체크포인트/커널/레시피 모두 공개된 커뮤니티 재현성 우수

주장별 신뢰도

#	주장	근거	신뢰도
1	14B까지 Transformer와 대등한 zero-shot 성능	Table 4 벤치마크 비교	🟢
2	추론 비용 O(T·d), 메모리 O(d)	설계상 자명 + 벤치 Fig. 7	🟢
3	긴 의존성 recall이 Transformer 수준	태스크가 일반 LM 벤치 중심, needle-recall 미포함	🟡
4	프롬프트 민감성은 “경미한” 문제	정량 측정 부족	🟡

읽기 난이도: ⭐⭐⭐

선형 어텐션 / AFT / RNN BPTT / prefix-sum 병렬화에 대한 선지식이 필요.

축	RWKV	Linear Attention (Katharopoulos 2020)	RetNet (Sun 2023)	Mamba (Gu&Dao 2023)
핵심 접근	채널별 exp-decay + receptance 게이트 RNN	커널 φ(Q)φ(K)ᵀV 재귀	Multi-scale retention (복소수 감쇠 행렬)	Selective SSM, 입력 의존 A,B
문제 정의	Transformer급 LM, 선형 추론	선형 복잡도 어텐션	LM + O(1) 추론	긴 시퀀스 modeling
데이터	Pile-like, 332B tok, 14B 파라미터	이미지/짧은 NLP 위주	책/웹 텍스트, 최대 6.7B	Pile, 최대 2.8B
핵심 메트릭	zero-shot LM acc ≒ Pythia	품질 저하 크고 소규모	LM ppl Transformer와 동급	SSM이 Transformer 상회 보고
확장성	14B 실증	소규모 한정	6.7B 실증	수 B 실증
한계	고정 감쇠, 정밀 recall 취약	품질 열세	복소수 파라미터 튜닝	커널 최적화 필요
코드 공개	✅	✅	✅	✅

원자적 인사이트 (Zettelkasten)

💡 학습-추론 이중 형식은 아키텍처 설계의 1차 원리다

출처: RWKV - Reinventing RNNs for the Transformer Era (Peng et al., 2023)
유형: 방법론적 / 이론적

동일한 연산을 “시간-병렬 prefix-sum”과 “상수 메모리 재귀”로 양쪽 모두 계산 가능하도록 수식을 설계하면, 학습 처리량과 배포 비용을 동시에 최적화할 수 있다. RWKV의 WKV는 이를 지수 감쇠 분자/분모로 구현한 최초의 대형 LLM급 사례다.

핵심 조건/맥락: 재귀의 누적 연산자가 결합 법칙(associativity)을 가질 때 — exp 감쇠 + 합은 이 조건을 만족.
연결: Mamba - Linear-Time Sequence Modeling with Selective State Spaces, RetNet parallel scan.
활용 가능성: 새 시퀀스 모델 설계 시 “내 연산이 parallel scan으로 풀리는가?”를 디자인 체크리스트로.

💡 소프트맥스 어텐션 = 지수 가중 RNN (채널별 감쇠를 도입하면)

출처: RWKV - Reinventing RNNs for the Transformer Era (Peng et al., 2023)
유형: 이론적

소프트맥스의 exp(·)을 “시간 축 지수 감쇠 + 키 보너스”로 재배치하면 분모/분자가 각각 RNN 상태로 누적된다. 어텐션과 RNN은 다른 구조가 아니라 같은 식의 다른 정규화 방식일 뿐이라는 관점을 제공한다.

핵심 조건/맥락: 감쇠가 데이터 독립일 때. 데이터 의존 감쇠로 가면 Mamba/Selective SSM으로 이행.
연결: Linear Attention kernel trick, RetNet retention.
활용 가능성: 어텐션 변형을 비교할 때 “감쇠가 어디서, 얼마나 data-dependent한가”를 축으로 분류.

💡 Receptance 게이트는 “과거를 얼마나 믿을지”를 채널별로 학습시킨다

출처: RWKV - Reinventing RNNs for the Transformer Era (Peng et al., 2023)
유형: 방법론적

고정 시간감쇠만 있으면 경직된 망각 곡선을 갖지만, $σ (R_{t})$ 를 곱해 주면 토큰/채널별로 과거를 읽을지 말지 동적 게이팅이 가능해진다. LSTM output gate의 단순 재해석이지만, 선형 어텐션식 집계 위에 얹었다는 점이 새롭다.

핵심 조건/맥락: R 프로젝션이 현재 입력 기반이어야 효과적 — token-shift가 받쳐 줌.
연결: LSTM output gate, GLU, Selective SSM의 입력 의존 gating.
활용 가능성: 선형 어텐션에 안정성이 부족할 때 receptance-style 게이트 추가가 싼 개선책.

💡 고정 exp 감쇠의 한계 = 정밀 recall의 구조적 병목

출처: RWKV - Reinventing RNNs for the Transformer Era (Peng et al., 2023)
유형: 실패/한계

과거가 단일 벡터 상태 $(a_{t}, b_{t})$ 에 지수 속도로 압축되면, “t-1000 스텝 전 특정 토큰” 같은 정밀 look-up은 수학적으로 손실된다. 이는 이후 Mamba의 selective mechanism, Based/Hybrid 모델의 등장 동기가 된다.

핵심 조건/맥락: 상태 차원 $d$ 와 시퀀스 길이 $T$ 의 비율이 충분히 작을 때만 근사 recall 가능.
연결: Transformer KV 캐시의 “lossless” recall과 대비, Mamba selective SSM.
활용 가능성: Long-context benchmark에서 RNN류 vs attention류의 trade-off를 설명하는 1차 근거.

핵심 용어 정리

용어	정의
RWKV	Receptance-Weighted Key-Value. 네 가지 학습 벡터(R, W, K, V) 중심의 선형-어텐션형 RNN.
Time-mixing block	시간 축으로 과거 토큰을 지수 감쇠 누적하는 블록 (어텐션 역할).
Channel-mixing block	위치 독립 GLU-변형 MLP (FFN 역할).
WKV	Time-mixing의 핵심 연산. 지수 가중치 $w$ + 현재 보너스 $u$ + key $k$ 로 value $v$ 를 누적.
Receptance (R)	σ(R)로 게이팅 — “과거 집계를 얼마나 받아들일지”를 채널별로 결정.
Token shift	$μ x_{t} + (1 - μ) x_{t - 1}$ 선형 결합. 얕은 1-D conv와 동치.
Time-parallel mode	학습 시 시퀀스 전체를 prefix-sum으로 병렬 처리.
Recurrent mode	추론 시 O(d) 상태 $(a_{t}, b_{t})$ 로 토큰당 상수 연산.
Linear Attention	softmax를 커널 분해로 근사해 O(Td²) 복잡도로 만든 어텐션 일가.
AFT	Attention Free Transformer. 학습 가능 위치 bias로 어텐션을 대체한 RWKV의 직접 전신.

RWKV - Reinventing RNNs for the Transformer Era