Leren om de middelste lagen van Transformers over te slaan

Learning to Skip the Middle Layers of Transformers

June 26, 2025

Auteurs: Tim Lawson, Laurence Aitchison

cs.AI

Samenvatting

Conditionele berekening is een populaire strategie om Transformers efficiënter te maken. Bestaande methoden richten zich vaak op individuele modules (bijvoorbeeld mixture-of-experts lagen) of slaan lagen onafhankelijk van elkaar over. Interpretatieonderzoek heeft echter aangetoond dat de middelste lagen van Transformers meer redundantie vertonen en dat vroege lagen informatie aggregeren in tokenposities. Geleid door deze inzichten stellen we een nieuwe architectuur voor die dynamisch een variabel aantal lagen van het midden naar buiten overslaat. In het bijzonder bepaalt een geleerd gatingmechanisme of een symmetrisch blok van centrale lagen moet worden overgeslagen op basis van de invoer, en een gated attentionmechanisme voorkomt dat volgende tokens aandacht besteden aan overgeslagen tokenposities. Residunormen worden beheerst met een 'sandwich'- of 'perilayernorm'-schema en gatesparsity met een adaptief regularisatieverlies. We hadden als doel de rekenvereisten voor 'eenvoudigere' tokens te verminderen en mogelijk een opkomende multi-level representatiehiërarchie te bevorderen, maar op de onderzochte schalen bereikt onze aanpak geen verbeteringen in de afweging tussen validatie-kruisentropie en geschatte FLOPs vergeleken met dichte baselines met minder lagen. We hebben onze code vrijgegeven op 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.