Parallax: 새로운 어텐션 기술

Parallax는 LLM의 핵심 연산인 어텐션 구조를 개선한 매개변수화된 로컬 선형 어텐션(LLA) 메커니즘이다.

기존 LLA는 수치적 안정성과 확장성 문제로 LLM 사전 학습에 적용하기 어려웠으나, Parallax는 수치 해석기를 제거하고 KV 공분산을 탐색하는 추가적인 쿼리형 프로젝터를 학습하여 이 문제를 해결했다.

주요 특징은 다음과 같다:

• 하드웨어 인식 알고리즘: FlashAttention보다 높은 연산 강도(arithmetic intensity)를 구현하여 어텐션을 연산 중심(compute-bound) 영역으로 전환했다.
• 성능: 프로토타입 디코드 커널은 다양한 배치 크기와 컨텍스트 길이에서 FlashAttention 2/3와 대등하거나 더 뛰어난 성능을 보인다.
• 확장성: 0.6B 및 1.7B 규모의 사전 학습을 통해 Perplexity 개선과 다운스트림 벤치마크 성능 향상을 입증했다.

또한, Muon 옵티마이저와의 결합이 Parallax의 역량을 극대화한다는 점을 발견하며, 아키텍처와 옵티마이저 간의 강력한 공동 설계(co-design) 효과를 입증했다.