Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design (2305.13035v5)

Abstract: Scaling laws have been recently employed to derive compute-optimal model size (number of parameters) for a given compute duration. We advance and refine such methods to infer compute-optimal model shapes, such as width and depth, and successfully implement this in vision transformers. Our shape-optimized vision transformer, SoViT, achieves results competitive with models that exceed twice its size, despite being pre-trained with an equivalent amount of compute. For example, SoViT-400m/14 achieves 90.3% fine-tuning accuracy on ILSRCV2012, surpassing the much larger ViT-g/14 and approaching ViT-G/14 under identical settings, with also less than half the inference cost. We conduct a thorough evaluation across multiple tasks, such as image classification, captioning, VQA and zero-shot transfer, demonstrating the effectiveness of our model across a broad range of domains and identifying limitations. Overall, our findings challenge the prevailing approach of blindly scaling up vision models and pave a path for a more informed scaling.

Analysis of "Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design"

The paper "Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design" proposes a novel approach to determining the optimal architectural shape for compute-efficient Vision Transformers (ViTs). The emphasis is on deriving compute-optimal model shapes—specifically their width and depth—by advancing the existing methodologies that focus solely on model sizes quantified by parameter counts. This approach introduces the shape-optimized ViT, or SoViT, which demonstrates competitive performance with significantly larger models trained under similar compute budgets. Specifically, SoViT-400m/14 attains a fine-tuning accuracy of 90.3% on the ILSVRC2012 benchmark, surpassing models like ViT-g/14 and approaching the performance of ViT-G/14 while incurring less than half the inference cost.

Contributions and Methodology

The paper addresses the growing trend of scaling model size without taking computational constraints into account by aligning model design with predefined compute budgets. It critiques the conventional reliance on parameter count as the sole indicator of model size, advocating instead for consideration of "shape dimensions" such as depth and width. Previous studies highlighted that transformer shape had minimal impact on scaling behavior in LLMing, but this paper emphasizes its critical role in vision tasks.

The authors employ a structured methodology to derive a scaling strategy for ViT architecture:

1. Joint Functional Formulation: The authors propose a joint functional form that considers both compute and individual shape dimensions (width, depth, MLP size). This approach extends existing literature on compute-optimal scaling by encompassing multiple dimensions rather than focusing on a single parameter count.
2. Star and Grid Sweeps: By employing a two-step approach of star and grid sweep, the methodology significantly reduces the need for extensive computationally costly experiments. The star sweep variance allows each dimension to be optimized independently, while the grid sweep helps identify a baseline compute-optimal architecture for small compute budgets.
3. Analytical Scaling Laws: The methodology involves deriving scaling exponents from empirical experiments, which inform the extent to which each shape dimension should be scaled relative to compute increase.

Experimental Results

The practical impact of their optimized approach is empirically validated across various vision tasks. The SoViT-400m/14 shows competitive performance not only in standard image classification benchmarks but also in diversified tasks like image captioning, optical character recognition, and multimodal VQA (Visual Question Answering). Furthermore, despite being smaller, SoViT models demonstrate the ability to match or exceed much larger ViT models in zero-shot classification and multi-task decoding efficiency. These results challenge existing norms suggesting that simply enlarging models leads to better performance, highlighting instead a more sophisticated balance between model shape and compute resource allocation.

Implications and Future Directions

This research prompts a reevaluation of strategies employed in designing neural architectures for large-scale vision tasks. By optimizing model shape alongside compute, deployers can realize improvements in model efficiency, which is pertinent given the increasing computational and environmental costs associated with training large AI models.

The findings encourage broader exploration of multi-dimensional scaling laws in other AI domains beyond vision and propose an alternative path in model development that emphasizes efficiency over mere scale. It would be interesting to see future research that further generalizes this approach and possibly extends it to LLMs or other architectures like recurrent or convolutional neural networks. Additionally, deploying such optimized models in real-world systems where computational resources are at a premium could be a fruitful area for exploration.

In conclusion, the paper brings forward an important perspective on how AI models, particularly transformers, can be reimagined to better align with the realities of available computational power while maintaining state-of-the-art performance.