Abstract
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
한국어 요약
📋 한 줄 요약
**[Transformer/긴 컨텍스트]** SDPA를 감싸는 학습 전용 hierarchical selection attention(Lighthouse)으로 긴 시퀀스 사전학습을 가속하고, 학습 종료 시 full attention 모델로 복원.
🎯 핵심 기여도
- Lighthouse Attention: 학습 시에만 사용하고 종료 시 제거 가능한 wrapper
- Subquadratic 계층적 pre/post-processing으로 sequence를 adaptive 압축·복원
- Symmetrical compression: query·key·value 동시 pooling, left-to-right causality 유지
- 두 단계 학습: 대부분 시간 Lighthouse로 사전학습 후 짧은 full-attention recovery
💡 핵심 아이디어
극단적 시퀀스 길이의 causal transformer 학습은 SDPA의 quadratic 비용에 막혀 있다. 학습 시에만 hierarchical selection으로 시퀀스를 압축하고 backward는 gradient-free로 두면 효율적이고, 학습 끝에 짧게 full attention으로 회복하면 표준 모델을 얻을 수 있다.
🔬 기술적 접근법
- **모델/방법론**: Causal Transformer LLM, ordinary SDPA를 wrapping
- **핵심 기법**: Subquadratic hierarchical adaptive compression, Q·K·V symmetrical pooling (causality 보존), gradient-free hierarchical selection, two-stage 사전학습 + 회복
📊 주요 결과
- Small-scale LLM 사전학습 실험에서 full attention 대비 학습시간 단축
- 회복 단계 후 더 낮은 final loss
- 다른 모든 설정 동일 조건에서 효과 검증
💭 의의 및 한계
**의의**: 긴 컨텍스트 LLM 학습 비용 절감, 표준 SDPA와 호환되어 추론 단계 변경 없음. **한계**: 현재 small-scale 검증에 그쳐 매우 큰 모델·매우 긴 시퀀스에서의 효과는 추가 실험 필요.
🚀 실용적 활용
- 긴 컨텍스트 LLM 사전학습 비용 절감
- Long-document·코드베이스 학습 모델 확장
- 기존 Transformer 학습 파이프라인에 drop-in 통합