Адаптивное сокращение вычислений для трансформера с забыванием

Adaptive Computation Pruning for the Forgetting Transformer

April 9, 2025

Авторы: Zhixuan Lin, Johan Obando-Ceron, Xu Owen He, Aaron Courville

cs.AI

Аннотация

Недавно предложенный Forgetting Transformer (FoX) включает в себя механизм забывания (forget gate) в softmax-внимание и демонстрирует стабильно лучшие или сопоставимые результаты по сравнению со стандартным Transformer на основе RoPE. Примечательно, что многие головы внимания в FoX склонны быстро забывать, что приводит к тому, что их выход на каждом временном шаге в основном зависит от локального контекста. На основе этого наблюдения мы предлагаем Adaptive Computation Pruning (ACP) для FoX — метод, который динамически обрезает вычисления, связанные с зависимостями вход-выход, которые сильно ослабляются механизмом забывания. Это достигается с помощью динамически устанавливаемого порога обрезки, который гарантирует, что обрезанные веса внимания остаются незначительными. Мы применяем ACP к предварительному обучению языковых моделей с использованием FoX и показываем, что он стабильно сокращает количество FLOP в softmax-внимании примерно на 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.