라이트하우스 어텐션을 활용한 장기 문맥 사전 학습

초록

극도로 긴 시퀀스 길이에서 인과적 트랜스포머를 훈련할 때, 스케일드 닷-프로덕트 어텐션(SDPA)의 이차 시간 및 메모리로 인해 병목 현상이 발생합니다. 본 연구에서는 일반 SDPA를 감싸는 훈련 전용 대칭 선택 기반 계층적 어텐션 알고리즘인 Lighthouse Attention을 제안하며, 이는 훈련 종료 시 쉽게 제거할 수 있습니다. 또한, 우리의 계층적 선택은 그래디언트가 없으므로 복잡하고 잠재적으로 비효율적인 역방향 패스 커널을 다루지 않아도 됩니다. 본 연구의 기여는 세 가지입니다: (i) 시퀀스의 적응형 압축 및 압축 해제를 수행하는 서브쿼드라틱 계층적 전처리 및 후처리 단계, (ii) 좌에서 우로의 인과성을 보존하면서 쿼리, 키, 값을 동시에 풀링하여 병렬성을 크게 향상시키는 대칭 압축 전략, (iii) 대부분의 훈련 시간 동안 Lighthouse Attention으로 사전 훈련하고 짧은 훈련을 통해 마지막에 전체 어텐션 모델을 복구하는 2단계 훈련 접근법. 우리는 다른 모든 설정을 동일하게 맞춘 전체 어텐션 훈련과 비교하여 우리 방법의 효과를 보여주는 예비 소규모 LLM 사전 훈련 실험을 수행했으며, 복구 단계 이후 더 빠른 총 훈련 시간과 더 낮은 최종 손실을 달성했습니다. 전체 코드는 https://github.com/ighoshsubho/lighthouse-attention 에서 확인할 수 있습니다.

English

Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention