Learning to Skip the Middle Layers of Transformers

2025-06-26

AI Summary

To address redundant computation in Transformer middle layers, this paper proposes a dynamic symmetric skip-layer architecture: (1) a learnable gating mechanism adaptively skips symmetric layer blocks outward from the network center; (2) gated attention masks out invalid attentions to already-skipped positions; and (3) sandwich or per-layer normalization is adopted for residual connections, coupled with an adaptive regularization loss to control gate sparsity. The method reduces FLOPs while preserving representational capacity, establishing a novel paradigm for conditional computation and hierarchical skipping. Although it does not surpass shallower dense baselines on the validation-set cross-entropy–FLOPs trade-off curve, it introduces a structurally innovative, conceptually transparent, and learnable skip-layer framework for efficient Transformers.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Responsible Web: Measurement, analysis, and circumvention of Web censorship
• Security and Privacy: Large-scale security measurements

Abstract

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.

Problem

• Reducing redundancy in middle Transformer layers
• Dynamic skipping of central Transformer blocks
• Improving efficiency via conditional computation

Innovation

• Dynamic skipping of middle Transformer layers
• Learned gating mechanism for layer skipping
• Gated attention prevents attending to skipped tokens