기술 요약: Parallax - 언어 모델링을 위한 매개변수화된 지역 선형 어텐션

문제 제기

대형 언어 모델 (LLM) 은 Softmax 어텐션을 핵심 계산 원시 연산으로 의존하며, 이 메커니즘은 Transformer 도입 이후 구조적으로 변경되지 않았습니다. 선형 어텐션과 상태 공간 모델 (SSM) 과 같은 효율적인 대안들은 2 차 미만의 복잡도를 제공하지만, 컨텍스트 내 정보 검색 작업에서 Softmax 어텐션보다 일관되게 낮은 성능을 보입니다. 이러한 성능 격차는 설계상의 근본적인 트레이드오프를 시사합니다.

이론적 프레임워크, 특히 테스트 시간 회귀 프레임워크는 어텐션을 컨텍스트 내 회귀 솔버로 해석합니다. 이 프레임워크 내에서 **지역 선형 어텐션 (Local Linear Attention, LLA)**은 Softmax 어텐션에 대한 우수한 대안으로 제안되었습니다. LLA 는 Softmax 어텐션에서 사용되는 지역 상수 추정기 (Nadaraya-Watson) 를 지역 선형 추정기로 대체하여, 이론적으로 더 나은 편차 - 분산 트레이드오프와 연관 기억 용량을 제공합니다. 그러나 LLA 는 세 가지 주요 과제로 인해 LLM 사전 학습에 성공적으로 확장되지 못했습니다:

1. 계산 오버헤드: 정확한 LLA 는 매 토큰마다 병렬 켤레 기울기 (CG) 솔버를 통해 선형 시스템을 풀어야 하므로, 상당한 입출력 (I/O) 및 계산 비용이 발생합니다.
2. 수치적 불안정성: CG 솔버는 정밀도 형식에 민감하여 저정밀도 하드웨어 환경에서 관리하기 어렵습니다.
3. 정규화 트레이드오프: 실제 사전 학습 설정에서 수치적 안정성을 보장하면서도 모델을 Softmax 어텐션으로 퇴화시키지 않도록 정규화 매개변수 $\lambda$를 균형 있게 조정하는 것은 단순하지 않습니다.

방법론: Parallax

저자들은 LLM 사전 학습에 확장 가능하고 효율적이며 수치적으로 안정적인 매개변수화된 지역 선형 어텐션 버전인 Parallax를 소개합니다.

핵심 메커니즘

Parallax 는 LLA 를 Softmax 어텐션에 대한 가산 보정으로 재형식화합니다. 쿼리당 최적 보정 벡터 $\rho^*_i$를 풀기 (이는 공분산 행렬의 역행렬 계산이 필요함) 대신, Parallax 는 레이어 입력 $x_i$에서 프로브 벡터 $\rho_i = W_R x_i$로의 직접 매핑을 학습합니다.
출력은 다음과 같이 계산됩니다:
$$o^{PLX}_i = o^{SA}_i - \Sigma^{(i)}_{KV} \rho_i$$
여기서 $o^{SA}_i$는 표준 Softmax 어텐션 출력이고 $\Sigma^{(i)}_{KV}$는 KV 공분산 행렬입니다.

주요 설계 선택 사항은 다음과 같습니다:

• 솔버 제거: $\rho_i$를 매개변수화함으로써 토큰별 켤레 기울기 솔버의 필요성을 제거했습니다.
• 경계 증폭 제거: 저자들은 경계 증폭 인자 $\eta_i$를 0 으로 설정했습니다. 정확한 LLA 에서 $\eta_i$는 경계 편향을 보상하지만, 매개변수화된 설정에서는 학습 불안정성을 유발할 수 있습니다. 이를 제거하면 공분산 보정 능력을 유지하면서 순전파를 단순화할 수 있습니다.
• 통합 프레임워크: Parallax 는 대역폭, 프로브 구성, 아핀 구조로 연결된 어텐션 메커니즘 계열 내에서 위치합니다. 이는 비모수적 방법 (LLA) 과 모수적 방법 (선형 어텐션, MesaNet) 간의 격차를 연결합니다.

하드웨어 인식 알고리즘

효율성을 해결하기 위해 저자들은 FlashAttention(FA) 과 호환되는 스트리밍 알고리즘을 개발했습니다.

• 산술 강도: 이 알고리즘은 동일한 KV 스트림을 재사용하면서 계산을 약 2 배 증가시켜 FlashAttention 대비 메모리 트래픽 바이트당 산술 연산 (FLOPs) 인 산술 강도를 높입니다. 이는 연산을 현대 GPU 아키텍처에 유리한 더 계산 집약적 영역으로 이동시킵니다.
• 커널 최적화: NVIDIA Hopper GPU 를 위해 CuTeDSL 로 커스텀 디코드 커널이 구현되었습니다. 이는 두 가지 브랜치 (Softmax 및 공분산) 가 동일한 KV 스트림을 공유한다는 사실을 활용하여 동일한 텐서 코어 지시어 (WGMMA) 내에서 $QK$및$RK$ 곱의 결합 계산을 가능하게 하여 HBM 트래픽을 최소화합니다.

주요 기여

1. 아키텍처: 토큰별 솔버를 제거하고 학습 가능한 프로브를 도입함으로써 지역 선형 어텐션의 이론적 이점을 대규모 사전 학습으로 확장하는 Parallax 의 유도.
2. 효율성: 다양한 배치 크기와 컨텍스트 길이에서 FlashAttention 2/3 와 맞먹거나 능가하는 하드웨어 인식 스트리밍 알고리즘 및 커스텀 디코드 커널 개발.
3. 옵티마이저 - 아키텍처 상호작용: Parallax 아키텍처와 Muon 옵티마이저 간의 중요한 상호작용 식별. 논문은 Parallax 의 보정 브랜치가 옵티마이저의 기하학에 크게 의존함을 보여주며, Muon 하에서는 상당한 이득을 보이지만 AdamW 하에서는 Softmax 어텐션과 유사한 성능을 보임을 입증합니다.

실험 결과

저자들은 0.6B 및 1.7B 파라미터 규모에서 합성 벤치마크와 LLM 사전 학습을 통해 Parallax 를 검증했습니다.

• 합성 벤치마크: MAD-Benchmark 에서 Parallax 는 컨텍스트 내 회상 (In-Context Recall), 퍼지 회상 (Fuzzy Recall), 노이즈 회상 (Noisy Recall), 선택적 복사 (Selective Copying) 와 같은 회상 중심 작업에서 Softmax 어텐션, Mamba, Gated DeltaNet 보다 일관되게 우수한 성능을 보이며 전체 정확도에서 최고를 기록했습니다.
• 언어 모델링:
  • 퍼플렉시티: Muon 으로 학습된 Parallax 모델은 Transformer 베이스라인보다 LAMBADA 및 WikiText 에서 더 낮은 퍼플렉시티를 달성했습니다.
  • 하류 작업: Parallax 는 제로샷 QA 및 상식 추론 벤치마크에서 일관된 개선을 보여주었습니다.
  • 통제: 파라미터 매칭 (베이스라인에서 쿼리 헤드 증가) 및 계산 매칭 (Parallax 에서 헤드 차원 감소) 통제 하에서도 개선이 지속되어, 이 개선이 추가 파라미터나 계산이 아닌 메커니즘 자체에서 비롯되었음을 확인했습니다.
• 옵티마이저 의존성: 중요한 발견은 Parallax 의 이점이 옵티마이저에 달려 있다는 것입니다. Muon 하에서는 보정 - 출력 비율 (COR) 이 레이어 깊이에 따라 증가하며, 모델은 공분산 보정을 활용하는 법을 학습합니다. 반면 AdamW 하에서는 모델이 보정 브랜치를 억제 (게이트 값 $\approx 0.26$) 하여 Transformer 베이스라인과 유사한 성능을 보입니다.
• 스펙트럼 분석: 가중치 행렬에 대한 분석은 Muon 이 AdamW 에 비해 투영 가중치 (특히 $W_R$ 및 $W_R K$ 회로) 에서 더 높은 안정적 랭크를 유지함을 보여주었으며, 이는 프로브가 주요 공분산 방향과 효과적으로 정렬되도록 합니다.

중요성 및 주장

이 논문은 어텐션 메커니즘에 대한 강력한 아키텍처 - 옵티마이저 공동 설계의 첫 번째 경험적 증명을 제공한다고 주장합니다. Parallax 와 같은 고급 어텐션 메커니즘의 성공은 단순히 아키텍처 설계의 함수가 아니라 최적화 풍경과 깊이 얽혀 있음을 강조합니다.

저자들은 Parallax 를 파레토 개선으로 위치시킵니다. 올바른 옵티마이저 (Muon) 가 사용된다면 효율성을 희생하지 않으면서 Softmax 어텐션보다 더 나은 성능을 제공합니다. 또한 Parallax 는 사전 학습된 Transformer 에서 $W_R=0$으로 초기화함으로써 사전 학습 후 적응이 가능하게 하며, 이는 선형 어텐션 계열이 공유하지 않는 기능입니다.

이 연구는 미래의 어텐션 메커니즘 연구가 학습된 프로브의 크기와 정렬이 중요한 아핀 구조의 경우 특히 "옵티마이저 - 아키텍처 상호작용"을 주요 설계 제약 조건으로 고려해야 함을 시사합니다.