Learning What Matters: Supervising Global Context Pruning with Causal Evidence Sets

Background Sparse attention lets each query read only a selected part of the context instead of all of it, cutting the cost of long-context inference. The selector (router) making that choice is usually trained by distilling a dense teacher model's attention distribution, resting on the assumption that where a teacher attends reveals what it actually depends on.

배경 희소 어텐션(sparse attention)은 각 쿼리가 문맥 전체가 아니라 선택된 일부만 읽게 해 긴 문맥 처리 비용을 줄이는 기법이다. 이런 선택기(라우터)는 보통 dense 교사 모델의 어텐션 분포를 그대로 모방하도록(distillation) 학습되는데, 이는 "교사가 어디에 주의를 기울이는지 보면 실제로 무엇에 의존하는지 알 수 있다"는 가정에 기대고 있다.

Problem That assumption has never really been tested against ground truth. Whether attention weights are trustworthy explanations has been debated for years, but without knowing exactly what evidence determines an answer, the debate had no way to be settled.

문제 문제는 그 가정 자체가 실증적으로 검증된 적이 없다는 것이다. 어텐션 가중치가 실제로 신뢰할 만한 설명(explanation)인지에 대한 논쟁은 오래됐지만, 정답을 결정하는 근거가 정확히 무엇인지 알 수 있는 정답(ground truth)이 없어 그 논쟁은 결론이 나지 않았다.

Prior limitations Prior sparse-attention work — fixed patterns (Longformer, BigBird), training-free selection methods, and learned routers trained via attention distillation — all used attention as the supervision signal directly, without ever checking that signal against known causal evidence.

기존 한계 기존 희소 어텐션 연구들 — 고정 패턴(Longformer, BigBird), 학습 없는 선택 기법, 어텐션 증류 기반 학습형 라우터 — 은 모두 어텐션을 라우팅의 지도 신호로 그대로 사용했을 뿐, 그 신호를 알려진 인과적 근거와 대조 검증한 적이 없다.

Goal Build synthetic tasks where the evidence for each answer is known exactly, and directly compare a router supervised with causal-masking-derived labels against one supervised with attention distillation.

목표 근거가 정확히 무엇인지 알 수 있는 합성 과제를 만들고, 인과적 마스킹으로 얻은 라벨로 지도한 라우터와 어텐션 증류로 지도한 라우터를 직접 대조하는 것이 목표다.

Method Causal evidence is defined by intervention: add -10^4 to the attention logits of a block's positions at every layer to "erase" it (without moving token positions), then check whether the answer flips. The resulting necessary/sufficient/interchangeable-family block labels become the router's training target directly, compared against a standard KL-distillation-from-teacher-attention router. It is like a blackout test — switching off each light in a room one at a time to see if the room stays lit tells you which light is actually load-bearing far more reliably than eyeballing how bright each one looks.

방법 인과적 근거는 개입(intervention)으로 정의한다 — 특정 블록의 위치에 해당하는 어텐션 로짓에 모든 레이어에서 −10⁴을 더해 그 블록을 "지운" 뒤(토큰 위치 자체는 옮기지 않는다), 답이 바뀌는지를 확인한다. 이렇게 얻은 필요/충분/대체가능(family) 블록 라벨을 라우터의 학습 목표로 직접 사용하고, 이를 표준적인 "교사 어텐션 분포에 대한 KL 증류" 라우터와 비교한다. 마치 정전 실험처럼 — 방 안의 각 조명을 하나씩 꺼보고 방이 계속 밝은지를 보는 것이 조명 하나하나가 "얼마나 밝게 보이는지"를 눈대중으로 재는 것보다 어떤 조명이 실제로 필요한지 더 정확히 알려주는 것과 같다.

Key idea The core move is using a non-differentiable but purely interventional signal — mask and recheck — instead of differentiable attention weights as the supervision target. This is deliberately blunt rather than a refined attribution technique, and that bluntness is the point: it carries no risk of re-inheriting some other model-internal artifact the way a fancier attribution method might.

핵심 아이디어 핵심은 미분 가능한 어텐션 가중치 대신, 미분 불가능하지만 순수한 개입(마스킹→재확인)을 지도 신호로 쓴다는 것이다. 이 방식은 정교한 귀속(attribution) 기법이 아니라 의도적으로 투박하다 — 그 투박함 덕분에 그 라벨이 모델 내부의 다른 아티팩트를 다시 상속받을 위험이 없다.

Validation Tested across five task families — unique retrieval, latest-write, duplicate evidence, multi-hop (up to 3 hops), and aggregation — plus cross-seed stability (Jaccard agreement), robustness under a 4x context-length shift, and a check of whether causal labels recovered without oracle access agree with ground-truth annotations (87-92%). The comparison is further extended to real pretrained models (Qwen2.5 at 0.5B/3B/7B, Gemma-2-9B) on a conflicting-fact task.

검증 단일 검색, 최신 갱신값 검색, 중복 근거, 최대 3-hop 멀티홉, 집계(count) 등 다섯 개 과제 계열에서 검증했고, 학습 시드 간 안정성(자카드 유사도), 4배 문맥 길이 이동에 대한 강건성, 그리고 정답 주석 없이 복원한 인과 라벨이 실제 정답 주석과 얼마나 일치하는지(87–92%)까지 확인했다. 또한 Qwen2.5(0.5B/3B/7B)와 Gemma-2-9B 같은 실제 사전학습 모델의 상충 사실 과제로 결과를 확장했다.

Results On a 2-hop chain task, the causally-supervised router reaches 0.99 routed accuracy (100/100/100% recall at each hop) versus 0.41 for the attention-distilled router — recall at the intermediate hop collapses to 0% because attention skips a step already "consumed" earlier in the computation. Causal labels are far more stable across training seeds (Jaccard 0.79-0.90 vs. 0.25-0.46 for attention-based selection) and hold up better under a 4x length shift (0.72-0.78 vs. 0.32-0.42). In the pretrained-model experiments, Qwen2.5-7B reaches 96% accuracy while still attending more to a stale record than the current one on 43% of examples.

결과 2-hop 체인 과제에서 인과 지도 라우터는 0.99의 라우팅 정확도(각 단계 재현율 100/100/100%)를 보인 반면, 어텐션 증류 라우터는 0.41에 그쳤다 — 이미 앞선 계산 단계에서 소비되어 "다 읽은" 중간 단계를 어텐션이 건너뛰기 때문에 그 단계의 재현율이 0%로 떨어진다. 인과 라벨은 시드 간 자카드 일치도가 0.79–0.90으로 어텐션 기반 선택(0.25–0.46)보다 훨씬 안정적이며, 4배 길이 이동에서도 더 잘 버틴다(0.72–0.78 대 0.32–0.42). 사전학습 모델 실험에서는 Qwen2.5-7B가 96%의 정답률을 보이면서도 여전히 43%의 사례에서 오래된 기록에 더 주의를 기울였다.

Comparison Against both attention distillation and an "adaptive attention" baseline (smallest block set covering 90% of attention mass), the three supervisions tie on simple single-record retrieval, but the causally-supervised router's advantage grows sharply as context lengthens or tasks require multi-hop chains.

비교 어텐션 증류 및 상위 90% 어텐션 질량을 덮는 "적응형 어텐션" 기준선 모두와 비교했을 때, 단일 레코드 검색처럼 단순한 과제에서는 세 방식이 비슷하지만 문맥이 길어지거나 멀티홉으로 복잡해질수록 인과 지도 라우터의 우위가 뚜렷해진다.

Significance From an efficient-AI standpoint, this reopens the question of what the correct training signal for a sparse-attention router should be — not imitating attention, but labels derived from causal intervention. It also shows that accuracy improves with scale faster than attention faithfulness does, which is a caution against treating raw attention weights as an interpretability signal more generally.

의의 효율적 AI 관점에서 이 논문은 희소 어텐션 라우터의 "정답 학습 신호"가 무엇이어야 하는지를 다시 묻는다 — 어텐션을 모방하는 것이 아니라 인과적 개입으로 얻은 라벨을 써야 한다는 것이다. 동시에 모델 규모가 커질수록 정답률은 개선되지만 어텐션의 신뢰도는 그보다 느리게(혹은 다르게) 개선된다는 점을 보여, 어텐션을 해석 가능성의 신호로 쓰는 관행 전반에 의문을 던진다.

Limitations Stated directly by the authors: the tasks are synthetic and the teacher models used in the main router experiments are small (10.7M parameters). Cost is measured as accuracy at a fixed block budget, not wall-clock time. Router transfer is checked across seeds and training length, not across tokenizers or architectures. The efficient label estimator handles redundancy and multi-hop chains separately but not combined in one example — only a costlier greedy search covers that mixed regime. The pretrained-model experiments cover only two tasks with forced-choice scoring at short (dozens-of-record) contexts, leaving free-form generation and document-scale context untested.

한계 저자들이 직접 밝힌 한계다: 과제는 합성이며 라우터 실험에 쓰인 교사 모델은 1,070만 파라미터로 작다. 비용은 블록 예산 대비 정확도로만 측정했고 실제 벽시계 시간(wall-clock)은 다루지 않는다. 라우터의 전이는 시드와 학습 길이에 대해서만 확인했고 토크나이저나 아키텍처 전이는 다루지 않는다. 중복(redundancy)과 체인이 한 예제 안에 함께 나타나는 혼합 상황은 효율적인 추정기가 처리하지 못하고, 더 비용이 큰 그리디 탐색만이 이를 다룬다. 사전학습 모델 실험은 두 과제, 강제 선택형 채점, 수십 개 레코드 수준의 문맥만 다뤄 자유형 생성이나 문서 규모 문맥은 열려 있다.

Future work No separate future-work section is given, but the limitations section names what remains open: extending to free-form generation and document-scale context, transferring the router across tokenizers and architectures, and finding a more efficient estimator for the combined redundancy-plus-chain regime.

향후 과제 논문은 별도의 향후 연구 절을 두지 않지만, 한계 절에서 다음이 미해결로 명시된다: 자유형 생성과 문서 규모 문맥으로의 확장, 토크나이저·아키텍처 간 라우터 전이, 그리고 중복과 체인이 함께 나타나는 혼합 상황을 다루는 더 효율적인 추정기.

Released resources No public code repository was found — no public release verified.

자원 공개 공개 코드 저장소는 확인되지 않았다 — 공개 링크 확인 안 됨.