Apprendre à sauter les couches intermédiaires des Transformers

Learning to Skip the Middle Layers of Transformers

June 26, 2025

Auteurs: Tim Lawson, Laurence Aitchison

cs.AI

Résumé

Le calcul conditionnel est une stratégie populaire pour rendre les Transformers plus efficaces. Les méthodes existantes ciblent souvent des modules individuels (par exemple, des couches de mélange d'experts) ou sautent des couches de manière indépendante. Cependant, les recherches en interprétabilité ont montré que les couches intermédiaires des Transformers présentent une plus grande redondance, et que les premières couches agrègent l'information dans les positions de tokens. Guidés par ces observations, nous proposons une nouvelle architecture qui saute dynamiquement un nombre variable de couches à partir du milieu vers l'extérieur. En particulier, un mécanisme de gating appris détermine s'il faut contourner une plage symétrique de blocs centraux en fonction de l'entrée, et un mécanisme d'attention gérée empêche les tokens suivants de prendre en compte les positions de tokens sautées. Les normes résiduelles sont contrôlées avec un schéma 'sandwich' ou 'perilayernorm', et la parcimonie des gates avec une perte de régularisation adaptative. Nous avions pour objectif de réduire les besoins en calcul pour les tokens 'plus simples' et potentiellement favoriser une hiérarchie de représentation multi-niveaux émergente, mais, aux échelles étudiées, notre approche n'améliore pas le compromis entre l'entropie croisée de validation et les FLOPs estimés par rapport aux modèles denses avec moins de couches. Nous publions notre code à l'adresse 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.