RWKV-X: 효율적인 장문 맥락 언어 모델링을 위한 새로운 아키텍처
더 길고 복잡한 시퀀스 처리에 대한 끊임없이 증가하는 요구는 대규모 언어 모델(LLM)의 경계를 넓혀왔습니다. 기존의 Transformer 기반 아키텍처는 강력하지만 시퀀스 길이에 따른 이차 복잡성으로 인해 상당한 확장 문제에 직면합니다. 이러한 제한은 특히 확장된 컨텍스트 입력을 처리할 때 분명해지며, 시퀀스의 먼 부분에서 정보를 효과적으로 캡처하고 활용하는 능력을 저해합니다. 이러한 문제에 대응하여 긴 시퀀스 처리에서 선형 복잡성을 달성하기 위한 혁신적인 접근 방식이 등장했습니다.
이러한 방법에는 선형 어텐션 모델, 상태 공간 모델 (Mamba), 선형 RNN (DeltaNet), RWKV가 포함됩니다. 이러한 각 아키텍처는 이차 복잡성 문제에 대한 고유한 솔루션을 제공하여 긴 시퀀스를 보다 효율적으로 처리할 수 있도록 합니다. 그러나 이러한 선형 아키텍처는 종종 긴 컨텍스트 정보를 완전히 이해하고 활용하는 데 어려움을 겪습니다.
예를 들어 RWKV-7 (2.9B 파라미터 모델)은 최대 28K 토큰까지의 패스키 검색 작업에서 높은 정확도를 보여줍니다. 그러나 이 임계값을 초과하면 성능이 빠르게 저하됩니다. 128K 길이의 데이터를 사용하여 지속적인 사전 학습을 수행하더라도 장문 맥락 제한이 지속됩니다. 이 문제는 RWKV에만 국한되지 않고 Mamba와 같은 다른 아키텍처로 확장되어 이 클래스 모델의 근본적인 문제를 나타냅니다. 확장된 컨텍스트에서 성능을 유지하기 위한 노력은 선형 복잡성 언어 모델에서 개선해야 할 중요한 영역을 강조합니다.
선형 복잡성 언어 모델의 현황
선형 복잡성 언어 모델은 긴 시퀀스 처리에 내재된 이차 계산 부담을 회피하여 Transformer 기반 아키텍처에 대한 매력적인 대안으로 부상했습니다. 이 영역에서 눈에 띄는 RWKV 모델 제품군은 훈련 중에 Transformer의 병렬 처리 가능성과 RNN과 유사한 순환 상태 표현을 능숙하게 결합합니다.
RWKV의 진화는 기본 RWKV-4에서 시작하여 RWKV-5, RWKV-6을 거쳐 RWKV-7에서 정점에 이르기까지 여러 반복에 걸쳐 있습니다. 각 반복은 모델의 기능을 향상시키고 제한 사항을 해결하는 개선 사항을 가져왔습니다. 또한 Jamba, Zamba 및 MiniMax와 같은 하이브리드 언어 모델은 고유한 하이브리드 디자인을 도입하여 선형 복잡성 모델의 현황을 더욱 풍부하게 만들었습니다.
효율적인 장문 맥락 처리를 위한 추구는 혁신적인 어텐션 메커니즘 개발로 이어졌습니다. 예를 들어 Native Sparse Attention은 토큰을 시간 블록으로 구성하여 세 가지의 고유한 어텐션 경로를 사용합니다. 전역 컨텍스트를 위한 압축된 거친 토큰, 로컬 세부 정보를 위해 선택적으로 유지된 미세 토큰, 로컬 컨텍스트 정보를 캡처하기 위한 슬라이딩 윈도우를 사용합니다. 다른 주목할만한 어텐션 메커니즘으로는 SeerAttention 및 블록 어텐션(MoBA)이 있으며, 각 메커니즘은 긴 시퀀스 내에서 관련 정보에 주의를 기울이기 위한 고유한 전략을 제공합니다.
RWKV-X: 향상된 장거리 컨텍스트 모델링을 위한 하이브리드 아키텍처
Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ), Shenzhen, Hohai University, Nanjing, Shenzhen University 및 Qinghai University, Xining의 연구원들은 RWKV-X라는 새로운 하이브리드 아키텍처를 도입했습니다. 이 아키텍처는 단거리 종속성을 모델링하는 RWKV의 효율성과 장거리 컨텍스트를 캡처하도록 특별히 설계된 희소 어텐션 메커니즘을 독창적으로 결합합니다.
이전의 하이브리드 접근 방식과 달리 RWKV-X는 훈련 중에 선형 시간 복잡성을 달성하고 추론 디코딩 중에 상수 시간 복잡성을 달성합니다. 따라서 긴 시퀀스를 처리하는 데 매우 효율적입니다. 이 모델은 64K 토큰 시퀀스에서 지속적으로 사전 훈련될 때 64K 패스키 검색 벤치마크에서 거의 완벽한 정확도를 보여줍니다. 짧은 컨텍스트 작업에서 강력한 성능을 유지하면서 장문 맥락 벤치마크에서 이전 RWKV-7 모델보다 지속적으로 뛰어난 성능을 보입니다.
RWKV-X의 혁신은 장문 맥락 언어 모델링의 문제를 해결하는 데 중요한 진전을 나타냅니다. 순환 모델과 희소 어텐션 메커니즘의 강점을 결합하여 RWKV-X는 효율성과 정확성 간의 균형을 이루어 확장된 시퀀스를 보다 효과적으로 처리할 수 있는 길을 열어줍니다.
RWKV-X: 아키텍처 및 훈련
RWKV-X는 RWKV-7 블록과 희소 어텐션 블록을 통합하여 두 접근 방식의 강점을 활용하는 하이브리드 아키텍처를 구현합니다. 처음부터 훈련하는 대신 RWKV-X는 LLaMA Pro에서 영감을 받은 인터리브 블록 확장 접근 방식과 제로 초기화 메커니즘을 사용하여 기존 모델을 기반으로 구축됩니다.
훈련 프로세스는 짧은 컨텍스트와 긴 컨텍스트 모두에서 모델의 성능을 최적화하도록 신중하게 설계된 두 단계로 구성됩니다.
- 단문 맥락 사전 학습: 처음에 모델은 MiniPile 데이터세트에서 추출한 짧은 1024 토큰 컨텍스트에서 훈련됩니다. 이 단계에서는 새로 추가된 블록의 파라미터를 제외한 모든 파라미터가 고정되어 기본 RWKV-7 모델에서 사전 훈련된 지식이 유지됩니다. 이를 통해 새로 추가된 블록은 사전 훈련된 표현을 방해하지 않고 기존 아키텍처에 적응할 수 있습니다.
- 장문 맥락 지속적 사전 학습: 두 번째 단계에는 ProLong-64K 데이터세트와 64K 토큰의 컨텍스트 길이를 사용하여 장문 맥락 지속적 사전 학습이 포함되며, 총 약 10억 개의 토큰을 처리합니다. 이 단계에서는 모든 파라미터가 고정 해제되고 공동으로 최적화되어 모델이 표현을 미세 조정하고 장거리 종속성을 학습할 수 있습니다. 훈련은 중요도에 따라 토큰의 가중치를 동적으로 조정하는 Long-context Cross-Entropy (LongCE) 손실을 사용합니다. 이 손실 함수는 모델이 시퀀스의 가장 관련성이 높은 부분에 집중하여 장거리 관계를 캡처하는 능력을 향상시키는 데 도움이 됩니다.
2단계 훈련 프로세스를 통해 RWKV-X는 단거리 모델링을 위한 RWKV-7의 효율성과 희소 어텐션 메커니즘의 장거리 컨텍스트 인식 기능을 효과적으로 결합할 수 있습니다. 먼저 짧은 컨텍스트에서 사전 훈련한 다음 긴 컨텍스트에서 미세 조정하여 모델은 시퀀스의 다른 부분에서 정보를 효과적으로 통합하는 방법을 학습합니다.
RWKV-X: 평가 및 성능
단문 맥락 평가는 RWKV-X가 표준 벤치마크에서 경쟁력 있는 성능을 유지하여 짧은 시퀀스를 효과적으로 처리할 수 있음을 보여줍니다. 더 작은 RWKV-X (0.22B)는 RWKV-7의 51.8과 비슷한 평균 점수 51.0을 달성합니다. 더 큰 규모에서 RWKV-X (3.6B)는 RWKV-7 (2.9B, 72.8) 및 Qwen2.5-3B (71.4)와 거의 일치하고 LLaMA3.2-3B (69.7)를 능가하면서 71.9에 도달합니다. 이러한 결과는 RWKV-X가 더 짧은 컨텍스트에서 성능을 희생하지 않고도 범용 LLM 백본으로서의 효율성을 확인시켜줍니다.
또한 효율성 분석은 긴 시퀀스에 대한 RWKV-X의 우수한 확장 특성을 보여줍니다. 128K 토큰에서 RWKV-X는 Flash-Attention v3보다 1.37배 빠른 속도 향상을 달성했으며, 컨텍스트 길이가 증가함에 따라 이러한 이점이 확장됩니다. 이는 시퀀스 길이가 증가함에 따라 RWKV-X가 다른 어텐션 메커니즘에 비해 점점 더 효율적이 된다는 것을 나타냅니다.
짧은 컨텍스트와 긴 컨텍스트 모두에서 RWKV-X의 강력한 성능은 언어 모델로서의 다재다능함과 효율성을 강조합니다. 더 짧은 시퀀스에서 경쟁력 있는 성능을 유지하면서 더 긴 시퀀스에서 상당한 속도 향상을 달성할 수 있는 능력은 광범위한 애플리케이션을 위한 유망한 아키텍처입니다.
RWKV-X: 제한 사항 및 향후 방향
RWKV-X는 단거리 종속성을 모델링하는 RWKV의 효율성과 장거리 컨텍스트 모델링을 위해 특별히 설계된 새로운 희소 어텐션 메커니즘을 성공적으로 결합한 하이브리드 언어 모델로 등장합니다. RWKV-X는 장문 맥락 언어 모델링에서 강력한 성능과 효율성을 보여주지만 몇 가지 제한 사항이 남아 있습니다.
첫째, 상위 k 청크 선택에 의존하는 희소 어텐션 메커니즘은 의미론적으로 관련된 종속성을 간과할 수 있는 휴리스틱 접근 방식을 사용합니다. 상위 k 선택 전략은 시퀀스에서 가장 중요한 정보를 항상 캡처하지 못할 수 있으며, 잠재적으로 차선의 성능으로 이어질 수 있습니다.
둘째, 현재 구현에서는 희소 어텐션 디코딩이 일반 RWKV보다 느리게 실행되어 성능을 최적화하기 위해 추가 엔지니어링 노력이 필요함을 나타냅니다. RWKV-X는 긴 시퀀스에서 다른 어텐션 메커니즘에 비해 상당한 속도 향상을 달성하지만 희소 어텐션 디코딩은 여전히 일반 RWKV보다 느리므로 구현에 개선의 여지가 있음을 시사합니다.
향후 연구에서는 보다 정교한 희소 어텐션 메커니즘을 탐색하고, 희소 어텐션 디코딩의 구현을 최적화하고, 대체 훈련 전략을 조사하여 이러한 제한 사항을 해결하는 데 집중할 수 있습니다. 이러한 문제를 극복함으로써 RWKV-X는 장문 컨텍스트 애플리케이션을 위한 더욱 강력하고 효율적인 언어 모델이 될 가능성이 있습니다.