ZHA는 KV 캐시를 트리로 요약해 어텐션 비용을 준(sub)이차로 낮춘다

무엇인가

어텐션의 이차 스케일링은 긴 문맥 LLM의 근본 병목이다. 고정 크기 상태를 쓰는 슬라이딩 윈도우·선형 어텐션 계열은 문맥이 길어지면 상태가 포화되어 성능이 떨어지고, 전체 어텐션을 일부 층에 남기는 하이브리드도 그 몇 개 층이 결국 비용을 지배한다. 희소 어텐션(DSA, MSA 등)은 질의마다 top-K KV를 골라 전체 어텐션 비용을 상수로 만드는 대신, 인덱서 자체가 KV 전체를 스코어링하므로 비용이 인덱서로 옮겨간다. 논문은 100만 토큰 근처에서 해당 층 시간의 60~90%가 여전히 인덱싱에 쓰인다고 지적한다. 저자들이 정리한 필요조건은 두 가지다. 질의 의존적 희소성, 그리고 선형보다 느리게 자라지만 멈추지 않는 활성 상태 크기다.

어떻게 동작하나

출발점은 관측이다. 긴 소설 데이터에서 전역 어텐션 층의 top-p KV 부분집합(어텐션 확률 합이 p 이상이 되는 최소 집합) 크기를 재 보면, 문맥 길이에 대해 p에 따라 지수가 달라지는 준선형 멱법칙으로 증가한다. 즉 어텐션의 실효 수용장은 선형보다 느리게 커진다. 반면 최근 W개 질의의 top-p 집합 합집합이 차지하는 프리픽스 비율 U_p(W)를 재면 W가 커질수록 1에 가까워진다. 어떤 질의에는 무시할 만한 키가 다른 질의에는 중요해진다는 뜻이므로, KV 캐시를 영구 삭제하는 대신 전부 보관하고 질의별로 골라야 한다. 여기서 전체 KV와 질의를 일일이 비교하지 않으려면 계층적 검색이 필요하다.

무엇과 다른가

ZHA의 트리는 연속된 키 청크를 재귀적으로 요약해 만든다. 청크 K_C의 비정규화 로그 어텐션 질량 F_C(q)=LSE_j(q·k_j/√d)를 기준 질의 q̂_r에서 1차 테일러 전개하면, 접평면 t_r(q)=q·k̂_r/√d + b̂_r이 된다. 여기서 압축 키 k̂_r은 기준 질의가 청크에 대해 수행한 어텐션 출력(Attn(q̂_r, K_C, K_C))이고, 바이어스 b̂_r은 F_C(q̂_r)에서 q̂_r·k̂_r/√d를 뺀 값이다. 미래 질의에 대한 청크의 로그 질량은 이 R개 접평면의 최댓값으로 근사된다. 상위 레벨 노드는 자식 요약들을 같은 방식으로 병합해 만들며, RoPE가 있는 모델에서는 부모의 회전된 기준 질의로 자식을 재평가할 때 1/√d⟨k̂_{r,c}, q̂_{r,P}−q̂_{r,c}⟩ 보정항을 더한다. GQA 계열은 KV 헤드마다 트리 하나를 만들어 질의 헤드를 공유한다. 청크 크기가 C라면 1레벨 요약은 C 디코드 스텝마다, 2레벨은 C² 스텝마다 추가되며, 저자들의 루프라인 모델에서 트리 구축 비용은 시퀀스 전체에 상각해 1% 이하다.

어떻게 쓰나

학습되는 것은 트리 전체에서 공유되는 기준 질의 R개뿐이다. 베이스 모델은 동결하고, 각 트리 레벨에서 원래 키 인덱스를 청크로 나눈 뒤 밀집 어텐션 가중치를 청크 단위로 합산한 목표 확률 p_{ℓ,C}(q)와 근사 확률 p̂_{ℓ,C}(q)의 KL 발산을 모든 레벨에 대해 최소화한다. 실험에서 4B~35B 모델을 32K 문맥, 40억 토큰으로 학습했고 최대 학습률은 3.2×10⁻²였다. R=16, KV 헤드당 트리 하나인 Qwen3.6-35B-A3B 기준 추가 파라미터는 81,920개로 약 0.00024%다. 전체 모델 파인튜닝도, 어텐션 구조 변경도 필요 없다.

전제와 한계

추론 시 라우팅은 레벨마다 top-p로 이뤄진다. 후보 청크를 점수 순으로 정렬한 뒤 추정 어텐션 질량의 p 비율을 포함하는 최소 프리픽스 m_ℓ(q)를 선택하고, 선택된 청크의 자식과 부모 없는 청크를 다음 레벨 후보로 넘긴다. 고정 예산 top-K와 달리 질량이 소수 노드에 집중되면 적게, 거의 균등하면 많이 고른다. 이 적응성이 길이 외삽의 근거다. 또한 레벨마다 후보를 줄여주므로 top-p 정렬 비용이 감당 가능한 수준으로 유지된다. 인덱스 0번 토큰은 어텐션 싱크로 간주해 항상 포함하고, 현재 진행 중인 청크의 가시 키도 함께 넣는다.

실험은 p-RoPE를 쓰는 Qwen 계열(256K~512K 문맥에서 사전학습된 4B~35B)과 NoPE·하이브리드인 Nemotron3 계열에서 이뤄졌다. 저자들은 수학적 추론, 에이전틱 코딩(SWE-bench), 장문 검색 평가에서 사실상 성능 저하가 없었다고 보고한다. RULER에서는 개별 과제 성능이 74~82%로 나왔고, 라우터가 선택한 키 집합은 밀집 top-p 집합보다 항상 더 희소했다. 과제별 편차도 해석과 맞아떨어진다. 반복 필러가 깔린 variable_tracking·niah_single_1은 평균 과제의 절반 정도 키만 읽고, 에세이 속 바늘 찾기와 문서 QA는 0.82~1.26배, 후보를 훑어야 하는 cwe·fwe·niah_multikey_2/3는 약 1.6배를 읽는다. 트리 깊이 최대 6레벨(R=16, C=4, L=6, p=0.9) 같은 다양한 기하에서도 성능이 유지됐고, 학습된 기준 질의는 하나로 붕괴하지 않고 방향 원뿔을 이루며(20개 KV 그룹 평균 쌍별 코사인 유사도 0.44), 상위 레벨일수록 소수 기준 질의 사용이 집중되고 리프에 가까울수록 균형을 이룬다.

개발자 입장에서 이 논문의 실질적 가치는 변환 비용이다. 이미 pretrained된 모델에 소량의 KL 워밍업만 얹어 긴 문맥 어텐션을 바꿔 끼우는 절차를 제시하고, 파라미터 증가가 0.0002% 수준임을 수치로 보여준다. 다만 저자들의 루프라인 성능 모델은 MSA 대비 교차점을 p=0.75의 성능 좋은 기하에서 약 50만 토큰으로 예측한다. 즉 이득은 100만 토큰 이상에서 본격화되며, 그 아래 구간에서는 이득이 제한적일 수 있다. 자체 서빙 스택에서 KV 캐시를 전부 보관할 메모리가 있고, 100만 토큰급 문맥이 필요하며, RULER·SWE-bench 같은 과제에서 희소화 후 품질을 직접 재검증할 계획이 있을 때 검토할 만하다.

한계는 저자들이 분명히 밝힌다. 경험적으로 256K 문맥까지는 잘 동작하지만, 희소도가 높아지고 문맥이 길어질수록 성능이 저하된다. 더 근본적인 문제는 top-p의 깊이 스케일링이다. 레벨마다 같은 p를 고정하면 유효 p는 p^L로 0에 수렴해 검색 상태가 상수 크기로 포화되고, 반대로 p^{1/L}로 보정하면 p가 1로 가서 상위 레벨이 전부 선택되어 이차 스케일링으로 되돌아간다. 준선형을 유지하려면 데이터의 멱법칙을 따라 p를 깊이에 따라 조정해야 하는데, 저자들이 관측한 것은 256K까지이며 수억~수십억 토큰에서도 성립하는지는 미해결이다. 또한 수백만 토큰 규모의 평가 성능은 실제로 검증되지 않았고, NoPE 모델에서의 길이 외삽은 불확실하며, 라우팅에서 잘린 서브트리가 기여하도록 하는 값 요약(value summary)은 아직 구현되지 않았다. 1B 미만 소형 모델은 교사 어텐션 자체가 노이즈여서 변환이 어려웠다는 점도 규모 의존성을 시사한다.