CRISP, 구조 인식 라우팅으로 장문맥 희소 프리필 개선

들어가며

장문맥 LLM 추론에서 생성 단계 이전의 프리필은 중요한 병목입니다. 모델은 입력 전체에 대해 자기 어텐션을 계산해야 하며, 밀집 어텐션의 비용은 문맥 길이에 따라 제곱으로 증가합니다. 희소 어텐션은 계산량을 줄일 수 있지만, 고정 패턴은 입력마다 달라지는 중요 위치를 반영하기 어렵습니다. 오프라인 프로파일링에 의존하는 방법도 새로운 입력의 구조를 충분히 포착하지 못할 수 있습니다. CRISP（Cliff-awaRe Input-adaptive Sparse Prefilling）는 실행 중 어텐션 구조를 판단하는 동적 희소화 방식을 더 가볍고 정밀하게 만드는 데 초점을 둡니다.

두 가지 핵심 변화

• 구조에서 직접 라우팅한다. 기존 동적 방법은 풀링된 행렬 곱을 수행한 뒤 Jensen-Shannon Divergence（JSD）를 계산해 어텐션 헤드와 희소 패턴의 적합도를 추정할 수 있습니다. CRISP는 라우팅 판단이 프록시 어텐션 맵의 구조에 이미 드러난다고 봅니다. 제안된 C_struct는 Vertical-Slash 패턴과 호환되는 위치에 분포한 어텐션 질량을 측정합니다. 논문에 따르면 이 지표는 JSD의 라우팅 결정을 재현하면서 풀링 행렬 곱과 이어지는 KL 발산 계산을 제거합니다.

• sink와 노이즈를 반영해 선택량을 정한다. softmax 이후의 어텐션 질량은 소수의 위치에 큰 값이 몰린 뒤 낮은 값의 긴 꼬리가 이어지는 ‘질량 절벽’ 형태를 보일 수 있습니다. 누적 질량이 일정 비율에 도달할 때까지 위치를 추가하는 방식은 문맥이 길어질수록 배경 영역을 O(n) 규모로 계속 포함시킵니다. 그 결과 계산량은 늘지만 새롭게 얻는 정보는 제한적일 수 있습니다. CRISP는 노이즈 바닥에서 출발한 sink-aware threshold를 사용해 중요한 구조는 유지하고 배경 누적은 억제합니다.

결과와 의미

제공된 자료에 따르면, 두 모델 계열을 대상으로 한 InfiniteBench、RULER、LongBench 실험에서 CRISP는 전반적으로 가장 강한 희소 방법으로 평가되었습니다. 검색 중심 작업에서는 정확한 밀집 어텐션과 비슷하거나 더 나은 결과를 냈고, 기준선 대비 최대 28.0%포인트를 회복했습니다. 512k 토큰 문맥에서는 어텐션 부분을 최대 5.30배 빠르게 만들었다고 보고합니다. 논문은 이 성능 향상의 주된 이유로 선택 과정에서 O(n) 배경 노이즈를 제거하면서 어텐션의 구조적 특성을 보존한 점을 제시합니다.

CRISP의 시사점은 동적 희소 추론의 문제가 패턴 선택 하나로 끝나지 않는다는 데 있습니다. 어떤 패턴을 선택할지 결정하는 라우팅 자체에도 비용이 있으며, 얼마나 많은 위치를 보존할지 정하는 누적 질량 규칙도 초장문맥에서는 부적절할 수 있습니다. CRISP는 각각을 구조 기반 프록시와 sink 인식 임계값으로 다룹니다. 다만 현재 자료만으로는 모든 모델 구조, 하드웨어 환경, 희소 패턴에 대한 일반화 여부까지 확인할 수 없으므로 후속 검증이 필요합니다.

출처: Hugging Face Daily Papers