기술 요약: 멀티믹서 모델 (Oryx)

문제 정의

소프트맥스 어텐션은 현대의 대규모 언어 모델 (LLM) 에서 지배적인 메커니즘이지만, 완전한 키 - 값 (KV) 캐시를 유지해야 하므로 시퀀스 길이에 대해 이차적인 계산 비용과 선형적인 메모리 요구 사항을 초래합니다. 반면 선형 재귀 모델 (예: 선형 어텐션 및 Mamba-2 와 같은 상태 공간 모델) 은 선형 계산과 상수 메모리를 제공하지만, 강력한 검색이나 인-컨텍스트 학습이 필요한 작업에서는 어텐션 기반 모델보다 성능이 뒤처지는 경우가 많습니다. 기존 하이브리드 아키텍처는 계층 간 (inter-layer) 에 어텐션과 재귀 블록을 정적으로 교차 배치하거나 단일 계층 내 (intra-layer) 에서 융합함으로써 이러한 트레이드오프를 완화하려 시도합니다. 그러나 이러한 접근 방식은 토큰 또는 계층당 계산 경로를 고정하여 시퀀스 내 서로 다른 세그먼트의 특정 요구 사항에 맞춰 모델링 메커니즘을 적응시키지 못합니다.

방법론: Oryx 아키텍처

저자들은 단일 토큰 시퀀스 전반에 걸쳐 다양한 시퀀스 믹서 (예: 이차적 소프트맥스 어텐션과 선형 재귀 메커니즘) 간 유연하게 전환할 수 있는 시퀀스 축 하이브리드 아키텍처인 Oryx를 제안합니다.

핵심 설계 원칙

1. 결합 프로젝션을 통한 공유 표현:

   • Oryx 는 어텐션과 선형 모델을 키 - 값 (KV) 연관성을 유지하는 연관 기억 (associative memory) 관점 하에 통합합니다.
   • 아키텍처는 믹서 모드 간에 키 ($W_K$) 와 값 ($W_V$) 에 대해 **결합 프로젝션 (tied projections)**을 사용합니다. 이를 통해 단일 순전파 (forward pass) 로 공유 표현을 사용하여 어텐션 KV 캐시와 선형 재귀 상태를 동시에 업데이트할 수 있습니다.
   • 믹서별 쿼리: 키와 값과 달리 쿼리 프로젝션 ($W_Q$) 은 각 믹서마다 고유하게 유지됩니다. 실험 결과에 따르면 쿼리 가중치를 공유하면 성능이 저하되는데, 이는 어텐션과 재귀 메커니즘이 공유 상태로부터 모드별 정보를 추출하기 위해 서로 다른 읽기 벡터 (readout vectors) 가 필요함을 시사합니다.

2. 아키텍처 구성 요소:

   • 블록은 공유 키와 값에 적용되는 선형 모델 성능에 중요한 단단한 컨볼루션과 곱셈 게이트와 같은 구성 요소를 포함합니다.
   • 호환성을 보장하기 위해 Oryx 는 Mamba-2 와 같은 선형 모델이 일반적으로 다중 값 (MVA) 구조를 사용함에도 불구하고 두 믹서 모두에 대해 통일된 헤드 구조 (Multi-Head Attention, MHA) 를 강제합니다.
   • 블록은 모든 타임스텝에서 KV 캐시와 선형 재귀 상태를 모두 유지하며, 이를 공동으로 업데이트합니다.

3. 조각화된 혼합 모드 학습:

   • 추론 중 강력한 모드 전환을 가능하게 하기 위해 Oryx 는 조각화된 혼합 모드 전략으로 학습됩니다. 학습 시퀀스는 고정된 길이의 조각 (예: 128 토큰) 으로 분할되며, 각 조각은 무작위로 믹서 모드 (어텐션 또는 선형) 를 할당받습니다.
   • 이 학습 체계는 모델이 모드 간 원활한 전환을 가능하게 하는 호환 가능한 표현을 학습하도록 강제하여, 시퀀스 수준 혼합 학습에서 선형에서 어텐션 메커니즘으로 전환할 때 흔히 발생하는 성능 저하를 방지합니다.

주요 기여

• 시퀀스 축 하이브리드화: Oryx 는 새로운 하이브리드화 축을 도입하여 모델이 고정된 계층 또는 블록 아키텍처에 제한되는 대신 토큰 시퀀스 전반에 걸쳐 메커니즘을 동적으로 전환할 수 있게 합니다.
• 파라미터 효율성: 이 설계는 믹서 모드 간에 최소 90% 의 파라미터 (특히 $W_K, W_V, W_G, W_O$) 를 결합하여, 어텐션 및 재귀 모드가 상당한 파라미터 오버헤드 없이 공유 내부 표현 위에서 작동할 수 있게 합니다.
• 유연한 추론: 이 모델은 시퀀스의 서로 다른 세그먼트 (예: 컨텍스트 프리필 vs 생성) 를 서로 다른 믹서가 처리할 수 있는 추론을 지원합니다. 예를 들어, 모델은 효율적인 컨텍스트 처리를 위해 선형 재귀를 사용하고, 검색이 많은 생성 단계에서는 어텐션으로 전환할 수 있습니다.

실험 결과

실험은 130M 에서 1.4B 파라미터까지의 모델에서 수행되었으며, 선형 믹서로 Mamba-2 와 게이트드 델타넷 (GDN) 을 사용했습니다.

• 언어 모델링: 고정된 토큰 예산 하에서 Oryx 모델 (Oryx-TM 및 Oryx-TG 변형 모두) 은 단일 믹서 베이스라인 (순수 Transformer, 순수 Mamba-2, 순수 GDN) 과 비교해 동등하거나 더 나은 성능을 달성합니다. 1.4B 규모에서 Oryx 변형은 평균 언어 모델링 작업에서 각각의 베이스라인보다 최소 0.7 퍼센트 포인트만큼 더 나은 성능을 보입니다.
• 검색 능력:
  • Oryx 는 구성 믹서의 검색 능력을 유지합니다.
  • 특히, 혼합 추론 모드 (선형 재귀로 컨텍스트 처리, 어텐션으로 생성) 를 사용할 때, Oryx 는 10% 미만의 토큰을 어텐션 모드로 처리하면서도 Transformer 베이스라인과 비교해 동등한 성능을 달성합니다.
  • 실제 세계 검색 작업에서 혼합 모드 Oryx 는 Oryx-TG 의 경우 8.6 퍼센트 포인트, Oryx-TM 의 경우 13.5 퍼센트 포인트만큼 선형 베이스라인을 능가합니다.
  • 합성 "Haystack 속 바늘 (Needle-in-a-Haystack, NIAH)" 테스트에서 Oryx 는 선형 베이스라인보다 38.6 에서 40.3 퍼센트 포인트만큼 더 나은 성능을 보여주어, 선형 믹서가 시퀀스의 대부분을 처리하더라도 공유 표현이 긴 컨텍스트 검색을 효과적으로 지원함을 입증합니다.
• 모드 전환: 퍼플렉시티 분석은 믹서 간 전환 (예: 어텐션에서 Mamba-2 로 또는 그 반대로) 후 모델의 퍼플렉시티가 선택된 모드의 베이스라인으로 빠르게 수렴함을 보여줍니다. 이는 공유 상태가 호환 가능하며 전환이 원활함을 의미합니다.

중요성 및 주장

이 논문은 Oryx 가 서로 다른 방법론에도 불구하고 어텐션과 선형 재귀 메커니즘이 내부 표현을 공유할 수 있음을 입증한다고 주장합니다. 이 발견은 이러한 메커니즘이 완전히 별도의 상태 공간을 필요로 한다는 가정에 도전합니다.

이 연구의 중요성은 다음과 같습니다:

1. 시퀀스 축 하이브리드화 검증: 시퀀스 내에서 믹서를 동적으로 전환하는 것이 적응형 계산을 위한 새로운 설계 공간을 제공하는 실현 가능하고 효과적인 전략임을 확립합니다.
2. 검색 격차 해소: 선형 모델이 전체 시퀀스에 대한 완전한 어텐션 메커니즘 없이도 특정 세그먼트에 대해 어텐션 모드를 활용할 수 있다면 Transformer 와 비교해 동등한 검색 성능을 달성할 수 있음을 보여줍니다.
3. 미래 방향: 저자들은 이 아키텍처가 학습된 라우팅 (예: FLOP 절약을 기반으로 RL 을 사용하여 믹서 선택) 과 어텐션 모드가 검증하기 위해 선형 모드가 추론 흔적을 초안으로 작성하는 추측적 디코딩과 같은 기회들을 열어준다고 제안합니다.

저자들은 KV 캐시와 선형 상태를 유지하는 것이 메모리 오버헤드를 초래한다고 겸손하게 지적하며, 이는 긴 컨텍스트 길이에서는 KV 캐시가 주된 요인이라고 덧붙입니다. 또한 멀티믹서 모델의 학습 역동성은 아직 크게 탐구되지 않았으며 규모가 커짐에 따라 특수한 목적 함수가 필요할 수 있다고 언급합니다.