ScaleNet: Scaling up Pretrained Neural Networks with Incremental Parameters

TL;DR

ScaleNet tackles the high cost of scaling Vision Transformers by enabling depth expansion through weight sharing and lightweight adapters. It inserts new layers that reuse pretrained weights, with per-layer adapters to preserve capacity, enabling a scaling factor of $2\times$ in depth without doubling parameters. On ImageNet-1K, ScaleNet achieves a $7.42\%$ accuracy improvement over training from scratch with only $\frac{1}{3}$ of the training epochs, and transfers to downstream tasks like object detection and semantic segmentation. This approach reduces training time and resource use, facilitating deployment on resource-constrained settings while preserving or enhancing performance.

Abstract

Recent advancements in vision transformers (ViTs) have demonstrated that larger models often achieve superior performance. However, training these models remains computationally intensive and costly. To address this challenge, we introduce ScaleNet, an efficient approach for scaling ViT models. Unlike conventional training from scratch, ScaleNet facilitates rapid model expansion with negligible increases in parameters, building on existing pretrained models. This offers a cost-effective solution for scaling up ViTs. Specifically, ScaleNet achieves model expansion by inserting additional layers into pretrained ViTs, utilizing layer-wise weight sharing to maintain parameters efficiency. Each added layer shares its parameter tensor with a corresponding layer from the pretrained model. To mitigate potential performance degradation due to shared weights, ScaleNet introduces a small set of adjustment parameters for each layer. These adjustment parameters are implemented through parallel adapter modules, ensuring that each instance of the shared parameter tensor remains distinct and optimized for its specific function. Experiments on the ImageNet-1K dataset demonstrate that ScaleNet enables efficient expansion of ViT models. With a 2$\times$ depth-scaled DeiT-Base model, ScaleNet achieves a 7.42% accuracy improvement over training from scratch while requiring only one-third of the training epochs, highlighting its efficiency in scaling ViTs. Beyond image classification, our method shows significant potential for application in downstream vision areas, as evidenced by the validation in object detection task.

Figures (6)

• Figure 1: Comparison between (a) progressive training and (b) the proposed ScaleNet method. Progressive training initializes new layers with pretrained layers, and each layer is treated as independent during subsequent training. In contrast, ScaleNet shares the same parameters between the new layers and the pretrained layers, with only a small fraction of adjustment parameters to achieve layer-specific functions, surpassing progressive training in terms of parameter efficiency.
• Figure 2: Overview of the ScaleNet method. (a) A pretrained ViT model consists of a patch embedding layer, intermediate layers, and a prediction head, where the intermediate layers are the target for model scaling. (b) Scaling the pretrained $L$-layer model to $L'$ layers. In the figure example, new layers are added at the tail of the pretrained layers, with each new layer sharing weights with a corresponding pretrained layer. Other layer mapping strategies are explored in our experiments. Layer-specific adjustment modules are adopted to differentiate the weight-shared layers, improving model capacity. (c) Parallel adapter-based adjustment module, which consists of a down-projection layer, an up-projection layer, and a non-linear activation function. This module is placed at all linear layers in MLP modules of the ViT.
• Figure 3: Comparison of the $l_2$ norm of gradients for each layer in the DeiT-B model scaled using different methods. Our method presents more stable gradients during training.
• Figure 4: Feature distribution comparison between the original pretrained DeiT-Base model, ScaleNet before training, and ScaleNet after training. The histograms show the activation distributions for various layers, illustrating how the feature distribution evolves as the model is scaled and fine-tuned. The density in each figure on the y-axis is scaled to the same scale, and the x-axis is truncated to the range $[-10, 10]$.
• Figure 5: CKA similarity analysis between the original pretrained DeiT-Base model and ScaleNet before and after finetuning. The heatmaps show the alignment of feature representations across layers, with values closer to 1 indicating higher similarity. The results demonstrate that the scaled model before finetuning is already closely aligned with the original model, and finetuning further enhances this alignment. Coordinate axes indicate the layer indices.

Theorems & Definitions (1)

• Theorem 1