[논문 리뷰] Routing Absorption in Sparse Attention: Why Random Gates Are Hard to Beat

Keston Aquino-Michaels|arXiv (Cornell University)|2026. 02. 11.

Advanced Memory and Neural Computing인용 수 0

한 줄 요약

엔드-투-엔드 per-query 희소 주의 게이트 학습은 Q/K/V의 공동 적응으로 인해 게이팅 신호를 흡수하므로 대부분 비효과적이다. 사후 증류 또는 분리된 학습은 거의 조밀한 성능에 근접할 수 있으며, 무작위 라우팅은 학습된 라우팅과 종종 동등한 성능을 보인다.

ABSTRACT

Can a transformer learn which attention entries matter during training? In principle, yes: attention distributions are highly concentrated, and a small gate network can identify the important entries post-hoc with near-perfect accuracy. In practice, barely. When sparse attention is trained end-to-end, the model's Q/K/V projections co-adapt to whatever mask is imposed, absorbing the routing signal until learned gates perform little better than frozen random gates. We call this routing absorption and present four independent lines of evidence for it in a controlled 31M-parameter transformer: (1) differentiable soft gating converges to nearly the same perplexity whether the gate is learned or random (48.73 +/- 0.60 vs. 49.83 +/- 0.04 over 3 seeds); (2) hard top-k gating receives exactly zero gradient through the mask; (3) a gate distilled onto co-adapted Q/K/V achieves high F1 against oracle masks but catastrophic perplexity when deployed (601.6 vs. 48.6 on mask-agnostic Q/K/V); and (4) stochastic mask randomization during training fails to prevent co-adaptation (78.2 ppl deployed dense vs. 37.3 baseline). We connect routing absorption to the same phenomenon in Mixture-of-Experts, where random routing matches learned routing because experts co-adapt to any router, but show that attention exhibits a structurally more severe form: shared Q/K/V parameters enable cross-layer compensation pathways absent in MoE, where experts are self-contained modules. The implication is that end-to-end sparse attention methods employing per-query token-level gating face absorption pressure proportional to the parameter asymmetry between the gate and the model, and that post-hoc approaches, which decouple representation learning from sparsification, sidestep this entirely.

연구 동기 및 목표

• 엔드-투-엔드 학습이 per-query 희소 주의 게이트를 통해 라우팅 이점을 거의 제공하지 않는다는 것 증명
• 다양한 학습 체계에서 게이트 학습이 무작위 또는 오라클 마스킹에 비해 얼마나 비교되는지 정량화
• 라우팅과 표현 학습을 분리하면 효과적인 희소 주의가 가능하다는 것을 보이고, 그 이유를 설명
• 희소 주의 방법에서 라우팅 흡수가 혼합-전문가(Mixture-of-Experts)와의 유사 현상과 관련되어 있음을 밝히고 실용적 함의를 논의

제안 방법

• WikiText-103에서 6-층, 256-차원, 4-헤드 트랜스포머를 학습(약 31M 파라미터)하며 희소도 수준 k=64에 대해 미분 가능 소프트 게이팅 사용
• 다수 시드에 걸쳐 엔드-투-엔드 학습된 게이트와 고정 무작위 게이트를 비교
• 마스크를 통한 그래디언트 흐름을 확인하기 위해 하드 top-k 게이팅 수행
• 밀집 및 공동 적응 Q/K/V 체크포인트에 게이트를 증류하여 배포 동작을 시험
• 공동 적응 방지를 위한 확률적 마스크 학습 평가
• 스케일 업 테스트(Qwen3-1.7B) 및 한 층/소수 층 흡수 실험으로 스케일에 대한 강건성 평가

Figure 2: Convergence dynamics under decoupled vs. co-adapted training. (a) Post-hoc gate-only training on the frozen 31M model: the learned gate converges from 46.8 to 37.3 ppl in 500 steps ( $>$ 99% of total improvement), while the frozen random gate stays flat. (b) Single-layer co-adaptation at Q

실험 결과

연구 질문

• RQ1 엔드-투-엔드 학습이 per-query 게이팅 메커니즘에서 무작위나 오라클 마스크에 비해 라우팅 이점을 향상시키는가?
• RQ2 게이트를 통한 그래디언트 신호가 Q/K/V와 게이트 간의 공동 적응을 극복하기에 충분한가?
• RQ3 분리된(사후) 라우팅 학습이 거의 조밀한 성능을 회복할 수 있는가, 그리고 왜 성공하는가?
• RQ4 스케일이 라우팅 흡수와 공동 적응의 강도에 어떤 영향을 미치는가?
• RQ5 학습된 라우팅에 의존하는 기존의 희소 주의 및 MoE 방법에 대한 시사점은 무엇인가?

주요 결과

• 엔드-투-엔드 소프트 게이팅은 고정 무작위 게이트와 거의 동일한 퍼플렉서리 언어 모델 성능을 보이며(48.73±0.60 대 49.83±0.04; 조밀도 37.32), 대부분의 라우팅 이점이 Q/K/V의 공동 적응에 흡수된 것을 시사
• 하드 top-k 게이팅은 마스크를 통한 그래디언트를 0으로 만들어 학습된 게이트와 거의 같은 성능을 보여 그래디언트가 없을 때도 동일한 결과
• 공동 적응된 Q/K/V에 증류된 게이트는 배포 성능(601.6)으로 수렴하지만 조밀도(48.6) 대비 마스크-특정 공동 적응이 강하게 작동하여 오라클 마스크에 대해 높은 F1을 보이더라도 배포 성능이 떨어짐
• 학습 중 확률적 마스크 무작위화는 흡수를 방지하지 못하고 표현력을 저하시킴(배포 거친 dense 78.2 ppl 대 기본 37.3)
• 사후 증류(분리된 라우팅)는 게이트 학습 단계가 훨씬 적은 수로도 거의 조밀한 성능을 달성할 수 있음(예: k=64에서 KL 증류로 1,000 단계 등, 128–256에서도 가능)
• 흡수는 스케일이 커질수록 강해진다: Qwen3-1.7B에서 오라클 perplexity는 k=64에서 조밀도와 0.4% 이내이나 엔드-투-엔드 라우팅은 공동 적응하에서 여전히 취약함; 파라미터 비대칭성(게이트 대 모델)이 이 현상을 주도

Figure 3: The absorption gradient at Qwen3-1.7B scale. As more layers unfreeze (increasing co-adaptation capacity), the random gate’s perplexity drops toward the learned gate’s level. The shaded area shows the gap shrinking from 31.5 (post-hoc, no co-adaptation) to 6.9 (29% of layers unfrozen). The