기존 대규모 언어 모델의 제약이었던 「입력량의 한계」를 제거한 「RWKV」란?

 

ChatGPT나 Bard 등, 상용으로 이용되고 있는 대규모 언어 모델은 거의 모두가 트랜스포머라는 아키텍처를 이용한 모델입니다만, 트랜스포머형 모델은 입력 길이의 제곱에 비례하여 계산량이 증가하기 때문에 입력 크기가 제한된다는 문제가 있습니다. 이러한 문제에 대응해, 대용량 데이터에의 대응이나 추론 시의 메모리 사용량 절감을 통하여 트랜스포머에 필적하는 성능을 낼 수 있는 아키텍처 「RWKV」에 대하여 저자 중 한 사람이 해설하였습니다.

The RWKV language model: An RNN with the advantages of a transformer | The Good Minima
https://johanwind.github.io/2023/03/23/rwkv_overview.html

The RWKV language model: An RNN with the advantages of a transformer

How the RWKV language model works | The Good Minima
https://johanwind.github.io/2023/03/23/rwkv_details.html


RWKV는 RNN(회귀형 신경망)이 베이스가 되고 있습니다. RNN은 입력된 토큰을 하나씩 순서대로 처리하여 상태 벡터를 업데이트하는 메커니즘입니다.

텍스트를 토큰화하고 순차적으로 읽어 들여 단일 상태 벡터로 만들거나 최종 결과를 바탕으로 다음 토큰을 예측하여 텍스트 생성 태스크를 처리할 수 있습니다. 텍스트가 토큰화되는 모습에 대해서는 아래를 참고하세요.

ChatGPT가 어떻게 문장을 인식하는지 한눈에 알 수 있는 'Tokenizer'

ChatGPT가 어떻게 문장을 인식하는지 한눈에 알 수 있는 'Tokenizer'

RNN은 "입력을 하나씩 순서대로 처리한다"는 성질을 갖기 때문에 GPU 등을 이용하여 대규모로 병렬화하는 것이 어려웠습니다. 그래서 입력을 순차적으로 처리하는 RNN 대신 "어텐션"이라는 구조를 이용하여 모든 입력 토큰을 동시에 처리하는 "트랜스포머"가 만들어졌습니다. 계산을 대규모로 병렬화할 수 있게 되었기 때문에, 대형 모델과 데이터 세트에서도 고속으로 트레이닝할 수 있어, ChatGPT 등 다수의 뛰어난 대규모 언어 모델을 만들어 냈습니다.

그러나 어텐션의 계산에 있어서, 입력 토큰 간의 모든 쌍에 대해서 계산할 필요가 있기 때문에, 처리에 필요한 시간이 입력 토큰 양의 제곱에 비례해 버린다는 것 외에, 텍스트를 생성할 때에는 모든 토큰의 어텐션 벡터를 이용하기 때문에 메모리를 대량으로 필요로 하는 등, 입력 토큰 길이에 제한이 있었습니다. 반면에 RNN의 경우 계산량이 입력 토큰 양의 1승에 비례하기 때문에 상당히 긴 문장을 읽을 수 있습니다.

RWKV는 트랜스포머처럼 입력 토큰을 동시에 처리하는 것으로 병렬화를 가능케 하면서, RNN과 같이 방대한 입력을 실시해도 고속 계산이 가능해지는 「양쪽의 좋은 것만 택한 아키텍처」라는 것으로, 아래 그림과 같은 구조를 가지고 있으며 "Time Mixing" 블록과 "Channel Mixing" 블록으로 나뉘는 것이 특징입니다.

이 아키텍처의 「RWKV」라는 이름은 과거 정보의 수용도를 나타내는 R 벡터나 시간 계수를 나타내는 W 벡터, 그리고 어텐션에서 사용되는 「Key」, 「Value」와 같은 벡터가 이용되고 있다는 것에서 각각 이니셜을 취한 것입니다.

RWKV는 구조상 레이어를 분리하여 RNN의 이전 노드 계산이 끝나기 전에 다음 노드 계산을 시작할 수 있습니다. 아래 그림과 같이 먼저 좌측 열의 「My → name」의 Time Mixing 블록의 계산을 한 단계에서, 가운데의 「name → is」 열의 Time Mixing 블록의 계산을 시작할 수 있습니다. 이렇게 RNN임에도 불구하고 대규모 병렬화를 가능하게 하고 있다는 것입니다. 덧붙여 이렇게 병렬화를 실시하는 것은 학습 시 뿐이며, 추론 시에는 순차 계산을 하는 것으로 메모리의 소비량을 대폭 삭감할 수 있다고 합니다.

LAMBADA 나 PIQA 등을 이용한 성능 평가에서는, RWKV는 작은 파라미터수부터 큰 파라미터수까지 트랜스포머형 모델에 필적하는 성능을 낼 수 있었다고 합니다.

RWKV의 구현에 대해서는, RWKV 논문의 저자 중 한명인 조한 윈드 씨가 약 100행의 최소 구현을 해설과 함께 공개하고 있습니다.

How the RWKV language model works
https://johanwind.github.io/2023/03/23/rwkv_details.html