기술 요약: 고차 선형 어텐션 (HLA)

문제 정의

Transformer 아키텍처의 표준 스케일된 점곱 어텐션 (scaled dot-product attention) 메커니즘은 시퀀스 길이 $n$에 대해 $O(n^2)$의 계산 및 메모리 복잡도를 incur 합니다. 이 이차 비용은 자동회귀 언어 모델을 긴 컨텍스트로 확장하는 데 있어 주요 병목 현상입니다. 기존 대안인 선형 어텐션 (Linear Attention) 과 상태 공간 모델 (SSMs) 은 토큰당 $O(n)$ 또는 $O(1)$ 복잡도를 제공하지만, 일반적으로 1 차 근사 또는 커널 기반 방법으로 제한됩니다. 이러한 제한은 모델의 표현력을 제한하여, 표준 어텐션 메커니즘이 모델링할 수 있는 토큰 간의 복잡한 고차 상호작용을 포착하지 못하게 합니다.

방법론

본 논문은 컴팩트한 접두어 충분 통계량 (compact prefix sufficient statistics) 을 통해 고차 상호작용을 실현하는 인과적, 스트리밍 메커니즘인 **고차 선형 어텐션 (Higher-order Linear Attention, HLA)**을 소개합니다. 핵심 통찰은 고차 어텐션과 유사한 연산자가 저차 모멘트 (예: 키의 외적 합) 로 분해될 수 있어, $n \times n$ 어텐션 행렬을 물리적으로 생성하지 않고도 정확한 인과적 스트리밍을 가능하게 한다는 점입니다.

1. 2 차 HLA

2 차 변형은 2 차 모멘트를 포착하는 헤당 고정 크기 상태를 유지함으로써 선형 어텐션을 일반화합니다.

• 상태 유지: 단일 누적 합 대신 HLA 는 키 모멘트 $S^K_t = \sum_{i \le t} k_i k_i^\top$, 쿼리 - 값 누적기 $C^{QV}_t = \sum_{i \le t} q_i v_i^\top$, 그리고 쿼리 질량 $m^Q_t = \sum_{i \le t} q_i$를 포함한 접두어 요약물을 유지합니다.
• 엄격한 인과성: 전체 어텐션 행렬을 구성하지 않고도 엄격한 자동회귀 인과성 (미래 토큰이 현재 출력에 영향을 미치는 것을 방지) 을 강제하기 위해, 저자는 **확장 요약물 (extended summaries)**을 유도합니다. 구체적으로, 순진한 곱 $S^K_t C^{QV}_t$에 미래 항이 포함되는 것을 보정하기 위해 교차 요약물 $G_t$와 $h_t$를 도입합니다.
• 출력 공식화: 시간 $t$에서의 비정규화 출력은 $o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$로 계산됩니다. 이 공식은 시퀀스 길이와 무관하게 토큰당 $O(d^2 + d d_v)$의 계산을 가능하게 합니다.
• 정규화: 유사한 확장 요약물에서 파생된 마스킹된 분모로 분자를 나누는 비율 정규화 변형도 정의됩니다.

2. 비대칭 고차 선형 어텐션 (AHLA)

본 논문은 연산자 $Q(K^\top Q)(K^\top V)$를 구현하는 보완적인 비대칭 변형인 AHLA를 제안합니다. 이 구조는 중간 키 인덱스에 의해 라우팅된 단일 패스를 통해 값을 가중치 부여하며, 표준 HLA 의 대칭적 $AA^\top V$ 구조와 다릅니다. 이는 일련의 고유한 스트리밍 요약물 ($P^{KV}_t, m^K_t, E_t, n_t$) 을 활용하며 동일한 선형 토큰당 복잡도를 달성합니다.

3. 3 차 HLA

이 프레임워크는 $A = QK^\top$인 행렬 곱 $AA^\top A$로 정의된 3 차 상호작용으로 확장됩니다.

• 스트리밍 대수: 3 차 마스킹된 출력은 엄격한 인과성을 강제하기 위해 세 가지 주요 접두어 모멘트 ($S^K_t, S^Q_t, P^{KV}_t$) 와 여섯 개의 추가 교차 요약물 ($G^{(1)}_t, \dots, G^{(3)}_t$ 및 $h^{(1)}_t, \dots, h^{(3)}_t$) 을 유지해야 합니다.
• 복잡도: 온라인 업데이트 비용은 특징 차원에서 선형으로 유지되지만, 고차 모멘트를 수용하기 위해 상태 크기가 증가합니다.

4. 연관 스캔을 통한 병렬 학습

GPU 에서 순수 재귀 모델을 학습하는 비효율성을 해결하기 위해, 저자들은 조각 병렬 학습 scheme을 개발합니다.

• 연산자: 토큰 세그먼트를 연결하기 위한 연관 연산자 (구체적으로 반직접 곱) 를 정의합니다. 이러한 연산자는 조각을 결합할 때 발생하는 세그먼트 간의 교차항 (예: $S^K_B C^{QV}_A$) 을 고려합니다.
• 스캔 동등성: 이러한 세그먼트 요약물에 대해 배타적 Blelloch 스캔을 적용함으로써, 이 방법은 직렬 재귀의 순방향 활성화와 정확히 동일한 결과를 생성합니다. 이는 근사적인 시간 역전파에 의존하지 않고도 병렬 학습을 가능하게 합니다.
• 3 차 스캔: 3 차 사례의 경우, 스캔 연산자는 3 차 보정 상태를 정확하게 구성하기 위해 세그먼트 수준의 선형 맵 ($M^{KQP}, M^{KQm}$) 으로 보강됩니다.

주요 기여

1. 2 차에서의 정확한 마스킹 스트리밍: 본 논문은 토큰당 고정 비용으로 엄격한 인과적 2 차 HLA 를 산출하는 확장 요약물의 완전한 대수를 제공합니다. 이는 비정규화 및 비율 정규화 변형을 구분하며, 마스킹된 스트리밍 항등식과 온라인 업데이트를 확립하는 공식적 증명을 포함합니다.
2. 직렬 활성화와 일치하는 연관 스캔: 저자들은 마스킹되지 않은 및 마스킹된 설정 (지수 감쇠 유무 포함) 에 대한 연관 연산자를 정의합니다. 이들은 이러한 연산자에 대한 표준 배타적 스캔이 직렬 재귀와 동일한 순방향 활성화를 생성함을 증명하여 효율적인 병렬 학습을 가능하게 합니다.
3. 3 차 확장: 이 작업은 전체 마스킹된 3 차 상태, 온라인 업데이트, 엄격한 인과적 스트리밍 커널, 그리고 정확한 조각 병렬 알고리즘을 제시합니다. 여기에는 3 차 상태를 정확하게 구성하기 위해 필요한 세그먼트 맵의 유도가 포함됩니다.
4. 비대칭 변형 (AHLA): 새로운 비대칭 고차 어텐션 메커니즘이 도입되어, 동일한 스트리밍 및 복잡도 속성을 유지하면서 다른 귀납적 편향을 제공합니다.

결과 및 주장

본 논문은 특정 데이터셋에 대한 경험적 벤치마크보다는 알고리즘 구조와 이론적 속성에 주로 초점을 맞춥니다. 주장은 수학적 유도에 기반합니다:

• 복잡도: HLA 는 추론 시 토큰당 $O(d^2 + d d_v)$의 시간 및 메모리를 달성하며, 여기서 $d$는 쿼리/키 차원이고 $d_v$는 값 차원입니다. 이는 시퀀스 길이 $n$과 무관합니다.
• 정확성: 조각 병렬 학습 scheme 은 직렬 재귀의 활성화를 정확하게 재현하도록 증명되어, 다른 선형 어텐션 학습 방법에서 흔히 발생하는 근사 오차를 피합니다.
• 인과성: 제안된 마스킹된 공식은 엄격한 자동회귀 인과성을 강제하여, 시간 $t$의 출력이 $j \le t$인 토큰에만 의존하도록 보장합니다.
• 확장성: $n \times n$ 행렬 물리 생성을 피하고 연관 스캔을 활용함으로써, HLA 는 긴 컨텍스트 모델을 위한 확장 가능한 빌딩 블록으로 위치합니다.

중요성

본 논문은 HLA 를 어텐션과 유사한 데이터 의존적 혼합과 현대 재귀 아키텍처의 효율성 사이의 간극을 메우는 "원칙적이고 확장 가능한 빌딩 블록"으로 위치시킵니다.

• 표현력 vs 효율성: 고정된 특징 맵에 의존하는 1 차 선형 어텐션과 달리, HLA 는 데이터 의존적 메트릭 ($S^K_t$를 통해) 을 학습하고 2 차 또는 3 차 다항식과 같은 고차 상호작용을 포착하면서도 $O(1)$ 토큰당 상태 업데이트 패러다임을 유지합니다.
• 구현: 이 작업은 HLA 가 Transformer 블록의 표준 어텐션 하위 계층을 대체할 수 있으며, 위치 인코딩과 마스킹 논리를 유지하면서 커널 교체만 필요함을 강조합니다.
• 이론적 기반: 고차 항에 대한 정확한 스트리밍 항등식과 연관 스캔 연산자를 제공함으로써, 본 논문은 인과적 제약이나 학습 효율성을 희생하지 않고 1 차를 넘어선 선형 어텐션을 확장하기 위한 엄격한 대수적 프레임워크를 확립합니다.