Long Context Pre-Training with Lighthouse Attention

Lighthouse Attention은 긴 맥락 트랜스포머의 효율적 사전학습을 위한 4단계의 커널 분리형 계층적 어텐션 메커니즘을 제공하여, 스케일드 닷-프로덕트 어텐션의 2차 복잡도를 크게 줄이는 동시에 완전한 밀집 어텐션 복원 가능성을 유지한다.

Lighthouse Attention은 긴 맥락 트랜스포머의 효율적 사전학습을 위한 4단계의 커널 분리형 계층적 어텐션 메커니즘을 제공하여, 스케일드 닷-프로덕트 어텐션의 2차 복잡도를 크게 줄이는 동시에 완전한 밀집 어텐션 복원 가능성을 유지한다. 이 접근법은 98K 맥락 길이의 530M 파라미터 모델에서 최대 1.69배의 실측 시간 가속과 더 낮은 최종 훈련 손실을 달성했으며, 맥락 병렬화를 통해 1M 토큰까지 더 확장된다.

핵심 요약

• 스케일드 닷-프로덕트 어텐션(SDPA)의 2차 시간·메모리 복잡도(Θ(N²))는 극히 긴 시퀀스 길이(예: 128K, 1M+)에서 인과적 트랜스포머를 훈련하는 것을 연산적으로 감당 불가능하게 만든다.
• 사전학습용 기존 희소 어텐션 기법들은 비대칭적(키와 값만 풀링)이거나 구조적으로 얽혀 있어, FlashAttention 같은 고도로 최적화된 밀집 어텐션 커널의 재사용을 막는 맞춤형 희소 어텐션 연산자를 요구하는 경우가 많다.
• 희소 어텐션으로 사전학습된 모델이 추론 시 완전한 밀집 어텐션 모델로 효과적으로 기능하도록 보장하는 것은 큰 난제다.
• 어텐션 커널을 감싸 표준 FlashAttention이 동적으로 압축된 하위 시퀀스에서 작동할 수 있게 하는 4단계 파이프라인(피라미드 구성, 점수화 및 선택, 수집-시퀀스 어텐션, 산포-복원 재구성)을 도입한다.
• 여러 피라미드 레벨에 걸쳐 쿼리·키·값에 대칭적 평균 풀링을 적용하고, 매개변수 없는 투영-노름 점수기와 미분 불가능한 상위 k개 선택 과정을 결합하여 준-2차 처리를 수행한다.
• 두 단계 훈련 방식을 사용한다. 즉, 효율을 위해 Lighthouse Attention으로 1차 훈련한 뒤, 추론을 위해 완전한 밀집 SDPA를 사용한 짧은 재개 단계로 모델의 완전한 어텐션 역량을 복원하고 최적화한다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.