Aprendiendo a omitir las capas intermedias de los Transformers

Learning to Skip the Middle Layers of Transformers

June 26, 2025

Autores: Tim Lawson, Laurence Aitchison

cs.AI

Resumen

La computación condicional es una estrategia popular para hacer que los Transformers sean más eficientes. Los métodos existentes suelen enfocarse en módulos individuales (por ejemplo, capas de mezcla de expertos) o saltan capas de manera independiente entre sí. Sin embargo, investigaciones sobre interpretabilidad han demostrado que las capas intermedias de los Transformers presentan una mayor redundancia, y que las primeras capas agregan información en posiciones de tokens. Guiados por estas ideas, proponemos una arquitectura novedosa que salta dinámicamente un número variable de capas desde el centro hacia afuera. En particular, un mecanismo de compuerta aprendido determina si se debe omitir un tramo simétrico de bloques centrales según la entrada, y un mecanismo de atención con compuerta evita que los tokens subsiguientes atiendan a las posiciones de tokens omitidas. Las normas residuales se controlan con un esquema de 'sándwich' o 'perilayernorm', y la dispersión de las compuertas con una pérdida de regularización adaptativa. Nuestro objetivo era reducir los requisitos de cómputo para tokens 'más simples' y potencialmente fomentar una jerarquía representativa emergente de múltiples niveles; sin embargo, en las escalas investigadas, nuestro enfoque no logra mejoras en la relación entre la entropía cruzada de validación y los FLOPs estimados en comparación con líneas base densas con menos capas. Publicamos nuestro código en 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.