On-Demand Attention은 모델이 필요할 때만 전체 문맥을 다시 읽는다

무엇인가

긴 문맥 추론과 에이전트 작업이 늘면서 디코딩 비용에서 어텐션이 차지하는 몫이 커지고 있다. 문제는 전체 어텐션(full attention) 디코딩이 다음 토큰 예측에 그 정보가 실제로 도움이 되는지와 무관하게 매 스텝마다 계속 커지는 과거 전체를 읽는다는 점이다. 반면 먼 정보가 필요한 정도는 생성 스텝마다 다르다. 이 논문은 전역 어텐션을 수행하기 전에, 즉 글로벌 읽기를 하기 전에 이미 프리트레인 모델의 디코딩 상태 안에 그 이득을 예측할 수 있는 정보가 들어 있다는 관찰에서 출발한다.

어떻게 동작하나

제안 방법인 On-Demand Attention(ODA)은 로컬 우선(local-first) 디코딩이다. 각 스텝에서 먼저 Local 어텐션을 계산한다. Local은 StreamingLLM의 초기 토큰 + 최근 윈도우 패턴을 따라 처음 s개 위치와 폭 w의 최근 윈도우(현재 토큰 포함)만 본다. 그다음 경량 recall head가 세 가지 입력, 즉 직전 위치에서 넘어온 은닉 상태 h_{t-1}, 현재 입력 토큰의 임베딩 E(x_t), 현재 Local 계산의 최종 정규화 은닉 상태 h_t^L을 받아 스칼라 점수 q_t를 만든다. q_t가 임계값 θ(기본값 0)를 넘으면 같은 스텝 직전 이력에서 Full 어텐션으로 현재 스텝을 다시 계산하고, 그렇지 않으면 Local 결과를 그대로 받아들인다. 선택된 은닉 상태는 공유 LM head를 거쳐 다음 스텝으로 전달되고, 해당 위치의 KV는 유지된 이력에 추가된다. 학습되는 것은 recall head 파라미터뿐이고 프리트레인 가중치는 그대로 두며, 과거 KV 캐시 전체를 보존하므로 한 스텝에서 건너뛴 정보가 나중 스텝의 recall에서 다시 접근 가능하다.

무엇과 다른가

recall head의 학습 목표는 같은 이력에서 Local과 Full을 짝지어 계산한 예측 이득이다. g_t = NLL_L(x_{t+1}) - NLL_F(x_{t+1}) = log p_t^F(x_{t+1}) - log p_t^L(x_{t+1})로 정의되며, 양수면 Full이 정답 토큰에 더 높은 확률을 준다는 뜻이다. Full 호출에는 추가 비용이 들므로 비용 페널티 λ를 빼서 d_t = g_t - λ를 만들고, 부호와 순서를 보존하면서 큰 값을 압축하는 부호 로그 변환 y_t = sign(u) log(1+|u|)를 적용한 뒤 Huber 손실로 회귀한다. 추론 시 절약은 ΔF = (1-ρ)F_F - (F_L + F_R)로 정리된다. 여기서 ρ는 Full을 호출하는 스텝 비율, F_F와 F_L은 각각 Full과 Local의 백본 FLOPs, F_R은 recall head의 FLOPs다. Local의 스텝당 어텐션 비용은 O(1), Full은 O(n)이고, 길이 n 시퀀스 전체로 누적하면 각각 O(n)과 O(n²)이 된다. 저자들은 이 절약을 실제 속도로 바꾸기 위해 vLLM에 GPU 측 조건부 실행을 구현해 Local 계산, recall 결정, 조건부 Full 재계산을 하나의 실행 흐름 안에 두었다.

어떻게 쓰나

실험의 주 모델은 Qwen3-1.7B이고, RULER16K(13개 과제, 과제당 100개 예시)와 LongBench v1(13개 과제, 2,550개 예시)에서 평가했다. Local과 ODA는 초기 4개 위치와 2,048토큰 최근 윈도우를 쓴다. RULER16K에서 ODA는 Full 호출률 41.6%로 81.17점을 기록해, Full의 81.94와 Local의 19.23 사이에서 Full-Local 격차의 대부분을 회복했다. LongBench v1에서는 Full 호출률 70.6%로 36.82점을 얻어 Full의 37.94에 근접했고, 두 경우 각각 격차의 98.8%와 92.2%를 되찾았다. 회복 정도는 과제마다 다르다. 영어 지문 검색은 10.33점 차이가 남았지만 중국어 지문 검색은 Full과 동일했다. 더 큰 Qwen3-8B는 Full 호출률 47.3%로 91.07점(Full 92.59, Local 25.43)을 냈고, 하이브리드 백본인 Qwen3.5-2B에서는 전체 이력을 읽는 6개 어텐션 레이어에만 제한을 걸어 94.08점(Full 94.35, Full 호출률 44.3%)을 기록했으며 나머지 18개 레이어는 자체 Gated DeltaNet 계산을 유지했다. Gemma-4-12B-it도 step-1,024 head로 95.27점(Full 호출률 57.8%, Full 96.61)을 냈다. 각 모델은 별도로 학습한 head를 쓰며, head 전이 가능성은 시험하지 않았다.

전제와 한계

호출 시점 선택 자체가 중요하다는 것을 보이기 위해, 저자들은 RULER16K의 5개 과제에서 ODA와 무작위 Full 호출을 비교했다. 임계값 0에서 ODA는 과제 평균 Full 호출률 41.33%로 88.96점을 냈고, 같은 호출률의 무작위 정책은 32.79 ± 0.95에 그쳤으며, 새로 측정한 Full 기준선은 88.63이었다. 이 격차는 시험한 8개 예산 모두에서 유지됐다. 즉 전역 호출 빈도만으로는 성능을 설명할 수 없고, 마주친 상태에 따라 호출을 배분하는 것이 결과를 좌우한다. 비용 측면에서는 128K 입력, 약 12.5% Full 호출 조건에서 ODA가 Full의 24.26% FLOPs만 사용해 75.74%를 줄였고 계산 비율로는 4.12배였다. vLLM에서 A100-SXM4-80GB 한 장, 배치 크기 1, CUDA Graphs 활성화 조건으로 같은 128K/12.5% 지점에서 중앙값 처리량이 75.54토큰/초에서 149.52토큰/초로 1.98배 올랐다. 다만 4K에서는 ODA가 15.8% 더 느렸다. 짧은 전역 읽기를 피해 얻는 이득이 Local 시도와 조건부 실행 오버헤드를 상쇄하지 못하기 때문이다.

저자들은 recall 결정을 이해하기 위한 진단도 제시한다. 직전 은닉 상태(H), 현재 토큰 임베딩(E), 현재 Local 은닉 상태(L)의 부분집합으로 7개의 head를 학습시킨 결과, H+E만 쓰면 정책이 거의 항상 Full을 호출했고(99.73%), L을 더하면 81.57점에 Full 호출률 76.22%가 됐다. 임베딩만 쓰는 head는 Full을 87.32%나 호출하면서도 49.88점에 그쳤다. 별도로 2,048개 외부 참조 위치에서 head 점수로 40%를 선택했을 때 위치당 순이득 0.150 nats를 얻어 균일 선택의 0.053 nats를 앞섰다. Local의 top-1 예측이 틀린 675개 위치에서는 이득 부호 AUROC가 0.643으로 Local 엔트로피의 0.486보다 높았고, 차이의 95% 구간은 [0.069, 0.230]이었다. 또한 두 모드가 모두 top-1을 맞힌 1,342개 위치에서도 Full은 718개에서 확률을 올리고 437개에서 내렸으며, Local이 틀린 345개 위치에서 Full이 양의 NLL 이득을 줬지만 실제로 교정한 것은 60개뿐이었다. 저자들은 이로부터 recall이 Local의 오류를 잡아내는 것이 아니라 상대적 이득을 예측한다고 주장한다.

실무 관점에서 ODA는 기존 체크포인트를 건드리지 않고 긴 문맥 디코딩 비용을 줄이는 선택지다. 백본 가중치는 고정한 채 recall head만 학습하면 되고, KV 캐시 전체를 유지하므로 한 번 건너뛴 정보를 나중에 다시 볼 수 있다. 다만 도입 전에 확인할 것이 있다. 이득은 문맥 길이와 recall 빈도에 함께 의존하므로 4K처럼 짧은 문맥에서는 오히려 손해일 수 있고, head는 모델별로 따로 학습해야 하며 전이 성능은 검증되지 않았다. 또한 vLLM 조건부 실행 구현이 전제이므로 서빙 스택에서 이를 지원하는지 확인해야 한다.

저자들이 밝힌 한계와 전제도 분명하다. 효율 측정은 고정된 continuation과 미리 정한 recall 스케줄로 실행 비용을 비교한 것으로, 학습된 정책의 품질 평가와는 별개다. Qwen3-1.7B head는 28.3M 파라미터에 196,608개 예시 학습 풀을 사용했고, 전체 자원 회계에는 동결 백본의 특징·라벨 계산이 포함되어야 하며 GPU 시간 측정은 아직 부록에 미완으로 남아 있다. Qwen3.5의 소규모 독립 테스트에서는 실패 부분집합 판별 우위가 재현되지 않았고, head 전이는 시험하지 않았다. 완전한 익명 구현 공개도 아직 보류 상태다.