PISA는 피라미드 블록 선택으로 긴 문맥 주의 비용을 로그 선형으로 줄인다.

무엇인가

긴 문맥 언어모델을 키울 때 가장 큰 걸림돌은 셀프 어텐션 비용이 시퀀스 길이의 제곱으로 늘어난다는 점이다. 블록 희소 어텐션은 이 비용을 줄일 대안이지만, 어떤 블록을 남길지 고르는 과정이 다시 병목이 된다. 기존 방식은 모든 쿼리-블록 조합을 점수화해야 해서 시퀀스 길이에 대해 여전히 제곱 복잡도를 가진다. PISA는 바로 이 선택 단계를 공략한다.

어떻게 동작하나

PISA는 피라미드 Top-K 선택 전략을 쓰는 블록 희소 어텐션이다. 핵심은 여러 계층을 두고 후보를 점진적으로 좁혀 가장 관련 있는 키를 더 효율적으로 찾는 것이다. 이를 위해 키를 거친 수준에서 세밀한 수준으로 이어지는 계층으로 구성하고, 가장 거친 수준부터 선택을 시작한다.

무엇과 다른가

각 계층에서는 제한된 후보 집합에 LogSumExp 점수를 적용해 다음의 더 세밀한 계층으로 넘길 후보를 고른다. 이 과정을 가장 세밀한 계층에 도달할 때까지 반복한다. 풀링을 통해 O(log N)개의 키 계층을 만들고, 전체 복잡도는 O(N log N)이 된다. 여기서 N은 시퀀스 길이다.

어떻게 쓰나

저자들은 학습과 추론 모두를 위한 하드웨어 인식 Triton 커널도 개발했다. 이 커널은 계층적 라우팅과 LogSumExp 점수 계산을 융합하며, 쿼리-키 점수 행렬을 실제로 메모리에 만들지 않는다. 이는 긴 문맥에서 메모리 사용량과 커널 실행 효율을 함께 겨냥한 설계로 읽힌다.

전제와 한계

실험은 언어 모델링 과제에서 수행됐다. 초록에 따르면 베이스라인과 비교해 상식 추론 같은 벤치마크에서는 비슷한 성능을 보였고, 검색 과제에서는 더 나은 결과를 냈다. 구체적인 수치나 데이터셋 이름, 베이스라인 세부 구성은 초록에 제시되지 않아 여기서도 다루지 않는다.

개발자 입장에서는 긴 문맥을 다루는 LLM의 학습·추론 비용을 낮추는 후보 기술로 볼 수 있다. 특히 검색 증강 생성이나 긴 문서 검색처럼 관련 키를 골라내는 작업에서 이점이 있는지 확인할 가치가 있다. 실제 도입 전에는 Triton 커널 지원 환경, 계층 구성과 후보 집합 크기, 풀링 방식이 자신의 시퀀스 길이와 배치 크기에서 어떤 메모리·지연 시간을 만드는지 봐야 한다.

초록은 명시적인 한계를 길게 밝히지는 않는다. 다만 O(log N) 계층을 만들기 위한 풀링, 제한된 후보 집합, 하드웨어 인식 커널이라는 전제가 결과에 영향을 준다. 따라서 공개된 벤치마크 외의 도메인, 다양한 시퀀스 길이, 베이스라인 대비 공정성, 그리고 학습·추론 양쪽의 실제 처리량을 추가로 확인해야 한다.