[论文解读] Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design

Ibrahim Alabdulmohsin, Xiaohua Zhai+2 more|arXiv (Cornell University)|May 22, 2023

Domain Adaptation and Few-Shot Learning被引用 4

一句话总结

本文提出了 SoViT，一种通过联合优化深度、宽度和 MLP 维度而得到的计算效率最优的视觉 Transformer 架构，在仅使用一半推理成本的情况下，实现了与 ViT-g/14 等更大模型相当的性能。通过将缩放定律应用于模型结构而非仅参数量，SoViT-400m/14 在 ImageNet-1K 上达到 90.3% 的准确率，并在相同计算预算下，在零样本迁移和下游任务中超越更大的模型。

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.

研究动机与目标

• 为解决盲目扩大视觉模型带来的低效问题，通过联合优化参数量以外的架构形状（深度、宽度、MLP 大小）以实现计算效率的提升。
• 开发一种方法，能够在远少于暴力搜索所需实验次数的情况下，识别出计算效率最优的模型结构。
• 挑战现有假设——即更大的模型始终优于更小的模型——通过实证表明，在相同计算预算下，经过结构优化的小型模型可达到甚至超越大型模型的性能。
• 为视觉 Transformer 中不同架构维度在多样化视觉任务中的缩放行为，提供实证与理论洞察。

提出的方法

• 作者推导出联合优化视觉 Transformer 深度、宽度和 MLP 维度的缩放定律，仅通过少量初始实验即可外推得出最优配置。
• 通过系统性地对结构维度进行网格搜索，识别出计算效率最优的边界，从而最小化所需训练次数。
• 该方法使用幂律外推来预测不同模型结构下的性能表现，实现在无需完整训练所有候选模型的情况下实现高效搜索。
• 通过训练 SoViT-400m/14（一个 4 亿参数的模型）来验证该方法，使其在相同计算预算下达到 ViT-g/14（18 亿参数）的性能水平。
• 在图像分类、图像字幕生成、视觉问答、零样本迁移和全景分割等多个任务上评估模型泛化能力。
• 应用 Flexification 方法测试不同图像块大小下的鲁棒性，证实 SoViT-400m/14 在原始配置之外仍保持优异性能。

实验结果

研究问题

• RQ1 模型结构（深度、宽度、MLP 维度）能否被联合优化，以实现在视觉 Transformer 中的计算效率最优？
• RQ2 当使用相同的计算预算进行训练时，一个更小但经过结构优化的视觉 Transformer 是否能超越更大的模型？
• RQ3 在图像分类、图像字幕生成和视觉问答等不同视觉任务中，结构维度的缩放定律有何差异？
• RQ4 最优模型结构是否对输入分辨率或图像块大小的变化具有鲁棒性，如通过 Flexification 所验证的那样？
• RQ5 在密集预测任务（如全景分割）中，计算效率最优的结构存在哪些局限性？

主要发现

• SoViT-400m/14 在 ImageNet-2012 上微调准确率达到 90.3%，超过 ViT-g/14，并在相同计算设置下接近 ViT-G/14 的性能。
• 在零样本迁移中，SoViT-400m/14 在 LiT 基准上达到 82.2% 的准确率，优于 ViT-L/16，且与 ViT-g/14 持平。
• 在多任务解码中，SoViT-400m/14 在 CIDEr（67.7 vs. 68.8）、VQAv2（56.0% vs. 58.0%）和 GQA（52.9% vs. 52.5%）的准确率上达到或超过 ViT-g/14。
• 在全景分割任务中，SoViT-400m/14 达到 43.7 PQ，略低于 ViT-g/14（44.8 PQ），表明在密集预测任务中可能存在结构方面的局限性。
• 对 SoViT-400m/14 进行 Flexification 后，在不同图像块大小下仍保持强劲性能，证实其在原始配置之外的鲁棒性。
• 该方法仅需 115 次实验即可识别出最优缩放配置，远少于以往工作在单维优化中所需超过 400 次实验的规模。