고차원 선형 어텐션

초록

스케일된 점곱 어텐션의 2차 비용은 장기간 문맥으로 자기회귀 언어 모델을 확장하는 데 있어 핵심적인 장애물입니다. 선형 시간 어텐션과 상태 공간 모델(SSM)은 확장 가능한 대안을 제공하지만, 일반적으로 1차 또는 커널 기반 근사에 제한되어 표현력을 제한할 수 있습니다. 우리는 컴팩트한 접두사 충분 통계량을 통해 고차 상호작용을 실현하는 인과적 스트리밍 메커니즘인 고차 선형 어텐션(HLA)을 소개합니다. 2차 경우에서 HLA는 일정한 크기의 상태를 유지하며 n x n 행렬을 구체화하지 않고 선형 시간에 토큰별 출력을 계산합니다. 우리는 폐쇄형 스트리밍 항등식, 두 개의 추가 요약 통계량을 사용하는 엄격한 인과적 마스킹 변형, 그리고 직렬 순환의 활성화를 정확히 재현하는 결합 스캔에 기반한 청크 병렬 훈련 기법을 제시합니다. 우리는 더 나아가 3차 및 고차로의 확장을 개요합니다. 종합적으로, 이러한 결과는 HLA를 어텐션과 유사한 데이터 기반 혼합과 현대적 순환 구조의 효율성을 결합한 원칙적이고 확장 가능한 구성 요소로 위치시킵니다. 프로젝트 페이지: https://github.com/yifanzhang-pro/HLA.

English

The quadratic cost of scaled dot-product attention is a central obstacle to scaling autoregressive language models to long contexts. Linear-time attention and State Space Models (SSMs) provide scalable alternatives but are typically restricted to first-order or kernel-based approximations, which can limit expressivity. We introduce Higher-order Linear Attention (HLA), a causal, streaming mechanism that realizes higher interactions via compact prefix sufficient statistics. In the second-order case, HLA maintains a constant-size state and computes per-token outputs in linear time without materializing any n times n matrices. We give closed-form streaming identities, a strictly causal masked variant using two additional summaries, and a chunk-parallel training scheme based on associative scans that reproduces the activations of a serial recurrence exactly. We further outline extensions to third and higher orders. Collectively, these results position HLA as a principled, scalable building block that combines attention-like, data-dependent mixing with the efficiency of modern recurrent architectures. Project Page: https://github.com/yifanzhang-pro/HLA.