잊기 변환기(Forgetting Transformer)를 위한 적응형 계산 가지치기

Adaptive Computation Pruning for the Forgetting Transformer

April 9, 2025

저자: Zhixuan Lin, Johan Obando-Ceron, Xu Owen He, Aaron Courville

cs.AI

초록

최근 제안된 포게팅 트랜스포머(FoX)는 소프트맥스 어텐션에 포겟 게이트를 통합하여 표준 RoPE 기반 트랜스포머와 비교해 일관되게 우수하거나 동등한 성능을 보여주었습니다. 특히, FoX의 많은 어텐션 헤드는 빠르게 포게팅되는 경향이 있어, 각 타임스텝에서의 출력이 주로 로컬 컨텍스트에 의존하게 됩니다. 이러한 관찰을 바탕으로, 우리는 FoX를 위한 적응형 계산 프루닝(ACP)을 제안합니다. 이 방법은 포겟 게이트에 의해 강하게 감쇠된 입력-출력 의존성을 포함하는 계산을 동적으로 프루닝합니다. 이는 프루닝된 어텐션 가중치가 무시할 수 있을 정도로 작아지도록 동적으로 설정된 프루닝 임계값을 사용하여 달성됩니다. 우리는 FoX를 사용한 언어 모델 사전 학습에 ACP를 적용하고, 다양한 모델 크기와 컨텍스트 길이에서 소프트맥스 어텐션의 FLOP 수를 약 70% 감소시켜 훈련 처리량을 약 10%에서 35% 개선함을 보여줍니다. 또한, 더 긴 컨텍스트 길이는 더 큰 계산 절약을 가져옵니다. 이러한 모든 속도 개선은 성능 저하 없이 달성되었습니다. 우리는 또한 프루닝 패턴을 검토하고, 다양한 어텐션 헤드 간의 FLOP 절감 분포를 분석하는 등 우리의 방법에 대한 더 깊은 통찰을 제공하기 위해 여러 분석을 수행했습니다. 우리의 코드는 https://github.com/zhixuan-lin/arctic-fox에서 확인할 수 있습니다.

English

The recently proposed Forgetting Transformer (FoX) incorporates a forget gate into softmax attention and has shown consistently better or on-par performance compared to the standard RoPE-based Transformer. Notably, many attention heads in FoX tend to forget quickly, causing their output at each timestep to rely primarily on the local context. Based on this observation, we propose Adaptive Computation Pruning (ACP) for FoX, a method that dynamically prunes computations involving input-output dependencies that are strongly decayed by the forget gate. This is achieved using a dynamically set pruning threshold that ensures that the pruned attention weights remain negligible. We apply ACP to language model pretraining with FoX and show it consistently reduces the number of FLOPs in softmax attention by around 70% across different model sizes and context lengths, resulting in a roughly 10% to 35% improvement in training throughput. Furthermore, longer context lengths yield greater computational savings. All these speed improvements are achieved without any performance degradation. We also perform several analyses to provide deeper insights into our method, such as examining the pruning patterns and analyzing the distribution of FLOP savings across different attention heads. Our code is available at https://github.com/zhixuan-lin/arctic-fox.