[논문리뷰] Language Models Can Control Their Own Attention

링크: 논문 PDF로 바로 열기

I have browsed the paper. Now I will proceed with summarizing it according to the specified format and constraints.

Part 1: Summary Body

Authors: Namgyu Ho, Huzama Ahmad, Woosung Koh, Se-Young Yun, Tal Schuster, Cicero Nogueira dos Santos 키워드: Declarative Attention, Long-Context LLMs, Sparse Attention, KV Cache Efficiency, Zero-Shot Prompting, Inference Optimization, Attention Masking

1. Key Terms & Definitions (핵심 용어 및 정의)

• Declarative Attention (DA): LLM이 추론 과정(Chain-of-Thought)에서 다음으로 집중할 컨텍스트 영역을 명시적으로 선언하도록 유도하는 프로토콜입니다. 이를 통해 추론 엔진이 모델의 선언에 따라 동적으로 Attention Mask를 생성하여 KV Cache 접근을 최적화합니다.
• KV Cache: Transformer 모델의 Inference 과정에서 Attention 계산에 필요한 Key와 Value를 저장하는 메모리 캐시입니다. Long-context 시나리오에서 이 캐시의 접근 비용이 Inference Latency의 주요 원인이 됩니다.
• Attention Mask: Attention 계산 시 특정 토큰 쌍 간의 상호작용을 제어하거나 차단하는 데 사용되는 이진 마스크입니다. DA에서는 모델의 선언에 따라 KV Cache의 특정 부분을 Masking하여 메모리 접근을 줄입니다.
• Global Mode (): DA 프로토콜의 세 가지 Attention Mode 중 하나로, 모델이 전체 컨텍스트 세그먼트를 모두 참조할 수 있는 모드입니다. 주로 관련 정보를 찾기 위한 탐색(Navigation) 단계에서 사용됩니다.
• Focus Mode (): DA 프로토콜의 세 가지 Attention Mode 중 하나로, 모델이 특정 명명된 컨텍스트 세그먼트(Magic Chunk)에만 집중하는 모드입니다. 특정 정보 추출 또는 해당 영역에 대한 심층 추론에 활용됩니다.
• Local Mode (): DA 프로토콜의 세 가지 Attention Mode 중 하나로, 모델이 컨텍스트 세그먼트를 전혀 참조하지 않고, 지금까지 생성된 자신의 Output과 이미 추출된 정보에만 Attention하는 모드입니다. 주로 자체적인 계획(Planning)이나 정보 통합(Synthesis)에 사용됩니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

Long-context Language Model (LLM) Inference 과정에서 Transformer 아키텍처는 매 Decoding Step마다 모든 이전 토큰에 대해 Attention을 계산하므로, KV Cache 메모리 접근 Latency가 전체 Decoding Time의 상당 부분을 차지하며 이는 Long-context Regime에서 특히 심화됩니다. 예를 들어, Qwen-3.5-397B-A17B 모델이 1M 토큰 컨텍스트를 처리할 때, 매 Decoding Step마다 약 15GB의 KV Cache를 로드해야 합니다. 이는 모델의 Active Parameter인 17B를 로드하는 것과 유사한 메모리 대역폭 요구 사항입니다. 기존 연구들은 Attention-heavy 토큰을 사전 예측하여 Masking하는 접근 방식을 사용하지만, 이러한 Extrinsic Scoring 방식은 여전히 Step당 O(N)의 비용을 발생시킵니다. 본 논문은 모델이 어떤 컨텍스트 부분이 관련 있는지 스스로 알고 있을 것이라는 직관에서 출발하여, 외부적인 Scoring 없이 모델 자체가 Attention 범위를 명시적으로 선언하도록 유도하는 Declarative Attention (DA) 프로토콜을 제안합니다 [Figure 1].

Figure 1 — DA 프로토콜 아키텍처

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 Declarative Attention (DA)이라는 프로토콜을 제안하며, 이는 LLM이 Chain-of-Thought 내에서 스스로 Attention 범위를 선언하도록 유도하여 Inference Engine이 이를 파싱하여 동적으로 Attention Mask를 구성합니다 [Figure 1]. DA는 Generation을 세 가지 Attention Mode로 나눕니다: (전체 컨텍스트), (특정 지역), (최근 Output만). Inference Engine은 이 선언을 Tool Call처럼 해석하여 KV Cache Read의 대부분을 건너뛸 수 있습니다. 이러한 접근 방식은 보조 Scorer 없이 Off-the-shelf 모델에서 Zero-shot으로 작동하며, Inference Engine의 Decode-time Intervention을 통해 vLLM에 통합되었습니다. DA는 Attention Mask를 Block Granularity로 적용하여 FlashAttention과 같은 기존 커널의 효율성을 활용하고, KV Cache Block Table을 재작성하여 Keep-block만 Attention Kernel에 보이도록 합니다.

Zero-shot 평가 결과, DA는 15개의 Long-context Task에서 Vanilla Baseline 대비 상당한 효율성 향상을 보였습니다. Gemma-4-31B 모델에서는 평균 Decoding Attention Cost를 52.0% 감소시켰고, Qwen-3.6-27B 모델에서는 31.1% 감소시켰습니다 [Table 2]. Accuracy 손실은 Gemma-4-31B에서 1.27pp, Qwen-3.6-27B에서 2.75pp로 비교적 적었으며, 이는 모델 Scale이 커질수록 Accuracy Gap이 줄어드는 경향을 보였습니다 [Figure 3]. 특히, Attention Mask의 역할이 중요하며, DA-no-mask (DA-nm) Baseline과 비교했을 때 Masking이 Gemma에서 71.1%, Qwen에서 46.5%의 Attended Token 절감을 가져왔습니다. 또한, Context Length가 길어질수록 절대적인 Token Saving은 더욱 증가하여, 가장 긴 Context에서 최대 21M 토큰까지 절약되었습니다 [Figure 4]. Wall-clock Time 분석 결과, DA는 최적화된 Inference 환경에서 Gemma-4-31B의 Decode Wall-clock Time을 Vanilla 대비 0.71x, Qwen-3.6-27B는 0.77x로 줄일 수 있음을 이론적으로 제시했습니다 [Table 3].

4. Conclusion & Impact (결론 및 시사점)

본 연구는 Off-the-shelf LLM이 Chain-of-Thought의 일부로 스스로 Attention 범위를 선언하도록 하는 Zero-shot 프로토콜인 Declarative Attention (DA)을 소개합니다. 이 프로토콜을 통해 Inference Engine은 모델의 Output Text에서 Attention Mask를 직접 읽어들여 기존 Activation 기반의 근사치 방식보다 효율적인 KV Cache 처리를 가능하게 합니다. 15개의 Long-context Task에 걸친 실험에서, DA는 Vanilla 모델에 근접한 Accuracy를 유지하면서 Global Attention KV Read에서 상당한 Attention Cost 절감(평균 52.0% on Gemma-4-31B, 31.1% on Qwen-3.6-27B)을 달성했습니다.

이 연구는 Selective Attention을 네트워크 내부에서 추론되는 패턴이 아닌, 모델이 명시적으로 선언하는 형태로 전환함으로써 Long-context Inference의 가독성과 효율성을 동시에 향상시킬 수 있음을 보여줍니다. DA는 Zero-shot Prompting만으로 구현되었기 때문에, 향후 Fine-tuning이나 Post-training을 통해 모델의 Protocol Adherence와 전략적 Mode 사용을 개선할 경우 더욱 큰 성능 향상이 기대됩니다. 특히, Agentic System이나 Reversible Context Compaction과 같은 응용 분야에서 DA는 Context Management의 새로운 축을 제공하며, LLM의 Long-horizon Reasoning 및 복잡한 Task 처리 능력을 크게 발전시킬 잠재력을 가집니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.