Stratified Knowledge-Density Super-Network for Scalable Vision Transformers

TL;DR

This paper introduces a hierarchical super-network for vision transformers that enables scalable deployment across resource levels by organizing knowledge efficiently and applying novel importance-aware dropout techniques.

Contribution

It proposes a stratified knowledge-density super-network with WPAC for knowledge concentration and PIAD for knowledge stratification, improving scalability and efficiency of ViT models.

Findings

WPAC outperforms existing pruning methods in knowledge concentration.

PIAD effectively promotes knowledge stratification during training.

The combined approach offers a competitive alternative to traditional model compression and expansion.

Abstract

Training and deploying multiple vision transformer (ViT) models for different resource constraints is costly and inefficient. To address this, we propose transforming a pre-trained ViT into a stratified knowledge-density super-network, where knowledge is hierarchically organized across weights. This enables flexible extraction of sub-networks that retain maximal knowledge for varying model sizes. We introduce \textbf{W}eighted \textbf{P}CA for \textbf{A}ttention \textbf{C}ontraction (WPAC), which concentrates knowledge into a compact set of critical weights. WPAC applies token-wise weighted principal component analysis to intermediate features and injects the resulting transformation and inverse matrices into adjacent layers, preserving the original network function while enhancing knowledge compactness. To further promote stratified knowledge organization, we propose \textbf{P}rogressive \textbf{I}mportance-\textbf{A}ware \textbf{D}ropout (PIAD). PIAD progressively evaluates the importance of weight groups, updates an importance-aware dropout list, and trains the super-network under this dropout regime to promote knowledge stratification. Experiments demonstrate that WPAC outperforms existing pruning criteria in knowledge concentration, and the combination with PIAD offers a strong alternative to state-of-the-art model compression and model expansion methods.