연구진은 SDPA의 한계를 극복하기 위해 Lighthouse Attention이라는 새로운 어텐션 알고리즘을 제안했습니다. 이 기법은 훈련 과정에서만 사용되며, 기존의 SDPA를 감싸는 방식으로 작동합니다. Lighthouse Attention은 시퀀스를 압축하고 해제하는 과정을 통해 훈련 시간을 단축하고 최종 손실을 줄입니다.
Lighthouse Attention은 쿼리, 키, 값을 동시에 풀링하는 대칭 압축 전략을 사용하여 좌우 인과 관계를 유지하며 병렬성을 향상시킵니다. 이 기법은 Lighthouse Attention으로 대부분의 훈련을 진행한 후, 짧은 훈련을 통해 전체 어텐션 모델로 복구하는 2단계 훈련 방식을 사용합니다.
연구진은 Lighthouse Attention의 효과를 검증하기 위해 소규모 LLM 사전 훈련 실험을 진행했으며, 다른 설정은 동일하게 유지하면서 전체 훈련 시간을 단축하고 최종 손실을 낮추는 결과를 얻었습니다. 관련 코드는 GitHub에서 확인할 수 있습니다.