[Paper] 주관적 깊이와 시간 스케일 트랜스포머: 언제 어디서 계산할지 학습
Source: arXiv - 2511.21408v1
Overview
Transformers는 현대 AI의 핵심이 되었지만, “모두에게 동일하게 적용되는” 연산 패턴—모든 토큰이 모든 레이어에서 서로를 어텐션하는 방식—은 특히 긴 시퀀스나 거대한 모델에서는 비효율적일 수 있습니다. 이 논문은 언제, 어디서 연산을 수행할지 학습하는 두 가지 새로운 Transformer 변형을 제안하여 불필요한 연산을 줄이면서도 성능을 유지합니다.
Key Contributions
- Subjective Depth Transformers (SDT) – 베이지안 서프라이즈를 이용해 토큰마다 전체‑정밀도 transformer 블록을 사용할지, 저렴한 “prior” 근사만 사용할지를 결정하는 Decision 레이어와 Dynamic 레이어를 교차 배치합니다.
- Subjective Timescale Transformers (STT) – 이 아이디어를 시간 차원으로 확장하여, 학습된 “변화 가설”에 따라 각 토큰에 대해 전체 transformer 블록을 건너뛰거나 실행하도록 라우터를 구성합니다.
- 베이지안 서프라이즈 신호(Expected and Unexpected Change)를 게이팅 기준으로 사용해 데이터의 새로움과 예측 가능성을 원칙적으로 구분합니다.
- 정적 컴퓨트 그래프 – 동적 라우팅에도 불구하고 전체 그래프는 정적이며, 기존 하드웨어와 컴파일러 스택에 쉽게 배포할 수 있습니다.
- 효율성 향상 – 실험 결과, self‑attention FLOPs를 최대 75 % 감소시키고 KV‑cache 사용량을 약 50 % 절감했으며, 정확도 손실은 미미합니다.
- 학습 동역학에 대한 실증적 증거 – 모델은 초기 학습 단계에서는 새로움에 기반한 게이팅을, 이후에는 예측에 기반한 게이팅으로 점진적으로 전환하며, 이는 서프라이즈 기반 처리에 대한 이론적 기대와 일치합니다.
Methodology
-
Decision Layer (SDT) – 각 토큰에 대해 두 개의 병렬 표현을 계산합니다:
- posterior (전체 transformer 블록) – 풍부한 컨텍스트를 포착합니다.
- prior (경량 선형 프로젝션) – 저비용 대체 역할을 합니다.
또한 각 토큰에 대해 posterior가 prior와 얼마나 다른지를 측정하는 베이지안 서프라이즈 점수를 추정합니다.
-
Dynamic Layer (SDT) – 고정 용량 Top‑K 라우터를 사용해 서프라이즈 점수가 가장 높은 K개의 토큰에만 비용이 많이 드는 posterior 연산을 적용하고, 나머지는 prior를 사용합니다. 라우터 선택이 점수에 따라 결정적이므로 전체 그래프는 정적으로 유지됩니다.
-
Transition Network (STT) – 각 토큰에 대해 잔차 업데이트를 예측하여, 토큰 표현이 시간에 따라 어떻게 변할지에 대한 가설을 형성합니다.
-
Temporal Router (STT) – 예측된 변화와 실제 변화를 서프라이즈 메트릭으로 비교합니다. 토큰이 “안정적”이라고 판단되면 해당 타임스텝에서 transformer 블록을 건너뛰고 캐시된 KV‑values를 재사용하고, 그렇지 않으면 블록을 실행합니다.
-
Training – 두 아키텍처 모두 표준 언어 모델링 목표로 end‑to‑end 학습됩니다. 서프라이즈 기반 게이트는 straight‑through estimator를 통해 미분 가능하게 구현되어, 모델이 데이터로부터 최적의 라우팅 정책을 학습할 수 있습니다.
Results & Findings
| Model | Compute Reduction | KV‑Cache Reduction | Perplexity (relative) |
|---|---|---|---|
| Baseline Transformer | – | – | 0.0 % |
| SDT (Depth gating) | self‑attention FLOPs 약 75 % 감소 | KV‑entries 약 50 % 감소 | +2–3 % |
| STT (Timescale gating) | self‑attention FLOPs 약 70 % 감소 | KV‑entries 약 45 % 감소 | +2–4 % |
- Surprise dynamics: 초기 에포크에서는 새로운 토큰에 대해 높은 게이팅 활동을 보이며, 이후 에포크에서는 실제로 놀라운 입력만이 비용이 많이 드는 연산을 트리거하도록 라우터가 안정화됩니다.
- Accuracy trade‑off: 퍼플렉시티가 약간 상승하지만, 큰 부분의 연산을 제거해도 언어 모델링 품질에 큰 영향을 주지 않음을 보여줍니다.
- Hardware friendliness: 컴퓨트 그래프 형태가 변하지 않기 때문에 GPU/TPU에서 별도 커스텀 커널 없이 효율적으로 실행됩니다.
Practical Implications
- Cost‑effective inference: 긴 문서(예: 법률 계약서, 코드베이스)를 서비스하는 경우, 예측 가능한 구간에서 어텐션을 건너뛰어 지연 시간과 GPU 메모리를 절감할 수 있습니다.
- Scalable training: 대규모 디코더‑전용 모델을 일반 하드웨어에서 학습할 때, 각 배치가 self‑attention에 소비하는 시간이 감소하여 훈련이 보다 현실적으로 됩니다.
- Edge and mobile AI: 정적 그래프 설계 덕분에 기존 툴체인(TensorRT, ONNX Runtime)으로 컴파일해 리소스가 제한된 디바이스에서도 가변 길이 입력을 처리할 수 있습니다.
- Fine‑grained control: 개발자는 서프라이즈 임계값을 런타임 파라미터로 노출해 속도와 품질을 실시간으로 조절할 수 있습니다(예: 배치 처리 시 aggressive pruning, 인터랙티브 채팅 시 보수적 게이팅).
- Foundation for adaptive APIs: 클라우드 제공자는 요청당 실제 수행된 연산량에 따라 비용을 청구함으로써, 모델 사용량과 비용을 보다 정확히 맞출 수 있습니다.
Limitations & Future Work
- Modest accuracy gap: 연산 절감 효과는 크지만, 현재 구현에서는 퍼플렉시티가 몇 퍼센트 상승하는데, 이는 고위험 응용에서는 허용되지 않을 수 있습니다.
- Surprise estimator overhead: 베이지안 서프라이즈 계산 자체가 작은 상수 비용을 추가하므로, 양자화 등으로 최적화하는 방안이 필요합니다.
- Generalization to encoder‑decoder or multimodal models: 본 논문은 디코더‑전용 언어 모델에 초점을 맞추었으며, 비전‑언어 또는 음성 모델에 대한 라우팅 메커니즘 확장은 아직 탐색되지 않았습니다.
- Dynamic hardware support: 그래프가 정적이지만, 실제 이득은 효율적인 Top‑K 선택 및 캐시 관리에 달려 있으므로, 하드웨어 스케줄러와의 tighter integration이 추가적인 속도 향상을 가져올 수 있습니다.
- Long‑term training dynamics: 새로움 기반에서 예측 기반으로 전환되는 현상은 더 깊은 이론적 분석이 필요하며, 커리큘럼 학습 전략에 영감을 줄 수 있습니다.
Bottom line: Transformer에게 “정말 이 연산이 필요할까?”라는 질문을 베이지안 서프라이즈 신호로 스스로 판단하게 함으로써, SDT와 STT는 더 똑똑하고, 더 저렴하며, 더 적응 가능한 딥러닝 모델을 향한 유망한 길을 열었습니다—대규모 AI 시스템을 구축하는 모든 사람에게 흥미로운 발전이라 할 수 있습니다.
Authors
- Frederico Wieser
- Martin Benfeghoul
- Haitham Bou Ammar
- Jun Wang
- Zafeirios Fountas
Paper Information
- arXiv ID: 2511.21408v1
- Categories: cs.LG, cs.AI, cs.CL, cs.IT
- Published: November 26, 2025
- PDF: Download PDF