Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond Uniform Scaling: Exploring Depth Heterogeneity in Neural Architectures

Akash Guna R. T, Arnav Chavan, Deepak Gupta

TL;DR

The work tackles inefficiencies of uniform depth scaling by introducing depth heterogeneity guided by the loss landscape in vision transformers, leveraging Hessian information ($H$). It proposes a training-aware scaling mechanism that adds skip-connected neurons to selected layers based on Hessian-derived criteria via the Splitting Matrix eigenvalues, while preserving original functions at initialization. Applied to DeiT-S on ImageNet-100 and CIFAR-100, the method yields about a 2.5% accuracy gain with roughly 10% fewer parameters and demonstrates improved learning on medium-scale datasets from scratch. This approach provides the first intact scaling mechanism for vision transformers and shows that depth heterogeneity can enable more efficient, scalable neural architectures.

Abstract

Conventional scaling of neural networks typically involves designing a base network and growing different dimensions like width, depth, etc. of the same by some predefined scaling factors. We introduce an automated scaling approach leveraging second-order loss landscape information. Our method is flexible towards skip connections a mainstay in modern vision transformers. Our training-aware method jointly scales and trains transformers without additional training iterations. Motivated by the hypothesis that not all neurons need uniform depth complexity, our approach embraces depth heterogeneity. Extensive evaluations on DeiT-S with ImageNet100 show a 2.5% accuracy gain and 10% parameter efficiency improvement over conventional scaling. Scaled networks demonstrate superior performance upon training small scale datasets from scratch. We introduce the first intact scaling mechanism for vision transformers, a step towards efficient model scaling.

Beyond Uniform Scaling: Exploring Depth Heterogeneity in Neural Architectures

TL;DR

The work tackles inefficiencies of uniform depth scaling by introducing depth heterogeneity guided by the loss landscape in vision transformers, leveraging Hessian information (). It proposes a training-aware scaling mechanism that adds skip-connected neurons to selected layers based on Hessian-derived criteria via the Splitting Matrix eigenvalues, while preserving original functions at initialization. Applied to DeiT-S on ImageNet-100 and CIFAR-100, the method yields about a 2.5% accuracy gain with roughly 10% fewer parameters and demonstrates improved learning on medium-scale datasets from scratch. This approach provides the first intact scaling mechanism for vision transformers and shows that depth heterogeneity can enable more efficient, scalable neural architectures.

Abstract

Conventional scaling of neural networks typically involves designing a base network and growing different dimensions like width, depth, etc. of the same by some predefined scaling factors. We introduce an automated scaling approach leveraging second-order loss landscape information. Our method is flexible towards skip connections a mainstay in modern vision transformers. Our training-aware method jointly scales and trains transformers without additional training iterations. Motivated by the hypothesis that not all neurons need uniform depth complexity, our approach embraces depth heterogeneity. Extensive evaluations on DeiT-S with ImageNet100 show a 2.5% accuracy gain and 10% parameter efficiency improvement over conventional scaling. Scaled networks demonstrate superior performance upon training small scale datasets from scratch. We introduce the first intact scaling mechanism for vision transformers, a step towards efficient model scaling.
Paper Structure (19 sections, 6 equations, 2 figures, 4 tables)

This paper contains 19 sections, 6 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Method
  4. Workflow Description
  5. Using Hessian to Identify the Right Neurons
  6. Selecting Neurons to Scale
  7. Scaling neurons
  8. Training Details
  9. Dataset Description
  10. Hyperparameters
  11. Reducing Memory utilization using shared input tensors
  12. Experiments
  13. Comparison with unscaled DeiT-S
  14. Training CIFAR-100 from Scratch
  15. Identifying the ideal base model
  16. ...and 4 more sections

Figures (2)

  • Figure 1: The Basic Workflow of the Proposed Scaling Technique. (a) shows the regular DeiT architecture. (b) shows a reduced DeiT architecture where parameters are reduced from intermediate layers to form bottlenecks. (c) shows an example of scaled DeiT architecture grown from a reduced DeiT architecture. We scale only selected neurons (S) for scaling through skip connections and do not scale other neurons (NS) present in the layer. Our scaling technique is applicable to QKV, Projection and Fully Connected layers of DeiT.
  • Figure 2: Plots showing negative eigenvalues of neurons in QKV and FC1 layers of the first transformer block. Each neuron in the X-axis has its magnitude shown in the Y-axis. (Zoom to view X and Y axes).