Обучение пропуску промежуточных слоев в трансформерах

Learning to Skip the Middle Layers of Transformers

June 26, 2025

Авторы: Tim Lawson, Laurence Aitchison

cs.AI

Аннотация

Условные вычисления являются популярной стратегией для повышения эффективности трансформеров. Существующие методы часто нацелены на отдельные модули (например, слои смеси экспертов) или пропускают слои независимо друг от друга. Однако исследования интерпретируемости показали, что средние слои трансформеров обладают большей избыточностью, а начальные слои агрегируют информацию в позициях токенов. Руководствуясь этими выводами, мы предлагаем новую архитектуру, которая динамически пропускает переменное количество слоев, начиная с середины. В частности, обученный механизм затворов определяет, следует ли обходить симметричный диапазон центральных блоков на основе входных данных, а механизм внимания с затворами предотвращает доступ последующих токенов к пропущенным позициям. Нормы остатков контролируются с помощью схемы "сэндвич" или "perilayernorm", а разреженность затворов — с помощью адаптивной регуляризационной функции потерь. Мы стремились снизить вычислительные требования для "более простых" токенов и потенциально способствовать возникновению многоуровневой иерархии представлений, однако на исследованных масштабах наш подход не улучшает компромисс между кросс-энтропией на валидации и оценкой FLOPs по сравнению с плотными базовыми моделями с меньшим количеством слоев. Мы публикуем наш код по адресу https://github.com/tim-lawson/skip-middle.

English

Conditional computation is a popular strategy to make Transformers more efficient. Existing methods often target individual modules (e.g., mixture-of-experts layers) or skip layers independently of one another. However, interpretability research has demonstrated that the middle layers of Transformers exhibit greater redundancy, and that early layers aggregate information into token positions. Guided by these insights, we propose a novel architecture that dynamically skips a variable number of layers from the middle outward. In particular, a learned gating mechanism determines whether to bypass a symmetric span of central blocks based on the input, and a gated attention mechanism prevents subsequent tokens from attending to skipped token positions. Residual norms are controlled with a 'sandwich' or 'perilayernorm' scheme and gate sparsity with an adaptive regularization loss. We had aimed to reduce compute requirements for 'simpler' tokens and potentially foster an emergent multi-level representational hierarchy but, at the scales investigated, our approach does not achieve improvements in the trade-off between validation cross-entropy and estimated FLOPs compared to dense baselines with fewer layers. We release our code at https://github.com/tim-lawson/skip-middle.