기술 요약: RADLADS

문제 정의

선형 어텐션 변형 모델들(RWKV 및 Mamba와 같은)은 전통적인 소프트맥스 어텐션 트랜스포머에 비해 상당한 이점(구체적으로 KV 캐시를 피함으로써 얻는 $O(1)$의 추론 시간 및 일정한 메모리 사용량)을 제공하지만, 대규모 선형 모델을 처음부터 학습시키는 것은 여전히 매우 비용이 많이 듭니다. 최첨단(SoTA) 트랜스포머 모델들은 종종 10조 개 이상의 토큰으로 학습되어야 하며, 이는 대부분의 조직이 감당하기 어려운 비용입니다. 사전 학습된 소프트맥스 트랜스포머를 선형 또는 순환 모델로 변환하는 기존 방법들(예: T2R, SUPRA, LoLCats, MOHAWK)은 역사적으로 방대한 양의 토큰(200억 개에서 1,000억 개 이상)을 필요로 했거나, MMLU와 같은 벤치마크에서 낮은 성능을 보였습니다. 또한, 많은 변환 시도들이 하이브리드 아키텍처(일부 소프트맥스 어텐션을 유지함)에 의존하거나 원래의 티처(teacher) 모델의 품질을 따라잡는 데 실패했습니다.

방법론

저자들은 소프트맥스 어텐션 트랜스포머를 최소한의 계산 비용으로 선형 어텐션 디코더 모델로 변환하기 위해 설계된 프로토콜인 RADLADS(Rapid Attention Distillation to Linear Attention Decoders at Scale)를 제시합니다. 이 과정은 세 가지 주요 단계와 특정 아키텍처 적응 과정을 포함합니다.

1. RADLADS 프로토콜

변환은 단 3.5억~7억 개의 토큰(티처의 사전 학습 데이터의 0.005% 미만)만을 필요로 하는 3단계 프로세스입니다:

• 설정 (어텐션 가중치 전송): 티처 모델의 어텐션 관련 가중치($W_q, W_k, W_v, W_o$)를 학생 모델의 시퀀스 믹싱 레이어로 직접 전송합니다. 다른 가중치들은 표준 사전 학습 방법을 통해 초기화되거나, 즉각적인 효과 없이 티처의 동작을 모방하도록 설정됩니다.
• 단계 1: 어텐션 은닉 상태 정렬 (Attention Hidden State Alignment): 학생 모델의 시퀀스 믹싱 레이어가 대응하는 티처 어텐션 레이어의 은닉 상태 출력을 근사하도록 학습됩니다. 이는 L2 거리(또는 MSE) 목적 함수를 사용하여 수행됩니다. 저자들은 Gated Linear Attention 커널을 사용하는 것(오프-바이-원 디케이 및 보너스 항 제거)이 표준 RWKV-6보다 티처의 은닉 상태에 더 밀접하게 부합한다는 것을 발견했습니다.
• 단계 2: 지식 증류 (Knowledge Distillation): 전체 학생 모델이 KL 발산 손실을 사용하여 티처의 출력 로짓(logits)을 근사하도록 학습됩니다. 저자들은 사실적 지식이 주로 티처의 MLP와 임베딩에 존재한다고 가설을 세웠으며, 따라서 이 구성 요소들의 학습률을 낮게 유지하거나 고정하여 파괴적 망각을 방지하는 반면, 시퀀스 믹서는 더 공격적으로 학습시킵니다.
• 단계 3: 컨텍스트 길이 확장 (선택 사항): 긴 컨텍스트 능력을 향상시키기 위해 티처 모델 없이 표준 교차 엔트로피 손실을 사용하여 더 긴 시퀀스(최대 16k 토큰)로 미세 조정합니다. 또는, 저자들은 증류 단계 자체에서 시퀀스 길이를 4096으로 확장하는 단계 2a를 제안하며, 이는 별도의 단계 3의 필요성을 없앨 수 있습니다.

2. 새로운 아키텍처들

저자들은 표준 RWKV 아키텍처가 항상 변환에 최적은 아니라는 점을 확인했습니다. 그들은 두 가지 새로운 변형을 도입했습니다:

• RAD-RWKV6 ("RADFinch"): 더 나은 안정성과 적합성을 위해 Gated Linear Attention 커널과 상태 균형(state balancing) 기술을 사용한 RWKV-6의 수정 버전입니다.
• RAD-RWKV7 ("RADGoose"): (이 맥락에서 이점이 없었던) "토큰시프트(tokenshift)" 메커니즘을 제거하고 RoPE(Rotary Positional Embeddings)를 직접 적용한 RWKV-7의 수정 버전입니다. 이 아키텍처는 수정되지 않은 RWKV-7에 비해 더 빠른 수렴 속도와 더 낮은 증류 손실을 보여주었습니다.

3. 하이퍼파라미터 및 데이터

• 데이터셋: 저자들은 Qwen 모델을 변환하는 데 있어 FineWeb이나 FineWeb-Edu보다 DCLM(DataComp-LM)이 우수하다는 것을 발견하여 모든 변환 단계에 이를 사용하기로 결정했습니다.
• 학습률 (Learning Rates): 단계 1에서는 은닉 상태를 정렬하기 위해 높은 값($10^{-3}$)에서 시작하여 티처의 최종 사전 학습 학습률($10^{-5}$) 근처에서 끝나는 코사인 어닐링 스케줄을 사용합니다. 단계 2와 3은 고정된 학습률을 사용합니다.

주요 기여

1. RADLADS 증류 레시피: 최소한의 데이터로 고품질 변환을 가능하게 하는 특정 하이퍼파라미터, 토큰 수, 데이터셋 선택을 포함한 상세한 단계별 프로토콜을 제공합니다.
2. 새로운 아키텍처: 수정되지 않은 이전 모델들보다 더 빠른 추론과 더 나은 티처 모델과의 정렬을 제공하며, 변환 과정에 최적화된 RAD-RWKV6 및 RAD-RWKW7을 도입했습니다.
3. 대규모 모델: 7B, 32B, 72B 파라미터 규모의 인기 있는 오픈 소스 Qwen2.5 모델을 선형 어텐션 변형 모델로 성공적으로 변환했습니다.
4. 오픈 소스 공개: 다른 이들이 프로세스를 재현할 수 있도록 코드와 변환된 모델(QRWKV6/7-7B/32B/72B)을 Apache 2.0 라이선스(72B 모델의 경우 Qwen 라이선스 제한 적용) 하에 공개했습니다.

결과

변환된 모델들은 해당 크기의 순수 순환 모델 중에서 최첨단 성능을 달립니다:

• 효율성: 72B 모델을 변환하는 데 드는 비용은 2,000달러 미만이며 약 7억 개의 토큰만을 필요로 합니다.
• 성능: 표준 벤치마크(Lambada, MMLU, ARC 등)에서 RADLADS 모델은 다른 변환 방법(예: SUPRA, LoLCats, MOHAWK, ARWKV)을 지속적으로 능가합니다.
  • QRWKV7-7B-Instruct는 티처 대비 **92.4%**의 상대적 MMLU 점수를 달성하여, 다른 7B 변환 모델들을 크게 앞질렀습니다.
  • QRWKV6-72B-Instruct는 **89.9%**의 상대적 MMLU 점수를 달성하여, 해당 규모의 순수 RNN 언어 모델 중 새로운 SoTA를 기록했습니다.
• 추론 속도: 선형 어텐션 메커니즘 덕분에, 변환된 모델은 컨텍스트 길이가 증가함에 따라 티처 모델보다 상당한 속도 향상을 보여줍니다. 예를 들어, 8k 입력 및 256 출력 시, 32B QRWKV7 모델은 Qwen3-32B 티처보다 1.61배 빠르며, 6k 입력 및 2k 출력 시에는 속도가 3.41배에 달합니다.

한계 및 주장

논문은 다음과 같은 몇 가지 한계를 겸허히 인정합니다:

• 추론 및 긴 컨텍스트: 일반적인 벤치마크에서는 강력한 성능을 보이지만, 복잡한 추론 작업(예: Minerva Math) 및 매우 긴 컨텍스트 작업(예: RULER)에서는 티처 모델만큼 추가 출력 토큰에 따라 성능이 효과적으로 개선되지 않는 한계를 보입니다.
• 아키텍처 민감도: 각 새로운 아키텍처 설계는 RADLADS 프로토콜과의 호환성을 보장하기 위해 세심한 테스트를 필요로 합니다. 예를 들어, 일부 선형 어텐션 변형에서 흔히 쓰이는 GroupNorm/LayerNorm은 14B 이상의 규모에서 학습 불안정을 초래하여 pre-scaling 기술로 대체되었습니다.
• 데이터셋 정렬: 저자들은 반복적인 루핑(looping) 동작을 방지하기 위해 추론 데이터셋이 티처 모델의 분포와 더 밀접하게 정렬되어야 할 수도 있다고 언급했으며, 이는 DCLM과 OpenThoughts를 혼합함으로써 부분적으로 해결했습니다.

의의

이 논문은 RADLADS가 대규모 선형 어텐션 모델에 대한 접근을 민주화할 수 있는 비용 효율적인 경로를 제공한다고 주장합니다. 학습 토큰 요구량을 조 단위에서 수억 개 단위로 줄임으로써, 연구자와 소규모 조직이 극단적인 비용 없이 새로운 종류의 RNN 아키텍처를 대규모로 테스트, 학습 및 배포할 수 있게 합니다. 저자들은 이를 차세대 압축 상태 어텐션(compressive state attention) 변형 모델의 개발을 가속화하기 위한 도구로 포지셔닝합니다.