Papers·2주 전
Northwestern, LLM 사전학습 가능한 Local Linear Attention 'Parallax' 공개 — 1.7B에서 perplexity 개선
Northwestern 팀이 softmax attention의 local constant 추정을 local linear 추정으로 업그레이드한 Parallax를 제안했습니다. FlashAttention 대비 compute-bound로 전환해 디코딩 속도를 유지하면서, 0.6B와 1.7B 규모 사전학습에서 perplexity를 일관되게 낮추고 downstream 벤치마크로 전이되는 이득을 확인했습니다. 단, Muon optimizer가 Parallax의 성능을 unlocking하는 현상이 발견되어 optimizer 선택이 중요해졌습니다.
Northwestern 팀이 softmax attention의 local constant 추정을 local linear 추정으로 업그레이드한 Parallax를 제안했습니다.
핵심 결론
- 성능 — 0.6B와 1.7B 규모 사전학습에서 perplexity가 일관되게 개선되었고, 이득이 downstream 벤치마크로 전이됩니다.
- Pareto 개선 — 파라미터 수와 compute를 통제한 비교에서도 Parallax가 softmax attention을 능가하는 Pareto 개선을 보였습니다.
방법
- Local Linear Attention — 기존 softmax attention의 local constant 추정을 local linear 추정으로 대체해 associative memory의 bias-variance tradeoff를 개선합니다.
- Parallax — LLA의 수치 해석기를 제거하고, KV covariance를 probing하는 추가 query-like projector를 학습합니다.
- 하드웨어 최적화 — FlashAttention 대비 arithmetic intensity를 높여 attention을 compute-bound 영역으로 전환하는 하드웨어 인식 알고리즘을 설계했습니다.
- Muon과의 공진 — Muon optimizer가 Parallax의 성능을 unlocking하는 현상을 발견했습니다. 이는 architecture-optimizer codesign의 첫 실증 사례입니다.
한계·조건
- 스케일 — 사전학습은 1.7B까지 검증되었으며, 더 큰 규모(7B 이상)에서의 추세는 추가 실험이 필요합니다.
- 디코딩 속도 — 프로토타입 디코드 커널은 FlashAttention 2/3와 비슷하거나 약간 빠른 수준이며, 다양한 배치 크기와 컨텍스트 길이에서 테스트되었습니다.
- 코드 — 논문에서 코드 공개 여부는 명시되지 않았습니다.
편집자 한 줄
Muon optimizer가 특정 attention 구조와 시너지를 낸 점은 흥미롭습니다. 향후 더 큰 모델에서도 같은 패턴이 유지될지 지켜볼 만합니다.
- #attention
- #llm
- #pretraining
- #northwestern
Northwestern University