Scaling Adaptive Depth with Norm-Agnostic Residual Networks

2026-06-14

AI Summary

This work addresses the performance limitations of deep residual networks, which suffer from rapidly growing residual stream norms that suppress updates in deeper layers during scaling. To overcome this, the authors propose Norm-Agnostic Residual architecture (NAG), which decouples the magnitude and direction of residual streams to preserve effective layer contributions as depth increases. Additionally, they introduce a Mixture-of-Depth (MoD) mechanism that adaptively skips either attention or MLP layers. By treating depth sparsity as a novel scaling dimension under a fixed computational budget, the method achieves performance on par with full-depth baselines using only 20%–25% MoD rates, significantly reducing forward FLOPs and activation memory while enabling more efficient training of deeper models.

Abstract

Residual architectures are ubiquitous in deep learning, but they suffer from a subtle structural limitation: the norm of the residual stream can grow rapidly with depth. As a result, updates from later layers become small relative to the accumulated residual state. This reduces their impact on the representation and limits the benefits of scaling models in depth. To address this, we introduce NAG, a norm-agnostic residual architecture that separates magnitude from directional information in the residual stream, preserving meaningful layer contributions throughout depth and preventing later updates from being systematically suppressed by residual-norm growth. Importantly, NAG introduces only a negligible number of additional parameters and relies on simple operations that are easily kernel-fusible, preserving training efficiency in practice. We show that this architecture outperforms baseline Transformers, with gains that increase substantially as depth grows, enabling effective training of much deeper models. The norm-agnostic formulation also leads to an interpretable Mixture-of-Depths (MoD) mechanism that adaptively skips both attention and MLP layers. Beyond serving as a post-training accuracy-compute tradeoff, this mechanism can be used as a pretraining-time scaling strategy: under iso-FLOP training, compute saved by reducing per-token forward-pass cost can be reinvested into training on more tokens while keeping the total parameter count and KV-cache budget fixed. In our experiments, moderate Mixture-of-Depths rates of approximately 20%-25% match full-depth baseline performance under equal training compute while substantially reducing the number of executed layer parameters and forward-pass FLOPs. These results identify sparsity in depth as a new scaling axis for fixed-compute training, enabling very deep yet FLOP-efficient models.

Problem

• residual networks
• depth scaling
• norm growth
• layer contribution
• model depth

Innovation

• norm-agnostic residual
• Mixture-of-Depths
• depth scaling
• residual norm growth
• compute-efficient training