Gorio Tech Blog search

Language Models Can Control Their Own Attention 요약 설명

|

Contents

이번 글에서는 Language Models Can Control Their Own Attention 논문의 핵심 포인트만 간단히 정리한다.

  • 2026년 9월 2일(Arxiv)
  • Ho, Namgyu, Ahmad, Huzama, Koh, Woosung, Yun, Se-Young, Schuster, Tal, Santos, Cicero Nogueira dos.
  • KAIST AI, Google DeepMind
  • 논문 링크

영문판 보기


요약

  • Declarative Attention(DA)은 언어 모델이 생성 과정에서 파싱 가능한 attention 선언을 출력하는 zero-shot 추론 프로토콜이다. 15개 long-context task에서 DA는 Gemma-4-31B의 평균 attended token을 52.0%, Qwen-3.6-27B의 평균 attended token을 31.1% 줄였으며, 정확도 하락은 각각 1.27pp와 2.75pp였다. wall-clock 결과는 지연 시간 측정치가 아니라 roofline 추정치다.

1. Introduction

일반적인 long-context decoding은 각 autoregressive step에서 전체 KV cache를 읽는다. DA는 생성된 텍스트에서 mask를 도출하므로, 모델이 의도적으로 full-context global phase에 진입하는 경우를 제외하면 step마다 보조 선택 scan을 수행하지 않는다.

2. Declarative Attention

DA는 추론을 <global>, <focus>, <local> span으로 구성한다. global mode는 탐색을 위해 전체 context를 유지하고, focus mode는 정보 추출을 위해 이름으로 지정한 context segment를 유지하며, local mode는 scaffold와 생성된 응답을 보존하면서 context segment를 제거한다.

그림은 생성된 <global>, <focus>, <local> span이 scaffold와 생성된 응답을 유지하면서 점진적으로 더 좁은 segment-level mask를 만드는 과정을 보여준다. 예시에서 attended prompt token은 global mode의 25,466에서 focus mode의 3,435, local mode의 1,124로 감소한다.

Prompt 구조, 3가지 선언적 추론 mode, 이에 대응하는 segment-level KV attention mask
Prompt 구조, 3가지 선언적 추론 mode, 이에 대응하는 segment-level KV attention mask

2.1. Prompt overview

Prompt는 항상 보이는 scaffold와 선택적으로 보이는 context를 분리한다. scaffold에는 system instruction, question, DA instruction이 포함되어 있으므로, 어떤 masking mode에서도 task와 프로토콜을 계속 사용할 수 있다.

2.2. Context delivery

입력은 주소 지정이 가능한 “magic chunk”로 분할한다. 2048-token segment를 목표로 하며 paragraph, newline, sentence, clause, word 경계를 우선한다. chunk는 도구 사용을 모사한 transcript로 제시하지만, 추론 시점에는 도구를 실행하지 않는다.

2.3. Decode-time interventions

State machine은 기본 global mode에서 시작하고 <focus magic_chunks=”K”> 또는 <local>의 여는 tag가 끝나는 지점에서 전환한 뒤, 대응하는 닫는 tag 뒤에 global mode로 복귀한다. <global> tag는 전환을 일으키지 않고 prompt 구조를 나타낸다. vLLM hook은 KV-cache block 단위로 request의 block table을 다시 작성한다. global attention layer에는 mask를 적용하고, sliding-window attention 및 Gated DeltaNet layer는 변경하지 않는다.

3. When does DA pay off?

DA는 추가 decode step과 step당 더 적은 attended KV position을 교환한다. 유리한 조건은 global-attention read가 memory-bound이고 decode 비용을 지배하는 long-context, large-batch serving 환경이다. 논문은 이 trade-off를 분석하기 위해 roofline decomposition을 사용한다.

4. Experimental settings

실험은 Gemma와 Qwen 모델 6개를 평가하며, Gemma-4-31B와 Qwen-3.6-27B가 주요 모델이다. 15-source suite는 RULER, LongBench v1/v2, LooGLE, ZeroSCROLLS에서 가져왔다. 244K token 초과 context와 Gemma-4-E4B의 경우 116K token 초과 context는 제외하고, 출력은 8K token으로 제한했다.

표는 15개 평가 source의 domain, task 유형, QA provenance, context-length 분포를 정의한다. filtering 전 평균 source length가 1071.1K-token인 code repository를 포함하여, single-span retrieval/reasoning과 multi-span reasoning을 모두 포괄함을 보여준다.

Single-span retrieval/reasoning 및 multi-span reasoning task로 구분한 long-context benchmark source
Single-span retrieval/reasoning 및 multi-span reasoning task로 구분한 long-context benchmark source

5. Results

연구는 Vanilla full causal attention, DA-no-mask(DAnm), custom masking을 적용한 DA를 비교한다. 정확도는 생성된 rubric을 사용하는 LLM judge로 평가한다. attended token은 decoding 동안 읽은 KV position의 합이며, roofline wall-time 추정치는 deployment-scale decode 비용을 산정한다.

5.1. Main results

Gemma-4-31B에서 DA는 평균 정확도를 87.01%에서 85.74%로, sample당 attended token을 13.43M에서 6.45M으로 바꾼다. Qwen-3.6-27B에서는 정확도가 85.31%에서 82.56%로, attended token이 22.54M에서 15.52M으로 변한다. DAnm은 정확도에는 거의 영향을 주지 않지만 attended token을 늘리므로 masking의 효과를 분리해 보여준다.

표는 2개 주요 모델에서 Vanilla, DAnm, DA의 task별 정확도와 attended-token 총합, 그리고 macro-average를 보고한다. 집계된 attention 감소에도 task별 결과는 이질적이며, 더 긴 generation이 총 attended token을 증가시키는 사례도 있음을 보여준다.

15개 source에서 Vanilla, DA-no-mask, Declarative Attention의 정확도와 attended-token 결과
15개 source에서 Vanilla, DA-no-mask, Declarative Attention의 정확도와 attended-token 결과

정규화된 비교에서 DA는 Gemma-4-31B의 decode step을 35%, Qwen-3.6-27B의 decode step을 31% 늘리지만, attended token은 각각 vanilla의 48%와 69%만 사용한다. DAnm은 attended token을 vanilla의 166%와 129%까지 늘려 mask의 이점을 분리해 보여준다.

주요 모델에서 Vanilla, DA, DA-no-mask의 상대 정확도, decode step, attended token
주요 모델에서 Vanilla, DA, DA-no-mask의 상대 정확도, decode step, attended token

5.2. Model capability scaling

상대적인 DA 정확도는 2개 모델 계열 모두에서 규모에 따라 단조롭게 향상한다. Gemma-4-E4B에서 Gemma-4-31B까지 vanilla 정확도의 29%에서 99%로, Qwen-3.5-4B에서 Qwen-3.6-27B까지 64%에서 97%로 증가한다. Gemma-4-12B는 일부 응답이 8K generation budget을 소진하여 vanilla보다 attended token이 많으며, 이는 masking이 step당 attention을 늘렸기 때문은 아니다.

그림은 Gemma와 Qwen 계열 모두에서 모델 크기가 커질수록 DA의 상대 정확도가 향상함을 보여준다. 또한 Gemma-4-12B의 183% attended-token 결과는 non-termination 때문이며, generation limit에 도달한 응답을 제외하면 값이 98%로 낮아진다고 설명한다.

Gemma 4 및 Qwen 3.5/3.6 모델 크기별 DA의 상대 정확도와 attended token
Gemma 4 및 Qwen 3.5/3.6 모델 크기별 DA의 상대 정확도와 attended token

5.3. Context-length scaling

Gemma-4-31B에서 DA는 32K-token context까지 vanilla와 약 1 percentage point 이내의 정확도를 유지하고, 64–256K bin에서는 vanilla의 약 96%까지 하락한다. 절대 attended-token 절감량은 짧은 context의 약 1M에서 가장 긴 bin의 약 21M으로 증가한다.

Gemma-4-31B에서 DA는 32K token까지 vanilla에 가까운 정확도를 유지한 뒤 더 긴 context에서 하락한다. DAnm에서는 같은 하락이 나타나지 않는다. 절대 attended-token 절감량은 context 길이와 함께 증가하고, DAnm의 overhead도 점점 양의 방향으로 커진다.

Context-length bin별 vanilla 대비 상대 정확도와 절대 attended-token 변화
Context-length bin별 vanilla 대비 상대 정확도와 절대 attended-token 변화

5.4. Efficiency results

단일 B200에서 bf16을 사용하고 40% MFU 및 70% MBU를 가정하면, roofline model은 Gemma-4-31B에서 DA 192.3 ms와 vanilla 269.1 ms를 추정한다. Qwen-3.6-27B에서는 DA 237.3 ms와 vanilla 306.2 ms를 추정한다. 이 추정치는 prefill을 제외하며, 포화 상태의 최적화된 phase-disaggregated serving 조건을 모델링한다.

표는 모델링한 decode 비용을 compute, global-memory, local-memory 항으로 분해한다. DA는 global KV read를 줄이지만 더 긴 generation 때문에 compute와 고정 local-memory 작업이 늘어나며, B200 roofline 가정 아래 보고된 모델링 총비용이 산출된다.

B200 accelerator에서 vanilla와 Declarative Attention의 roofline 추정 응답당 decode 비용
B200 accelerator에서 vanilla와 Declarative Attention의 roofline 추정 응답당 decode 비용

6. Analysis

분석은 mode 사용으로 인한 절감과 선언 프로토콜을 따르지 못해 발생하는 실패를 분리한다. 특히 작은 backbone에서는 global-mode 사용과 유효하지 않은 focus reference가 실질적인 한계로 확인된다.

6.1. Efficiency impact of DA modes

Gemma-4-31B에서 global mode는 전체 생성 token의 약 27%를 차지하며, focus와 local mode가 나머지 73%를 차지한다. focus와 local mode는 vanilla step 대비 생성 token당 약 76–99%의 attention을 절감한다. 다만 가장 긴 context bucket에서는 global-mode 사용 비율이 약 45%까지 상승한다.

그림은 Gemma-4-31B의 generation을 context bucket별 global, focus, local mode로 분해한다. focus와 local mode가 절감의 원천임을 보이는 동시에, 가장 긴 context에서 global-mode 사용이 늘어나면서 총 절감량이 제한됨을 보여준다.

Gemma-4-31B의 context 길이별 mode-token 비중과 token당 attention 절감량
Gemma-4-31B의 context 길이별 mode-token 비중과 token당 attention 절감량

6.2. Protocol adherence

Focus-reference 성공률은 Gemma-4-E4B에서 58%, Gemma-4-31B에서 99%로 상승하며, Qwen-3.5-4B에서 89%, Qwen-3.6-27B에서 99%로 상승한다. Focus 시도 횟수는 응답당 약 1.4–1.9의 좁은 범위에 머무르므로, 큰 모델은 focus 호출을 줄이기보다 주로 reference 유효성을 개선한다.

그림은 모델 크기와 유효한 focus-reference 해석률, focus-call 빈도의 관계를 보여준다. 특히 Gemma에서 58%에서 99%로 상승하는 focus 성공률은 프로토콜 준수가 작은 모델의 주요 병목이라는 결론을 뒷받침한다.

모델 크기별 focus-reference 성공률과 응답당 focus 시도 횟수
모델 크기별 focus-reference 성공률과 응답당 focus 시도 횟수

7. Conclusion

논문은 모델이 생성한 제어 텍스트가 parameter update 없이 가역적인 KV-cache masking을 구동할 수 있다고 결론짓는다. 저자들은 모델 능력이 높아질수록 정확도 격차가 줄고 context 길이가 늘수록 절대 절감량이 커지므로, zero-shot 결과를 하한으로 해석한다.

8. Discussion, Limitations, and Future Work

DA는 zero-shot elicitation, static benchmark의 인위적인 segmentation, non-thinking-mode 평가, 비용이 큰 global navigation step이라는 한계를 가진다. 근거가 segment 사이에 파괴적으로 분산된 task나 출력 길이가 문서와 함께 증가하는 task에서는 정확도가 낮아지거나 end-to-end 절감 효과가 사라질 수 있다.

8.1. Relevance of DA in the present and future

저자들은 global attention이 decode-time 비용에서 큰 비중을 차지하는 환경에서는 DA가 여전히 유효하다고 주장한다. 1M-token roofline 계산은 Kimi-K3의 global-attention 비중을 94%로, 여러 indexer-based sparse architecture의 비중을 56–97%로 추정한다. 다만 이는 직접 측정치가 아니라 architecture별 모델링 값이다.

8.2. Further potential of DA with post-training and in agentic settings

Post-training은 DA의 추가 decode step을 줄이고 mode-selection policy를 개선할 수 있다. Agentic context는 tool output이나 dialogue turn처럼 자연스러운 segment를 제공할 수 있지만, 이 용도는 평가하지 않았다. 테스트한 모델이 thinking trace 내부에서 DA syntax를 안정적으로 따르지 못했기 때문에 thinking mode는 비활성화했다.

8.3. Synergy with modern efficient inference techniques

DA는 남은 global phase를 처리하는 lightweight-scan sparse attention 및 순차 decoding overhead를 줄이는 speculative decoding과 결합할 수 있다. 논문은 block-table masking이 이 방법들과 호환된다고 주장하지만, 결합 실험은 제공하지 않는다.

8.4. DA as system-2 sparse attention

저자들은 모델이 latent activation에서 선택을 유도하는 데 그치지 않고 attended region을 명시적으로 추론하고 선언한다는 점에서 DA를 “system-2 sparse attention”으로 규정한다. 선언된 focus span을 기반으로 가역적인 KV-cache offloading과 prefetching도 제안하지만, 이는 향후 system 방향으로 제시한다.

부록

  • Appendix A는 DA를 dynamic sparse attention, KV-cache eviction, attention-steering 방법과 비교한다. DA는 resident block을 영구적으로 evict하지 않고 가역적으로 mask한다. 부록에는 vLLM 통합, roofline 유도, dataset, judge validation, sampling, prompt 구성, 추가 결과와 segmentation에 민감한 근거 또는 문서 길이에 비례하는 출력을 포함하는 failure case도 문서화되어 있다.

짧은 생각

DA는 생성된 추론 구조와 추론 시점 memory access를 연결하는 구체적인 interface를 제공하며, DAnm ablation은 chunked prompting이 아니라 masking이 주요 attention 절감을 만든다는 점을 보여준다. 실용성은 강한 프로토콜 준수, 관리 가능한 global-mode 사용, 적절한 task decomposition, roofline 가정에 가까운 deployment 조건에 달려 있다.