효율적인 멀티헤드 어텐션 을 위한 적응형 헤드 예산 할당

트랜스포머 (Transformers) 는 다양한 도메인에서 지배적인 아키텍처로 자리 잡았으며, 이는 주로 다중 헤드 어텐션 (multi-head attention) 이 다양한 표현 서브스페이스를 포착하는 데 효과적이기 때문입니다. 그러나 표준 다중 헤드 어텐션은 작업 요구 사항이나 입력의 복잡성과 관계없이 모든 입력에 대해 모든 헤드를 균일하게 활성화합니다. 많은 시나리오, 특히 텍스트 분류와 같은 거친 수준의 작업 (coarse-grained tasks) 에서 관련 정보는 종종 전역적이며 전체적인 주의 헤드의 다양성을 필요로 하지 않습니다. 그 결과, 고정된 수의 헤드를 사용하면 할당이 입력과 일치하지 않을 때 불필요한 계산 비용이 발생할 수 있거나 성능이 최적화되지 않을 수 있습니다. 이 한계를 해결하기 위해 우리는 계산 자원을 동적으로 할당하는 적응형 다중 헤드 어텐션 메커니즘을 갖춘 트랜스포머 아키텍처인 BudgetFormer 를 소개합니다. 우리의 접근 방식은 각 입력에 대해 필요한 주의 헤드의 수에 해당하는 헤드 예산 (head budget) 과 가장 정보량이 많은 헤드를 선택하는 관련성 분포 (relevance distribution) 를 모두 학습합니다. 우리는 또한 모델이 효율적인 사용 패턴으로 수렴하기 전에 효과적인 헤드 구성을 발견할 수 있도록 탐험과 활용의 트레이드오프를 기반으로 한 훈련 전략도 제안합니다. 다양한 복잡성을 가진 텍스트 분류 작업에 대한 실험 결과, 우리의 방법은 FLOPs 와 메모리 측면에서 추론 비용을 줄이는 동시에 표준 전체 다중 헤드 어텐션을 능가하는 성능을 달성함을 보여줍니다. 이러한 결과는 적응형 헤드 할당이 트랜스포머 모델의 효율성과 효과성을 모두 개선하기 위한 원칙적인 접근 방식임을 시사합니다.