Gorio Tech Blog search

Language Models Can Control Their Own Attention 요약 설명

|

목차

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

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

요약

  • 장문 문맥의 전역 어텐션은 실제로 중요한 일부 토큰에만 어텐션 질량이 집중되어도, 응답 토큰마다 전체 KV 캐시를 읽어야 한다. 이 논문은 모델이 chain-of-thought 안에서 읽을 범위를 선언하게 하고, 추론 엔진이 그 선언을 KV-cache 마스크로 변환하는 Declarative Attention(DA)을 제안한다. <global>은 전체 문맥을 탐색하고, <focus>는 지정한 문맥 구간에서 사실을 추출하며, <local>은 문맥 구간을 보지 않고 앞선 응답을 바탕으로 추론한다. 별도 학습이나 파라미터 갱신 없이, 고정된 과제 불문 프롬프트를 사후학습된 모델에 적용한 제로샷 방법이다. 15개 장문 문맥 원천의 동등 가중 평균에서 Gemma-4-31B는 Vanilla 정확도 87.01%에서 DA 85.74%로 1.27퍼센트포인트 낮아졌고, 응답당 누적 attended tokens는 13.43M에서 6.45M으로 6.98M, 즉 52.0% 감소했다. Qwen-3.6-27B는 정확도 85.31%에서 82.56%로 2.75퍼센트포인트 낮아졌으며, 누적 attended tokens는 22.54M에서 15.52M으로 7.02M, 즉 31.1% 감소했다. 이 비교는 해당 두 모델과 15개 원천 평균에 한정되며, 저자들은 모델 규모가 커질수록 Vanilla와 DA의 정확도 격차가 줄어든다고 보고한다.
    • 기존 동적 희소 어텐션은 경량 프록시 점수로 매 디코드 단계 관련 KV를 찾지만, 선택을 위해 문맥 전체를 훑는 단계별 O(N) 비용이 남는다. DA는 생성된 선언에서 마스크를 직접 읽어 이 선택 스캔을 피하려는 접근이다.
    • 누적 attended tokens는 디코드 전 단계에서 읽은 KV 위치 수를 응답별로 합산한 지표다. 이는 종단간 지연시간의 실측값과 같지 않으며, 지연시간은 별도 roofline 추정으로 분석한다.

2. Declarative Attention

  • DA 프롬프트는 항상 보이는 scaffold와 가시 범위가 바뀌는 장문 context로 구성된다. scaffold는 attention sink를 차지하도록 둔 짧은 시스템 지시, 질문, DA 사용 지시, 그리고 지금까지 생성한 응답으로 이루어지며 모든 모드에서 유지된다. context는 약 2,048토큰을 목표로 magic chunk로 나뉜다. 분할기는 문단, 줄바꿈, 문장 끝, 절 경계, 단어 경계 순으로 가장 거친 가능한 경계를 택해 단어 중간을 자르지 않는다. 공백 없는 긴 문자열은 자르지 않고 상한을 넘는 단일 구간으로 남기며, 전체 구간을 연결하면 원래 입력이 복원되는 손실 없는 분할이다. 각 chunk는 실제 호출하지 않는 get_magic_chunk 도구 대화 형식으로 제공해 모델이 사후학습에서 익힌 메시지 경계를 활용하게 한다.
    • <global>은 모든 context chunk를 읽고 다음에 살필 chunk와 그 이유를 찾는 탐색 모드다. <focus magic_chunks=”K”>는 지정한 하나 이상의 chunk, scaffold, 기존 응답만 읽고 필요한 값을 원문대로 추출하도록 유도한다. <local>은 context chunk를 읽지 않고, 이미 추출한 값과 질문·응답을 사용해 계획하거나 답을 종합한다.
    • 프로토콜 자체는 모드의 순서와 횟수를 제한하지 않는다. 다만 제로샷 프롬프트는 적어도 한 번의 focus 사용, 최종 local 단계에서 답에 넣을 값을 명시하는 것, local에서 보지 않은 chunk 내용을 추정하지 않는 것을 권고한다.
    • DA에는 별도의 목적 함수, gradient 경로, 학습 데이터, 추가 학습이 없다. SFT나 강화학습으로 정확도와 어텐션 효율을 함께 최적화하는 것은 저자들이 향후 방향으로 제안한 내용이다.
세 가지 선언 모드와 구간별 어텐션 마스크의 동작 예시
세 가지 선언 모드와 구간별 어텐션 마스크의 동작 예시

2.3. Decode-time interventions

  • 디코드 시 DA 상태 기계는 기본 상태를 global로 두고 생성 스트림의 태그를 파싱한다. <focus …> 또는 <local>의 여는 태그가 닫히는 문자 `>`에서 제한 마스크로 전환하고, 대응하는 닫는 태그 뒤에는 global로 되돌아간다. <global> 태그는 기본 상태인 global을 다시 선언해 모델의 추론 구조를 돕는 표기일 뿐 상태를 전환하지 않는다. 즉 모델이 읽을 범위를 텍스트로 먼저 선언하면 런타임이 이를 즉시 KV 접근 범위에 반영한다.
    • vLLM의 KV 캐시는 일반적으로 16~32토큰 고정 블록으로 저장되고 커널도 블록 전체를 읽으므로, DA는 토큰이 아니라 블록 단위로 마스크를 구성한다. 유지할 span의 경계는 바깥 블록 경계로 확장하며, 각 경계에서 발생하는 추가 비용은 최대 한 블록이다.
    • 구현은 attention metadata builder의 hook에서 요청별 KV block table을 갱신해 유지 블록만 기존 커널에 노출한다. FlashAttention 등의 커널과 스케줄러는 수정하지 않는다. Qwen의 전역 어텐션 계층과 Gemma의 전역 어텐션 계층에만 적용하며, 문맥 길이에 무관한 Gated DeltaNet 및 sliding-window attention 계층은 그대로 둔다.
    • DA는 추가 선언과 추론으로 디코드 단계 수를 늘릴 수 있는 대신, 전역 어텐션 계층의 단계당 KV 읽기를 줄인다. 따라서 실제 이득은 문맥 길이, 생성 길이, 전역 어텐션이 전체 디코드 비용에서 차지하는 비중에 좌우된다.

4. Experimental settings

  • 평가는 Gemma-4-{31B, 12B, E4B}와 Qwen-3.6-27B, Qwen-3.5-{9B, 4B}의 두 계열 6개 모델로 구성하며, 주 결과는 Gemma-4-31B와 Qwen-3.6-27B에서 제시한다. Gemma-4-E4B의 입력 한도는 128K토큰이고 나머지는 256K토큰이다. 15개 문맥 원천은 단일 구간 검색·추론과 다중 구간 추론으로 분류했으며, RULER, LongBench v1/v2, LooGLE, ZeroSCROLLS에서 가져왔다. 원본 QA 11개 원천과 Gemini-3-Flash가 생성한 합성 QA 4개 원천을 사용했다. 8K 생성 토큰과 DA 프롬프트용 4K토큰을 남기기 위해 244K토큰 초과 사례를 제외했고, Gemma-4-E4B에는 116K 한도를 적용했다. 각 원천에서 고정 시드로 최대 128개 예제를 뽑았으며 같은 모델의 세 방법에는 동일 표본을 사용했다.
    • Vanilla는 원문 context와 질문을 인라인으로 제공하고 전체 causal attention을 사용한다. DAnm은 magic chunk와 세 모드 지시를 포함한 DA 프롬프트를 쓰되 전체 causal attention을 유지한다. 따라서 Vanilla 대 DAnm은 chunked 프롬프트 형식의 효과를, DAnm 대 DA는 맞춤형 마스크 효과를 분리한다.
    • 정확도는 정답 조건부 엄격 채점 rubric을 Gemini-3-Flash로 만들고, thinking-enabled Qwen-3.5-4B가 판정한 결과다. 형식 실패도 오답으로 처리한다. 부록의 2,993개 응답 검증 표본에서 이 판정기는 Gemini-3.1-Pro와 개별 판정 98.53% 일치했고 셀별 정확도는 Pearson r=0.992였지만, 본문 정확도는 여전히 LLM 판정 절차에 의존한다.
    • attended tokens는 디코드 전체에서 읽은 KV 위치 합을 응답별로 집계한 뒤 원천별로 평균하고 15개 원천을 동등 가중한 값이다. 모든 모델은 수정한 vLLM으로 NVIDIA B200에서 제공했으며, thinking mode는 태그 안에서 DA 프로토콜을 따르지 못해 비활성화했다. 최대 생성 길이는 8K토큰이고, 샘플링은 모델 계열별 권장 설정을 사용했다.
장문 문맥 벤치마크 원천과 과제 범주의 구성
장문 문맥 벤치마크 원천과 과제 범주의 구성

5.1. Main results

  • Table 2와 Figure 2의 15개 원천 평균에서 DA는 Vanilla보다 낮은 누적 어텐션 비용과 정확도 저하의 절충을 보였다. Gemma-4-31B의 정확도는 Vanilla·DAnm·DA 순서로 87.01%, 87.01%, 85.74%였고, attended tokens는 13.43M, 22.31M, 6.45M/sample이었다. DA의 Vanilla 대비 절감량은 13.43M−6.45M=6.98M/sample이며, 6.98/13.43×100=52.0%다. Qwen-3.6-27B의 정확도는 85.31%, 84.62%, 82.56%, attended tokens는 22.54M, 29.02M, 15.52M/sample이었다. DA는 22.54M−15.52M=7.02M/sample, 7.02/22.54×100=31.1%를 덜 읽었다. 정확도 차이는 각각 −1.27, −2.75퍼센트포인트다.
    • DAnm은 Gemma에서 Vanilla보다 attended tokens가 66.2% 많고, Qwen에서 28.8% 많다. DA는 DAnm 대비 Gemma에서 (22.31−6.45)/22.31×100=71.1%, Qwen에서 (29.02−15.52)/29.02×100=46.5%를 줄인다. 같은 DA 프롬프트에서 마스크만 달리한 이 비교를 근거로 저자들은 누적 토큰 절감의 주된 원인이 chunk 형식보다 마스크라고 해석한다.
    • 정확도 손실도 주로 DAnm과 DA의 차이에서 나타난다. DAnm은 Gemma에서 Vanilla와 같은 87.01%이고, Qwen에서는 Vanilla보다 0.69퍼센트포인트 낮다. DA는 DAnm보다 Gemma에서 1.27퍼센트포인트, Qwen에서 2.06퍼센트포인트 낮다. 이는 동일 프롬프트 아래 마스크가 효율 이득과 정확도 비용의 대부분을 만든다는 해당 비교의 결과다.
    • 과제별 결과는 균일하지 않다. DA는 Gemma에서 15개 중 7개, Qwen에서 5개 원천에서 Vanilla와 같거나 높은 정확도를 냈다. Gemma의 longdep_qa는 65.6%에서 68.8%로 3.2퍼센트포인트, Qwen의 code_repo는 72.2%에서 77.8%로 5.6퍼센트포인트 높았다. 다중 구간 추론 평균의 정확도 하락은 Gemma 2.28퍼센트포인트, Qwen 3.59퍼센트포인트로 단일 구간의 0.78, 2.34퍼센트포인트보다 컸다. 원문 5.1은 Qwen에서 DA의 attended tokens가 Vanilla보다 증가한 원천이 5개라고 주장하지만, Table 2에서 직접 확인되는 것은 niah_multikey_3(4.29→4.48M), qmsum(1.77→2.67M), multidoc_qa(16.17→20.89M), singledoc_qa(12.39→18.00M)의 4개다. 따라서 원문 서술과 표는 충돌하며, 여기서는 원천별 수치를 직접 제시한 Table 2를 채택한다.
두 대형 모델에서 Vanilla·마스크 없는 DA·DA의 정확도와 누적 어텐션 토큰 비교
두 대형 모델에서 Vanilla·마스크 없는 DA·DA의 정확도와 누적 어텐션 토큰 비교
주요 모델에서 정확도·디코드 단계·누적 어텐션 토큰의 상대 비교
주요 모델에서 정확도·디코드 단계·누적 어텐션 토큰의 상대 비교

5.2. Model capability scaling

  • Figure 3a에서 DA 정확도를 각 모델의 Vanilla 정확도로 나눈 상대 정확도는 모델 규모와 함께 두 계열에서 단조롭게 높아졌다. Gemma 계열은 E4B 29%, 12B 91%, 31B 99%였고, Qwen 계열은 4B 64%, 9B 89%, 27B 97%였다. 저자들은 DA가 태그·chunk 참조·모드 순서를 지키는 프로토콜 수행과 제한된 정보에서의 답변 생성을 함께 요구하므로, 일반 모델 역량이 높을수록 Vanilla와의 차이가 줄어든다고 해석한다.
    • Gemma-4-E4B의 낮은 상대 정확도에는 추론 실패뿐 아니라 태그 형식 실패가 관여한다. 유효 chunk 참조로 파싱된 focus 호출 비율은 58%로, 가장 큰 두 모델의 99%보다 낮다. 저자들은 작은 모델의 성능 저하 중 상당 부분을 프로토콜 미준수와 연결한다.
    • 누적 attended tokens는 6개 모델 중 5개에서 Vanilla보다 낮았고, Gemma 계열은 Vanilla의 46~48%, Qwen 계열은 50~69%를 읽었다. Gemma-4-12B만 183%였는데, DA 응답 약 6%가 8K 생성 한도 안에 종료하지 못한 생성 길이 효과다. 이 응답을 양쪽 arm에서 제외한 Figure 3b 값은 98%다.
    • 저자들은 마스크의 단계당 절감은 대체로 모델 규모와 무관하다고 보며, 누적 비용의 차이는 global·focus·local 모드의 사용 비율과 생성 길이가 함께 결정한다고 구분한다.
모델 크기에 따른 상대 정확도와 누적 어텐션 토큰 변화
모델 크기에 따른 상대 정확도와 누적 어텐션 토큰 변화

5.3. Context-length scaling

  • Gemma-4-31B에서 15개 원천을 문맥 길이 구간으로 묶은 Figure 4a는 DA의 상대 정확도가 32K토큰까지 Vanilla에서 약 1퍼센트포인트 이내에 머물다가, 64K~256K 구간에서 Vanilla의 약 96%까지 완만하게 내려감을 보인다. 같은 DA 프롬프트에서 마스크만 제거한 DAnm에는 이 하락이 나타나지 않았으므로, 저자들은 이 조건에서 장문 문맥 정확도 비용이 chunk 형식이 아니라 attention mask에서 비롯된다고 해석한다.
    • Figure 4b에서 Gemma-4-31B의 Vanilla 대비 DA attended-token 절감은 짧은 문맥 구간에서 약 1M/sample, 가장 긴 구간에서 약 21M/sample이다. 단계당 절감 비율이 대체로 일정하므로 문맥이 길어질수록 절대 절감량이 커진다는 결과다.
    • 이 절의 값은 Gemma-4-31B에서 15개 원천을 풀링하고 길이별로 구간화한 결과다. Qwen-3.6-27B의 대응 결과는 부록에 있으므로, Gemma 수치를 두 모델의 공통된 정확한 크기로 일반화할 수 없다.
문맥 길이별 상대 정확도와 절대 누적 어텐션 토큰 절감
문맥 길이별 상대 정확도와 절대 누적 어텐션 토큰 절감

5.4. Efficiency results

  • 저자들은 실측 지연시간 대신, 최적화된 대규모 배치·분리형 서빙을 가정한 roofline 기반 디코드 시간 추정을 제시한다. 단일 NVIDIA B200, bf16, 행렬곱 MFU 40%, 메모리 읽기 MBU 70%에서 응답 전체의 행렬곱·전역 메모리·로컬 메모리 비용을 합산한다. 전역 메모리는 문맥 길이에 비례하는 전역 어텐션 KV 읽기이며 DA가 줄이는 유일한 항이다. 로컬 메모리는 Gemma의 sliding-window attention KV 또는 Qwen의 Gated DeltaNet 상태처럼 문맥 길이에 무관한 읽기라서 DA가 줄이지 못한다.
    • Table 3에서 Gemma-4-31B의 Vanilla 총 추정 디코드 시간은 269.1ms, DA는 192.3ms로 192.3/269.1=0.71배다. 전역 메모리 읽기는 1,100GB·196.5ms에서 528GB·94.4ms로 줄었지만, 생성 단계 증가로 행렬곱은 22.9ms에서 30.9ms, 로컬 메모리는 49.7ms에서 67.1ms로 늘었다.
    • Qwen-3.6-27B는 Vanilla 306.2ms, DA 237.3ms로 237.3/306.2=0.77배다. 전역 메모리는 1,477GB·263.8ms에서 1,017GB·181.6ms로 감소하고, 행렬곱은 34.4ms에서 45.1ms, 로컬 메모리는 8.0ms에서 10.5ms로 증가한다. Vanilla의 전역 메모리 항 비중은 Gemma 73%, Qwen 86%다.
    • 이 수치는 명시한 활용률에서의 roofline 이론 추정이며 측정된 wall-clock 시간이 아니다. prefill은 별도 pool에서 수행한다고 가정해 제외했다. Gemma는 60개 층 중 50개가 SWA여서 줄일 수 없는 로컬 비용 바닥이 크므로, 전역 읽기 감소율이 더 크더라도 전체 시간 감소율이 비례해서 커지지 않는다.
최적화된 분리형 추론 조건에서의 roofline 기반 디코드 시간 분해
최적화된 분리형 추론 조건에서의 roofline 기반 디코드 시간 분해

6.1. Efficiency impact of DA modes

  • Gemma-4-31B의 Figure 5 분석에서 생성 토큰 중 global은 약 27%, focus와 local은 합계 약 73%다. global은 전체 문맥을 읽으므로 단계당 절감이 0%인 반면, focus와 local은 Vanilla의 같은 단계 대비 문맥 길이 구간별로 각각 76~98%, 88~99%의 어텐션 읽기를 절감한다. 따라서 전체 효율은 값 추출과 응답 내 추론을 제한된 범위의 모드에 얼마나 배치하는지에 의존한다.
    • Figure 5a에서 global 비율은 <32K, 32~64K, 64~128K, >128K 구간에서 각각 27%, 26%, 23%, 45%다. 가장 긴 문맥 구간의 global 비중 증가는 전체 절감률을 제한한다.
    • focus와 local의 단계당 절감률은 문맥이 길수록 커진다. 다만 저자들은 모드 사용 비율이 제로샷 프롬프트에 대한 모델 반응이므로, DA 메커니즘의 고정된 한계로 해석해서는 안 된다고 설명한다.
문맥 길이별 모드 토큰 비중과 focus·local 단계당 어텐션 절감률
문맥 길이별 모드 토큰 비중과 focus·local 단계당 어텐션 절감률

6.2. Protocol adherence

  • DA는 <focus> 태그가 유효한 magic chunk 번호로 파싱되어야 하므로 프로토콜 준수율을 별도로 측정한다. Figure 6a의 focus 성공률은 Gemma-4-E4B·12B·31B에서 각각 58%, 90%, 99%이며, Qwen-3.5-4B·9B·Qwen-3.6-27B에서는 89%, 95%, 99%다. 큰 모델은 거의 완전한 수준의 유효 chunk 참조를 생성하지만, 가장 작은 Gemma 모델의 낮은 준수율은 Figure 3의 낮은 상대 정확도와 함께 나타난다.
    • 응답당 focus 시도 횟수는 Gemma에서 1.37, 1.82, 1.57회, Qwen에서 1.59, 1.85, 1.50회로 약 1.4~1.9 범위다. 모델 크기에 따른 뚜렷한 추세는 없다. 저자들은 큰 모델의 이점이 focus 호출 수 감소보다 생성한 호출을 유효하게 파싱하는 능력에 있다고 해석한다.
    • 유효 태그 생성 능력은 답변 정확도와 같지 않지만, 생성 텍스트를 런타임 제어 신호로 쓰는 DA에서는 독립적인 병목이다. 준수율과 시도 횟수는 모두 제로샷 프로토콜 사용 방식에서 관측한 값이다.
모델 크기별 focus 참조 파싱 성공률과 응답당 focus 시도 횟수
모델 크기별 focus 참조 파싱 성공률과 응답당 focus 시도 횟수

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

  • 저자들은 본문 결과를 훈련 없는 제로샷 DA의 하한으로 규정하고 사후학습 및 에이전트 문맥으로의 확장을 논의한다. 현 방식은 Vanilla보다 디코드 단계를 약 3분의 1 더 사용하고, 고정 벤치마크 문맥을 인공적인 약 2K토큰 chunk로 나누며, thinking trace 안에서는 프로토콜을 따르지 못해 non-thinking mode만 평가했다. 정확도와 attention 효율을 함께 보상하는 SFT 또는 강화학습, 자연스러운 도구 호출·사용자 턴·검색 문서를 주소 단위로 활용하는 에이전트 문맥, 표준 도구 선언으로 노출한 thinking 단계는 실험으로 검증된 결론이 아니라 향후 개선 방향이다.
    • 분할이 증거 구조를 훼손하는 경우와 출력 길이가 문서 길이에 비례하는 경우가 중요한 실패 조건이다. 부록 Table 10에서 Gemma-4-31B의 RULER/cwe와 LBv2/structured_data는 전역 집계 또는 분할된 표가 필요해 Vanilla 대비 정확도가 크게 낮았다. fwe·summ_screen_fd·in_context_learning·book_sum_sort는 단계당 마스크 절감에도 긴 출력을 누적 생성해 평균 attended tokens가 Vanilla보다 높았다. 저자들은 전자에는 구조 보존 분할과 map-reduce식 누적을, 후자에는 긴 출력을 focus에서 생성하도록 학습하는 방안을 제안한다.
    • global 단계의 비용은 남아 있다. 저자들은 전역 탐색에 각 segment의 짧은 설명으로 만든 in-context index를 쓰거나, global 단계에 경량 스캔 기반 희소 어텐션을 결합하는 방안을 제시한다. speculative decoding은 단계당 KV 읽기를 줄이는 DA와 순차 단계를 줄이는 기술로서 상보적일 수 있지만, 이 논문에서 실험하지는 않았다.
    • 에이전트 시스템에서 retrieval은 무엇을 context에 넣을지 결정하고, DA는 이미 누적된 context 중 현 단계에서 무엇을 읽을지 결정한다는 구분을 제시한다. 선언이 span 전환 전에 텍스트로 나오므로, 초점 밖 KV를 호스트 메모리로 옮겼다가 다음 선언에 맞춰 prefetch하는 가역적 KV offloading도 가능성으로 논의한다.

부록

  • 부록은 동적 희소 어텐션, KV eviction, 모델 제어형 추론과의 관련 연구를 정리하고, vLLM 통합 방식과 roofline 비용 모델의 유도, 데이터·LLM 판정기·샘플링 설정, 추가 실험과 실패 사례, DA·Vanilla·채점 프롬프트 전문을 수록한다.

짧은 생각

가장 중요한 실험 설계는 DAnm 대조군이다. DA 프롬프트만 적용하고 마스크를 끈 DAnm은 Vanilla보다 생성 길이가 늘어 누적 attended tokens가 커지지만, 동일한 프롬프트에서 마스크를 켠 DA는 순절감을 보인다. 따라서 비용 감소를 출력 축소나 chunked 프롬프트 형식으로 설명하기 어렵고, Table 2의 DAnm 대 DA 비교가 마스크의 기여를 직접 뒷받침한다. 다만 서비스 적용 여부를 판단하려면 roofline 추정 외에도 같은 구현에서 배치 크기, prefill과 decode의 자원 경합, 태그 파싱과 block-table 갱신 비용을 포함한 실측 지연시간이 필요하다. 특히 구조화된 표나 전역 집계처럼 분할이 증거를 훼손하는 과제와, 긴 출력 때문에 누적 읽기량이 증가하는 과제에서 사후학습이 정확도와 종단간 비용을 함께 개선하는지를 확인해야 적용 범위를 판단할 수 있다.