HySparse2는 전체 어텐션층만 잇는 2단계KV공유 하이브리드 희소 어텐션이다.

무엇인가

이 논문이 푸는 문제는 에이전트 추론의 입력 편중 현상이다. 장기 문맥·다회차 에이전트는 짧은 행동을 생성하고 도구와 환경에서 훨씬 긴 관측을 받아들이는데, 상호작용 라운드마다 짧은 도구 호출이 그보다 훨씬 긴 검색 결과나 실행 트레이스, 문서를 반환한다. 관측이 누적될수록 에이전트는 확장하는 이력 전체에서 근거를 검색해 결합해야 하므로 어텐션 연산과 KV 캐시 저장이 함께 커진다. 따라서 효율적인 prefill, 압축된 KV 캐시, 정확한 장문맥 검색이 동시에 필요하다. 기존 HySparse는 전체 어텐션 레이어와 희소 어텐션 레이어를 교차 배치하고, 각 전체 어텐션 레이어가 선택 인덱스와 KV 캐시를 뒤따르는 희소 레이어에 공급하는 방식으로 연산과 저장을 줄였다. 그러나 prefill 단계에서 여전히 모든 레이어를 실행해야 했고, 블록 단위 희소성은 정밀도에 한계가 있었다.

어떻게 동작하나

제안 방법의 핵심은 2단계 KV 공유다. 바깥 단계인 KV Bridging은 YOCO처럼 백본을 self-decoder와 cross-decoder로 나누되, 전체 어텐션 레이어 사이만 연결한다. self-decoder는 전체 어텐션과 슬라이딩 윈도우 어텐션(SWA)을 혼합해 국소 모델링을 담당하고, cross-decoder는 전체 어텐션과 희소 어텐션을 혼합해 전역 검색을 담당한다. self-decoder의 i번째 전체 어텐션 레이어 입력 은닉 상태를 H_i^self, cross-decoder의 j번째 전체 어텐션 레이어 입력 은닉 상태를 H_j^cross라 하면, cross-decoder 레이어 j는 K_j^cross = Proj_j^K(H_i^self), V_j^cross = Proj_j^V(H_i^self), Q_j^cross = Proj_j^Q(H_j^cross)로 키·값·쿼리를 계산한다. 즉 키와 값은 self-decoder의 은닉 상태에서, 쿼리는 자신의 현재 은닉 상태에서 만든다. 각 cross-decoder 전체 어텐션 레이어가 자기만의 K/V 프로젝션을 가지므로 같은 은닉 상태를 공유해도 레이어마다 다른 KV 캐시가 만들어지고, 두 디코더의 하이브리드 비율이 달라도 되므로 self-decoder의 전체 어텐션 레이어 하나가 여러 cross-decoder 전체 어텐션 레이어에 공급될 수 있다. 안쪽 단계인 KV Reuse는 HySparse의 설계를 유지해, 각 하이브리드 블록 안에서 희소 레이어가 전체 어텐션 레이어의 KV 캐시와 선택 인덱스를 그대로 재사용한다.

무엇과 다른가

HySparse2는 여기에 두 가지 수정을 가한다. 첫째, 블록 단위 선택을 토큰 단위 선택으로 바꿔 전체 어텐션 레이어에서 개별 토큰에 top-k를 적용하고, 뒤따르는 희소 레이어가 선택된 토큰의 KV 항목을 재사용한다. 둘째, 희소 레이어에 있던 별도 SWA 분기를 제거하고 대신 최근 토큰의 슬라이딩 윈도우를 희소 선택에 강제로 포함시킨다. 최근 토큰 윈도우를 항상 먼저 선택하고, 그 바깥에서 점수가 가장 높은 토큰을 이어서 선택하며, 두 집합 모두 전체 어텐션 KV 캐시에서 읽어 뒤따르는 희소 레이어가 재사용한다. 별도 SWA 분기를 없앤 이유는 prefill 조기 종료 때문이다. cross-decoder에 SWA 분기가 있으면 자기 은닉 상태의 프로젝션으로 KV 캐시 접미사를 만들어야 하고, 이는 이전 레이어의 은닉 상태 집합에 다시 의존해 깊이에 따라 선형으로 커진다. 이 연쇄 의존성 때문에 cross-decoder는 윈도우 크기보다 훨씬 긴 토큰 접미사를 처리해야 하고, 그 상태를 prefill 중에 계산하거나 bounded replay 같은 방법으로 근사해야 한다. HySparse2는 이 의존성을 구조적으로 제거해 prefill 중 cross-decoder 연산을 전혀 요구하지 않는다. 또한 전체 어텐션 레이어는 소수만 유지하는데, 이 레이어들이 정확한 어텐션 점수로 이후 희소 레이어에 오라클 토큰 선택을 제공하는 인덱서 역할을 하므로 별도 인덱서 모듈이나 증류 목적함수 없이 네이티브 종단간 학습이 가능하다.

어떻게 쓰나

실험은 80B-A3B MoE 모델로 진행했다. 49개 트랜스포머 레이어, 은닉 크기 2,048, 잔차 혼합 행렬을 항등으로 고정한 단순화된 mHC 변형을 사용한다. 약 500B 토큰을 32k 문맥으로 사전학습하고, 약 100B 토큰의 가벼운 후학습에서 에이전트 데이터를 혼합하고 문맥을 256k로 확장했다. 세 모델이 같은 데이터 혼합과 스케줄을 공유하며, Muon 옵티마이저와 WSD 스케줄, 최대 학습률은 사전학습 10^-3, 후학습 5×10^-5다. Hybrid SWA는 전체 어텐션 레이어 9개를, HySparse와 HySparse2는 각각 5개만 쓰면서 전체 어텐션 기준선과 비슷한 수준을 유지한다. Hybrid SWA와 HySparse는 GQA를, HySparse2는 MQA를 사용해 KV 캐시가 더 작고 토큰 희소 어텐션 커널 효율이 좋다. HySparse2의 self-decoder SWA 레이어는 64개 로터리 차원, 베이스 10,000의 부분 RoPE를 쓰고 전체·희소 어텐션은 NoPE를 쓰므로 장문맥 확장 시 RoPE 조정이 필요 없다. HySparse2는 128개의 강제 로컬 토큰과 1,024개의 전역 토큰을 토큰 단위로 선택하고, HySparse는 64토큰 블록으로 1,024개 전역 토큰을 선택하며 별도 128토큰 SWA 분기를 게이트 융합으로 결합한다. Hybrid SWA는 128토큰 슬라이딩 윈도우를 쓴다.

전제와 한계

사전학습 결과에서 HySparse2는 일반 능력은 두 기준선과 대체로 비슷하면서 장문맥에서 가장 뚜렷한 이득을 냈다. RULER와 NoLiMa에서 최고 점수를 기록해 HySparse 대비 각각 5.88점, 9.49점 올렸고, Repo Code PPL도 1.1570으로 HySparse의 1.1588, Hybrid SWA의 1.1578보다 낮았다. BBH와 MMLU-Pro에서는 강했지만 DROP에서는 HySparse가 우세하는 등 일반 과제는 혼합된 양상을 보였다. 같은 100B 토큰 후학습 이후에는 MRCR-v2, RULER-v2, AgentPPL, LongPPL 네 지표 모두에서 평가한 모든 길이에서 두 기준선을 앞섰다. 검색에서 이득이 가장 컸는데, HySparse 대비 평균 MRCR-v2와 RULER-v2가 각각 11.30점, 19.81점 올랐고 Hybrid SWA 대비로는 6.44점, 18.65점 올랐다. 256k에서 HySparse2의 RULER-v2는 58.45로 HySparse의 32.61, Hybrid SWA의 35.74와 큰 격차를 보였다. AgentPPL과 LongPPL도 전 구간에서 더 낮았는데, 문맥이 길어질수록 AgentPPL은 오르고 LongPPL은 내려가는 대비가 나타난다. 이는 긴 문맥이 LongPPL의 핵심 토큰에는 관련 정보를 더 노출하지만, 다회차 궤적에서는 중간 턴과 도구 출력이 늘어 근거를 찾기 어렵게 만들기 때문으로 해석된다.

효율 수치는 1k에서 1M 토큰까지 FP8 KV 캐시 기준으로 측정했다. 1M 토큰에서 HySparse2는 prefill FLOPs를 HySparse 대비 2.92배, Hybrid SWA 대비 5.02배 줄였고, KV 캐시는 2.69GB로 HySparse의 6.72GB, Hybrid SWA의 12.09GB보다 작다. 배포 관점에서는 prefill–decode 분리 구조에서 prefill 노드가 self-decoder와 KV Bridging 프로젝션만 호스팅하면 되는데, 49레이어 모델 기준 앞쪽 25개 레이어만 배치하면 되므로 prefill 노드 메모리가 거의 절반으로 줄어든다. cross-decoder 전체 어텐션 KV 캐시는 prefill 노드에서 계산해 디코드 노드로 전송하는데, 프로젝션된 KV 캐시가 원본 은닉 상태보다 작기 때문이다. 추측 디코딩과 관련해서는 사전학습 시 self/cross 경계의 은닉 상태에 조건화된 단일 MTP 레이어를 써서 수렴을 돕고, 후학습에서는 이를 같은 은닉 상태에 조건화된 더 큰 DFlash 스타일 드래프터로 교체할 수 있다고 언급한다.

절제 실험은 세 갈래다. 선택 입자도 비교에서는 블록 크기 64의 블록 단위와 토큰 단위가 같은 백본 배치, 1,024개 전역 토큰, 별도 128토큰 로컬 윈도우를 공유하도록 맞췄고, 토큰 단위가 같은 어텐션 예산에서 RULER-v2를 6.57점, 2-needle MRCR-v2를 8.14점, GraphWalks를 5.55점 올렸다. 에이전트 궤적은 추론·도구 호출·관측이 반복적으로 교차하고 역할 구분자 같은 특수 토큰이 경계를 표시하는데, 블록 단위는 토큰 하나를 유지하려고 이웃까지 예산을 쓰는 반면 토큰 단위는 문맥 전반의 개별 위치에 예산을 배분할 수 있다는 설명이 붙는다. 로컬 윈도우 절제에서는 별도 게이트 128토큰 SWA 분기, 분기 제거, 최근 128토큰을 희소 선택에 강제하는 방식 세 가지를 비교했는데, 분기를 제거하면 성능이 전반적으로 떨어졌고 강제 방식은 여러 과제에서 게이트 방식과 경쟁력이 있었다. 강제 방식은 RULER, RULER-v2, GraphWalks에서 최고였지만 게이트 방식 대비 GSM8K가 5.08점, MRCR-v2가 4.99점 낮고 LongPPL이 1.50% 높았다. 저자들은 추가 프로젝션 파라미터와 로컬 KV 캐시가 필요 없어 조기 종료 prefill이 가능해진다는 점에서 이 맞바꿈을 수용 가능하다고 본다. KV Bridging 절제는 290B-A8B 모델을 약 1.8T 토큰, 32k 문맥으로 학습해 검증했는데, KV Bridging을 쓴 쪽이 MMLU와 TriviaQA는 소폭 오르고 RULER는 0.31점 차이, Repo Code PPL은 거의 동일했으며 LongPPL은 3.6053에서 3.4202로 개선됐다. 반면 BBH와 GSM8K는 약 1점 하락하고 DROP은 71.37에서 68.17로 떨어졌다. 연결 방식으로는 KV Mirror와 비교했는데, KV Mirror가 앞뒤 레이어를 역순으로 짝짓는 U자형 연결을 쓰는 반면 KV Bridging은 self/cross 디코더의 전체 어텐션 레이어만 잇는다. 학습 대부분 구간에서 KV Bridging이 우세했고 최종 87.65 대 81.28로 앞섰다. 저자들은 전체 어텐션이 보이는 문맥 전체에 걸쳐 점수로 역전파되는 반면 SWA는 직접 연결을 국소 윈도우로 제한하기 때문에, 전체 레이어 상태가 더 나은 소스라고 추측한다.

개발자 관점에서 이 논문이 유용한 지점은 prefill 비용과 KV 캐시 메모리를 동시에 줄이면서 검색 정확도를 올리는 구체적 경로를 제시한다는 것이다. 특히 prefill 노드에 모델 절반만 올리는 분리 배포, 프로젝션된 KV 캐시를 디코드 노드로 전송하는 방식, 장문맥 확장 시 RoPE 조정이 필요 없는 NoPE 조합은 서빙 스택 설계에 바로 참고할 수 있다. 다만 도입 전에 확인할 것은 세 가지다. 토큰 단위 희소 선택이 실제 워크로드에서 커널 효율을 유지하는지, 강제 로컬 윈도우가 GSM8K·MRCR-v2에서 보인 하락이 자신의 과제에서 허용 가능한지, 그리고 KV Bridging이 DROP 같은 일부 과제에서 보인 하락이 서비스 품질에 영향을 주는지다.

저자들이 밝힌 전제와 한계도 분명하다. 전체 어텐션이 모델 품질에 여전히 중요하다고 보기 때문에 소수의 전체 어텐션 레이어를 유지하며, 그 비율을 작게 가져간다. 이 레이어들은 계산 비용이 여전히 비싸므로, 후학습 단계에서 경량 인덱서와 희소 어텐션으로 근사하는 방향이나 전체 어텐션의 쿼리 헤드를 줄이고 희소 어텐션의 쿼리 헤드를 늘리는 방향은 미래 과제로 남긴다. 추측 디코딩도 더 큰 드래프터 학습과 비동기 실행·검증 피드백 조율이 아직 미해결 과제라고 명시한다. 또한 강제 로컬 윈도우는 게이트 SWA 대비 일부 과제에서 성능을 내주는 맞바꿈이며, KV Bridging도 DROP에서 품질 하락이 관측된다는 점을 저자들이 직접 수치로 제시한다.