PISA, 블록 희소 어텐션을 O(N log N) 복잡도로 개선

블록 희소 어텐션의 숨은 병목

장문맥 언어 모델의 확장을 어렵게 만드는 대표적인 요인은 자기 어텐션의 비용입니다. 표준 어텐션은 모든 쿼리와 키의 조합을 비교하므로 시퀀스 길이가 늘어날수록 계산량이 제곱 형태로 증가합니다. 블록 희소 어텐션은 중요한 키 블록만 남겨 실제 어텐션 계산량을 줄이는 접근입니다. 그러나 어떤 블록을 남길지 결정하기 위해 모든 쿼리와 블록을 먼저 점수화한다면, 선택 과정 자체가 여전히 이차 복잡도를 가질 수 있습니다.

ByteDance Seed 연구진이 발표한 PISA는 이 블록 선택 비용을 줄이기 위한 방법입니다. 핵심 아이디어는 전체 후보에서 한 번에 Top-K를 찾는 대신, 여러 해상도의 표현을 이용해 검색 공간을 단계적으로 줄이는 피라미드형 선택입니다.

거친 후보에서 세밀한 후보로

PISA는 풀링을 사용해 키 표현을 여러 단계의 계층으로 구성합니다. 상위의 거친 단계에서는 시퀀스의 넓은 영역을 적은 수의 표현으로 요약하고, 더 낮은 단계로 내려갈수록 후보를 세밀한 블록으로 확장합니다. 이 계층은 시퀀스 길이 N에 대해 O(log N)개의 단계로 만들어집니다.

각 쿼리에 대한 탐색은 가장 거친 단계에서 시작합니다. 현재 단계에서는 전체 키를 모두 검사하지 않고, 크기가 제한된 후보 집합에 대해 LogSumExp 기반 점수를 계산합니다. 그중 상위 후보만 다음 단계로 전달하고, 이 과정을 가장 세밀한 단계까지 반복합니다. 마지막에는 희소 어텐션에 사용할 키 블록이 선택됩니다.

이 방식은 두 가지 측면에서 비용을 줄입니다. 각 단계가 전체 시퀀스를 다시 검색하지 않기 때문에 후보 평가 범위를 통제할 수 있고, 대규모 전역 검색을 여러 번의 작은 지역적 선택으로 바꿀 수 있습니다. 논문은 이러한 계층적 라우팅을 통해 전체 복잡도를 O(N log N)으로 낮출 수 있다고 설명합니다.

하드웨어를 고려한 구현

이론적 복잡도만 줄이는 것으로 실제 효율이 보장되지는 않습니다. 연구진은 학습과 추론 모두를 대상으로 하드웨어 친화적인 Triton 커널을 구현했습니다. 계층적 라우팅과 LogSumExp 점수 계산을 하나의 흐름으로 결합하고, 전체 쿼리-키 점수 행렬을 메모리에 생성하지 않는 방식을 사용했습니다. 이 설계는 중간 결과의 메모리 부담을 줄이고 가속기에서의 실행에 더 적합하도록 돕습니다.

실험은 언어 모델링 과제를 중심으로 진행되었습니다. 초록에 따르면 PISA는 상식 추론과 같은 벤치마크에서 기준 방법과 비슷한 성능을 보였고, 검색 과제에서는 더 나은 결과를 기록했습니다. 긴 문맥 안에서 관련 정보를 찾아야 하는 검색 과제에서의 결과는 계층적 후보 선택이라는 설계와 잘 맞습니다. 다만 제공된 자료에는 모델 설정, 문맥 길이, 절대 점수, 전체 처리량이 제시되지 않았으므로 실제 가속 폭이나 모든 작업에 대한 일반성을 단정하기는 어렵습니다.

의미와 남은 과제

PISA의 중요한 점은 블록 희소 어텐션에서 블록을 고르는 라우팅 과정도 별도의 최적화 대상으로 다룬다는 데 있습니다. 희소 어텐션이 장문맥에서 실질적인 효율을 얻으려면 최종 어텐션 계산뿐 아니라 관련 블록을 찾는 과정에서도 조밀한 계산을 피해야 합니다.

다만 거친 표현을 이용한 초기 판단에는 정확도와 효율 사이의 절충이 있습니다. 중요한 영역이 초기 단계에서 탈락하면 이후의 세밀한 탐색으로 복구하지 못할 수 있습니다. 따라서 후보 예산, 풀링 방식, 과제별 민감도는 후속 검증에서 중요한 요소가 될 것입니다. PISA는 장문맥 어텐션을 전수 비교에서 계층적 검색으로 전환하는 하나의 실용적 설계 방향을 제시합니다.

출처：Hugging Face Daily Papers