💬 NLP

Long Context Pre-Training with Lighthouse Attention

본 논문은 극단적인 시퀀스 길이에서 인과적 트랜스포머의 효율적인 2 차 미만의 복잡도로 사전 학습을 가능하게 하기 위해 표준 스케일된 점곱 어텐션을 감싸는 학습 전용, 기울기 없는 계층적 선택 알고리즘인 라이트하우스 어텐션을 소개하며, 이는 짧은 복구 단계를 통해 완전히 제거되어 더 빠른 학습과 더 낮은 최종 손실을 갖는 전체 어텐션 모델을 생성할 수 있습니다.

원저자: Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

게시일 2026-05-08

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

한 번에 거대한 도서관의 책들(매우 긴 텍스트) 전체를 읽어서 이야기를 이해하려고 상상해 보세요. AI 세계에서는 이를 '긴 문맥 학습(long-context training)'이라고 부릅니다.

문제는 표준 AI 모델(트랜스포머)이 연결고리를 찾기 위해 모든 단어를 다른 모든 단어와 비교하려 한다는 점입니다. 10 만 단어가 있다면 모델은 100 억 번의 비교를 수행해야 합니다. 이는 도서관에서 특정 문장을 찾기 위해 모든 단어를 다른 모든 단어에게 동시에 외치는 것과 같습니다. 이는 영원히 걸리며, 막대한 전기 비용을 소모하고 컴퓨터의 메모리가 부족해집니다.

이 논문은 이를 해결하기 위해 '등대 어텐션(Lighthouse Attention)'이라는 새로운 방법을 소개합니다. 간단한 비유를 통해 작동 원리를 설명하겠습니다.

1. 문제: '만능 시선'은 너무 무겁습니다

표준 AI 어텐션은 스타디움의 모든 사람을 동시에 바라보며 누가 누구와 대화하는지 확인하려는 집착을 가진 경비원처럼 작동합니다. 스타디움이 커질수록 경비원은 압도당합니다.

2. 해결책: '등대' 전략

등대 어텐션은 어둠의 바다를 비추는 등대처럼 모든 것을 고해상도로 보는 대신, 바다를 무시하지 않고 가장 중요한 부분을 빠르게 찾기 위해 계층적으로 스캔합니다.

논문에 설명된 3 단계 과정은 다음과 같습니다.

단계 A: '피라미드'(군중 요약)

거대한 군중(텍스트)이 있다고 상상해 보세요. 개별 얼굴 하나하나를 보는 대신, 그들을 작은 클러스터(가족이나 팀과 같은) 로 그룹화합니다.

기법: 이 논문은 여기서 독특한 작업을 수행합니다. 대부분의 다른 방법들은 '대화의 대상'(키와 값) 만 요약하고 '대화를 주도하는 사람'(쿼리) 은 고해상도로 남깁니다.
등대의 움직임: 등대는 모두를 동등하게 요약합니다. 요약의 피라미드를 생성합니다.
- 레벨 0: 모든 단일 단어.
- 레벨 1: 4 개의 단어가 하나의 그룹으로 요약됨.
- 레벨 2: 16 개의 단어가 하나의 그룹으로 요약됨.
- 이어서 계속됩니다.
  이것은 '초세부'에서 '큰 그림'까지 텍스트의 다중 레벨 지도를 생성합니다.

단계 B: '스포트라이트'(가장 좋은 부분 선택)

이제 모델은 이 피라미드의 어떤 부분이 전체 세부 사항으로 읽기에 중요한지 결정해야 합니다.

선택: 추가 학습이 필요 없는 간단한 무료 규칙을 사용하여 모든 그룹에 점수를 매깁니다. "어떤 그룹이 가장 많은 '에너지'나 중요성을 가지고 있는가?"라고 묻습니다.
컷: 피라미드의 모든 레벨에서 가장 중요한 상위 $K$ 개의 그룹을 선택합니다.
결과: 10 만 단어를 읽는 대신, 모델은 이제 가장 중요한 '조각들'(약 5,000 단어일 수 있음) 의 작고 밀집된 목록만 읽으면 됩니다.

단계 C: '플래시'(선택된 부분 읽기)

모델이 상위 5,000 단어를 선택하면, 이를 표준의 초고속 '플래시 어텐션(FlashAttention)' 엔진에 통과시킵니다. 목록이 이제 짧기 때문에 이 단계는 놀라울 정도로 빠르고 컴퓨터 메모리에 쉽게 들어갑니다.

3. '마법' 회복 (2 단계 학습)

이것이 이 논문에서 가장 중요한 부분입니다. 저자들은 다음과 같이 우려했습니다: "만약 AI 를 요약만 보도록 훈련한다면, 나중에 완전한 문장을 읽는 방법을 잊어버리지 않을까?"

이를 해결하기 위해 2 단계 학습 레시피를 사용합니다.

1 단계 (등대 단계): 대부분의 시간 동안 등대 방법을 사용하여 모델을 훈련시킵니다. 모델은 효율성을 배우고 올바른 하이라이트를 선택하는 법을 배웁니다.
2 단계 (회복 단계): 훈련의 마지막 짧은 기간 동안 '요약'과 '선택' 부분을 끕니다. 표준 방법을 사용하여 모델이 다시 전체 텍스트를 읽도록 강요합니다.

결과: 모델은 효율적인 훈련에서 '깨어나' 완전한 어텐션을 사용하는 방법을 완벽하게 기억합니다. 논문은 이 짧은 회복 기간 후, 모델이 전체 텍스트를 전체 시간 동안 훈련받은 모델만큼 (또는 더) 잘 수행하지만, 훨씬 더 빠르고 저렴하게 그 지점에 도달했다고 주장합니다.

이것이 왜 중요한가요?

속도: 이 논문은 매우 긴 텍스트 (10 만 단어 이상) 의 경우 이 방법이 표준 방법보다 17 배에서 21 배 더 빠르다고 보여줍니다.
불필요한 '쓰레기' 코드: '선택' 과정을 처리하기 위해 맞춤형 복잡한 컴퓨터 칩(커널) 을 구축해야 하는 다른 방법들과 달리, 등대는 실제 읽기를 위해 표준의 상용 컴퓨터 부품을 사용합니다. 데이터만 전달하기 전에 재배열할 뿐입니다.
대칭성: 문장의 '질문' 부분과 '답변' 부분을 동등하게 취급하여 수학을 더 깔끔하고 안정적으로 만듭니다.

결론

등대 어텐션은 똑똑한 연구 조교를 고용하는 것과 같습니다. 조교는 보고서 1,000 페이지를 단어별로 읽는 대신, 문서 전체를 빠르게 스캔하여 가장 중요한 50 개 단락을 하이라이트한 후, 오직 그 50 개 단락만 깊이 있게 읽습니다.

이 논문은 이렇게 AI 를 훈련시킨 후, 마지막에 전체를 읽는 것에 대한 빠른 '보충 과정'을 제공하면, 지능을 잃지 않는 초고속 독서가 된다는 것을 증명합니다.

기술 요약: Lighthouse Attention 을 통한 긴 컨텍스트 전처리

1. 문제 제기

극단적인 시퀀스 길이 (128K~1M+ 토큰) 에서 인과적 트랜스포머를 학습시키는 것은 현재 확장된 점곱 어텐션 (scaled dot-product attention, SDPA) 의 이차 시간 및 메모리 복잡도 ( $\Theta(N^2)$ ) 에 의해 병목 현상을 겪고 있습니다. FlashAttention 과 같은 방법들은 상수를 최적화하지만 점근적 복잡도를 변경하지는 않습니다. 기존 희소 어텐션 접근법들은 쿼리가 키의 부분집합에만 어텐션하는 선택 메커니즘으로 조밀한 어텐션을 대체함으로써 이를 완화하려 시도합니다. 그러나 이러한 방법들은 전처리 (pre-training) 에 적용될 때 두 가지 주요 설계 결함을 겪습니다:

비대칭성: 대부분의 방법 (예: NSA, HISA) 은 키와 값을 풀링하지만 쿼리는 전체 해상도로 유지합니다. 이는 계층 구조를 다중 스케일 표현이 아닌 압축된 메모리로 취급하여, 모델이 스케일 간 일관된 상호작용을 학습하는 능력을 제한합니다.
아키텍처 얽힘: 선택 로직이 종종 커스텀 희소 어텐션 커널 내에 내장됩니다. 이는 FlashAttention 과 같이 매우 최적화된 조밀한 어텐션 커널의 재사용을 방해하고 커스텀 역전파를 요구하여 학습 및 추론 배포를 복잡하게 만듭니다.

또한, 학습 시 희소 방법의 경우 정확성이 중요한 우려 사항입니다. 추론 전용 희소성 (유능한 조밀한 백본을 가정함) 과는 달리, 학습 시 희소 방법은 희소성이 제거된 후 결과 가중치가 유능한 조밀한 어텐션 모델로 기능할 수 있음을 증명해야 합니다.

2. 방법론: Lighthouse Attention

Lighthouse Attention 은 어텐션 커널 자체를 수정하지 않고 일반 SDPA 를 감싸도록 설계된 전처리 전용, 대칭적, 선택 기반 계층적 어텐션 알고리즘입니다. 이는 4 단계 파이프라인으로 작동합니다:

A. 대칭 피라미드 구성

기존 작업들이 컨텍스트 (K/V) 만 풀링하는 것과 달리, Lighthouse 는 쿼리, 키, 값 ( $Q, K, V$ ) 에 동시에 대칭 평균 풀링을 적용합니다.

시퀀스는 풀링 인자 $p$ 를 가진 $L$ 단계 피라미드로 구성됩니다.
0 단계는 전체 해상도 시퀀스입니다. 각 후속 단계 $\ell$ 은 이전 단계 $\ell-1$ 의 $p$ 개의 연속된 엔트리를 요약합니다.
이는 모든 단계에서 일관된 삼중체 $(Q^{(\ell)}, K^{(\ell)}, V^{(\ell)})$ 를 생성하여, 풀링된 쿼리와 키가 동일한 표현 공간에 존재하도록 보장합니다.

B. 파라미터 없는 스코어링 및 선택

계층적 선택기가 피라미드 엔트리를 순위 매겨 어텐션할 토큰을 결정합니다.

스코어링: 이 방법은 파라미터 없는 스코어링 함수를 사용합니다. 각 엔트리에 대해 투영된 $Q$ 와 $K$ 벡터의 $\ell_2$ 노름을 계산합니다. 더 거친 단계는 최대 풀링을 통해 더 세밀한 단계에서 스코어를 상속받아, 거친 범위가 가장 강력한 토큰의 중요성을 유지하도록 보장합니다.
선택: 퓨즈된 chunked-bitonic top-K 커널이 모든 단계에 걸쳐 상위- $K$ 개 엔트리를 통합적으로 선택합니다.
미분 불가능성: top-K 선택 단계는 이산적이고 미분 불가능합니다. 스트레이트-쓰루 추정기나 보조 손실은 사용되지 않습니다. 기울기는 gather, FlashAttention, scatter-back 연산을 통해 투영 행렬 ( $W_Q, W_K, W_V$ ) 로만 흐릅니다. 결과적으로 투영은 점수자를 속이도록 학습하는 것이 아니라, 선택되었을 때 유용한 값을 생성하도록 학습합니다.

C. 조밀한 서브시퀀스 어텐션

선택된 엔트리는 길이 $S \ll N$ 인 연속적이고 인과적으로 일관된 서브시퀀스로 수집됩니다.

수집 (Gather): 선택된 삼중체는 위상적으로 정렬되어 조밀한 시퀀스를 형성합니다.
어텐션: 이 서브시퀀스에서 표준 FlashAttention(또는 일반 SDPA) 이 실행됩니다. 입력이 조밀하고 연속적이므로 커스텀 희소 커널이 필요하지 않습니다.
분산 (Scatter-Back): 결정론적 scatter 커널이 어텐션 출력을 전체 시퀀스의 원래 위치로 다시 분배합니다.

D. 2 단계 학습 레시피

모델이 완전한 어텐션 능력을 유지하도록 하기 위해, 저자들은 2 단계 학습 접근법을 제안합니다:

1 단계: 대부분의 시간을 Lighthouse Attention 을 사용하여 전처리합니다.
2 단계: 동일한 체크포인트를 사용하여 완전 조밀한 SDPA로 짧은 기간 동안 학습을 재개합니다. 이 "회복 단계"는 모델이 처음부터 시작하지 않고도 완전한 어텐션 메커니즘에 가중치를 적응할 수 있게 합니다.

3. 주요 기여

대칭 계층적 어텐션: $Q, K, V$ 를 대칭적으로 풀링하여 압축된 메모리가 아닌 진정한 다중 스케일 표현을 생성하면서 좌우 인과성을 유지하는 새로운 설계입니다.
커널 분리: 선택을 어텐션 경로 외부에 배치함으로써, 조밀한 서브시퀀스에서 일반 FlashAttention을 활용합니다. 이는 매우 최적화된 조밀한 커널의 재사용을 가능하게 하고 역전파를 단순화합니다.
학습 시 정확성: 이 논문은 엄격한 경험적 기준을 제시합니다: Lighthouse 로 전처리된 모델은 짧은 조밀한 SDPA 재개 후 조밀한-from-scratch 베이스라인과 일치하거나 능가해야 합니다. 저자들은 Lighthouse 로 학습된 모델이 이 기준을 충족함을 보여줌으로써, 계층적 학습이 모델의 완전한 어텐션 사용 능력을 "비워내지" 않음을 증명합니다.
효율적인 구현: 퓨즈된 chunked-bitonic top-K 및 커스텀 scatter-back 커널의 사용은 멀티 노드 Blackwell GPU 에서 극단적인 컨텍스트 길이 (최대 1M 토큰) 로의 효율적인 학습을 가능하게 합니다.

4. 실험 결과

저자들은 98,304 토큰의 컨텍스트 길이를 가진 C4 데이터셋으로 학습된 530M 파라미터 Llama-3 스타일 디코더를 사용하여 실험을 수행했습니다.

회복 가능성: Lighthouse 로 전처리된 후 4,000 스텝 동안 조밀한 SDPA 로 재개된 모델은 동일한 토큰 예산으로 학습된 조밀한-from-scratch 베이스라인 (0.7237) 보다 더 낮은 최종 학습 손실(0.6980–0.7102) 을 달성했습니다. 이는 계층적 학습 신호가 모델의 조밀한 어텐션 능력을 저하시키지 않음을 확인시켜 줍니다.
처리량 및 속도 향상:
- Lighthouse 는 전처리 단계에서 조밀한 SDPA 대비 2 배의 처리량 우위(GPU 당 초당 토큰 수) 를 달성했습니다.
- 종단간 (End-to-end) 으로, 2 단계 레시피는 동일한 토큰 예산에 대해 조밀한-from-scratch 학습 대비 1.40 배에서 1.69 배의 월클럭 속도 향상을 달성했으며, 최종 손실은 일치하거나 더 낮았습니다.
- 512K 컨텍스트 길이에서 Lighthouse 는 순방향 pass 에서 21 배, 순방향 + 역방향 pass 에서 17.3 배 더 빨랐습니다.
확장성: 이 방법은 컨텍스트 병렬성을 사용하여 32 개의 Blackwell GPU 로 1M 토큰까지 성공적으로 확장되었으며, 조밀한 베이스라인 대비 속도 향상 우위를 유지했습니다.
Ablation 연구:
- 스코어러: 파라미터 없는 $\ell_2$ 노름 스코어러는 더 비싼 dilated softmax 스코어러와 유사한 성능을 보였지만 약 9% 더 빨랐습니다.
- 하이퍼파라미터: 더 작은 풀링 인자 ( $p=2$ ), 더 얕은 피라미드 ( $L=3$ ), 그리고 더 작은 top-K 예산 ( $k=1536$ ) 이 일반적으로 손실/성능 트레이드오프에서 가장 좋은 결과를 제공했습니다.
검색 (Needle-in-a-Haystack): 조밀한 SDPA 재개 후, Lighthouse 로 학습된 모델은 조밀한 베이스라인과 같거나 더 나은 검색률 (예: 평균 검색률 0.76 대 0.72) 을 달성하여 긴 컨텍스트 유지 능력을 검증했습니다.

5. 중요성 및 주장

이 논문은 Lighthouse Attention 이 표준 트랜스포머의 아키텍처 무결성을 유지하면서 어텐션 복잡도를 이차에서 준이차 (유한한 $K$ 에 대해 실질적으로 $N$ 에 선형) 로 줄임으로써 긴 컨텍스트 전처리의 근본적인 병목 현상을 해결한다고 주장합니다.

저자들은 그들의 방법이 전처리 전용임을 강조합니다. 선택 메커니즘은 제거되거나 (또는 모델이 전환되어) 추론을 위해 완전한 조밀한 모델이 됩니다. 이는 조밀한 백본에서 "정확성 바닥"을 물려받지만 희소 학습 과정이 유능한 조밀한 가중치를 산출함을 증명하지 않는 추론 전용 희소 방법과 Lighthouse 를 구분합니다.

저자가 인정하는 한계:

대칭 풀링은 모든 쿼리가 하나의 순방향 pass 에 공존한다고 가정하는데, 이는 자기회귀 디코딩 중 위반됩니다. 따라서 이 방법은 추론 준비가 된 모델을 생성하기 위해 조밀한 SDPA 재개에 의존합니다.
내부 어텐션은 수집된 서브시퀀스에서 $\Theta(S^2 d)$ 입니다. $N$ 에 대해 준이차적이지만, 선택 예산 $K$ 가 $N$ 에 따라 확장되어야 한다면 엄격하게 선형은 아닙니다.
이 방법은 현재 2 단계 학습 레시피가 필요하며, DSA 나 NSA 와 같은 네이티브 희소 추론 목표는 아직 통합되지 않았으나, 저자들은 이를 향후 방향으로 제안합니다.

요약하자면, Lighthouse Attention 은 추론 시 모델이 완전한 조밀한 어텐션 메커니즘을 완전히 활용할 수 있음을 보장하면서, 계산 비용을 크게 줄인 상태에서 백만 토큰 컨텍스트로 모델을 학습할 수 있는 경로를 제공합니다.