Differentiable Model Scaling using Differentiable Topk

📄 arXiv: 2405.07194v1 📥 PDF

作者: Kai Liu, Ruohui Wang, Jianfei Gao, Kai Chen

分类: cs.CV, cs.AI

发布日期: 2024-05-12

备注: Accepted by ICML 2024

💡 一句话要点

提出可微模型缩放(DMS)方法，高效搜索网络最优宽度和深度。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 可微模型缩放 神经架构搜索 网络宽度 网络深度 可微Top-k 模型优化 自动化设计

📋 核心要点

  1. 现有网络架构设计依赖手动，导致配置欠优，神经架构搜索(NAS)方法效率低下。
  2. DMS以完全可微的方式建模网络宽度和深度，实现高效优化。
  3. 实验表明，DMS在图像分类、目标检测和语言建模任务上均优于现有方法。

📝 摘要（中文）

近年来，大型语言模型开启了智能涌现时代，网络缩放受到了越来越多的关注。目前，许多网络架构都是手动设计的，这通常会导致次优配置。虽然已经提出了神经架构搜索(NAS)方法来自动化这个过程，但它们的搜索效率很低。本研究引入了可微模型缩放(DMS)，提高了搜索网络最优宽度和深度的效率。DMS可以以直接且完全可微的方式对宽度和深度进行建模，从而易于优化。我们在各种任务中评估了DMS，包括视觉任务、NLP任务以及各种网络架构，包括CNN和Transformer。结果一致表明，DMS可以找到改进的结构，并且优于最先进的NAS方法。具体来说，对于ImageNet上的图像分类，DMS将EfficientNet-B0和Deit-Tiny的top-1准确率分别提高了1.4%和0.6%，并且优于最先进的零样本NAS方法ZiCo，同时仅需0.4 GPU天进行搜索。对于COCO上的目标检测，DMS将Yolo-v8-n的mAP提高了2.0%。对于语言建模，我们修剪后的Llama-7B优于先前的方法，具有更低的困惑度和更高的零样本分类准确率。我们将在未来发布我们的代码。

🔬 方法详解

问题定义：现有网络架构的设计通常依赖于手动调整，这不仅耗时耗力，而且难以找到全局最优的配置。神经架构搜索（NAS）旨在自动化这一过程，但现有的NAS方法往往面临搜索效率低下的问题，需要大量的计算资源和时间才能找到合适的网络结构。因此，如何高效地搜索最优的网络宽度和深度成为了一个重要的研究问题。

核心思路：DMS的核心思路是将网络宽度和深度的搜索过程转化为一个可微优化问题。通过引入可微的top-k算子，DMS能够直接优化网络的宽度和深度，而无需像传统NAS方法那样进行离散的架构选择。这种方法使得搜索过程更加高效，并且能够更好地利用梯度信息来指导搜索方向。

技术框架：DMS的技术框架主要包括以下几个部分：首先，定义一个包含多种宽度和深度选择的超网络；然后，使用可微的top-k算子来选择超网络中的一部分结构作为最终的网络；接着，使用梯度下降算法来优化top-k算子的参数，从而实现对网络宽度和深度的自动调整；最后，在验证集上评估优化后的网络性能，并选择最优的网络结构。

关键创新：DMS最重要的技术创新在于使用了可微的top-k算子来实现对网络宽度和深度的选择。传统的top-k算子是不可微的，因此无法直接用于梯度下降优化。DMS通过对top-k算子进行可微近似，使得整个搜索过程可以端到端地进行优化，从而大大提高了搜索效率。与现有方法的本质区别在于，DMS避免了离散的架构选择，而是直接优化网络的连续表示。

关键设计：DMS的关键设计包括：1) 使用Gumbel-softmax技巧来近似可微的top-k算子；2) 设计合适的损失函数来平衡网络性能和复杂度；3) 使用warm-up策略来稳定训练过程；4) 针对不同的任务和网络架构，调整搜索空间和超参数。

🖼️ 关键图片

fig_0

📊 实验亮点

DMS在多个任务上取得了显著的性能提升。在ImageNet图像分类任务中，DMS将EfficientNet-B0和Deit-Tiny的top-1准确率分别提高了1.4%和0.6%，并且优于最先进的零样本NAS方法ZiCo 1.3%，同时仅需0.4 GPU天进行搜索。在COCO目标检测任务中，DMS将Yolo-v8-n的mAP提高了2.0%。在语言建模任务中，修剪后的Llama-7B优于先前的方法，具有更低的困惑度和更高的零样本分类准确率。

🎯 应用场景

DMS具有广泛的应用前景，可用于各种深度学习模型的自动设计和优化。例如，可以应用于图像分类、目标检测、自然语言处理等领域，帮助研究人员和工程师快速找到性能优异的网络结构。此外，DMS还可以用于模型压缩和加速，通过自动调整网络的宽度和深度，在保证性能的前提下，降低模型的计算复杂度和存储空间。

📄 摘要（原文）

Over the past few years, as large language models have ushered in an era of intelligence emergence, there has been an intensified focus on scaling networks. Currently, many network architectures are designed manually, often resulting in sub-optimal configurations. Although Neural Architecture Search (NAS) methods have been proposed to automate this process, they suffer from low search efficiency. This study introduces Differentiable Model Scaling (DMS), increasing the efficiency for searching optimal width and depth in networks. DMS can model both width and depth in a direct and fully differentiable way, making it easy to optimize. We have evaluated our DMS across diverse tasks, ranging from vision tasks to NLP tasks and various network architectures, including CNNs and Transformers. Results consistently indicate that our DMS can find improved structures and outperforms state-of-the-art NAS methods. Specifically, for image classification on ImageNet, our DMS improves the top-1 accuracy of EfficientNet-B0 and Deit-Tiny by 1.4% and 0.6%, respectively, and outperforms the state-of-the-art zero-shot NAS method, ZiCo, by 1.3% while requiring only 0.4 GPU days for searching. For object detection on COCO, DMS improves the mAP of Yolo-v8-n by 2.0%. For language modeling, our pruned Llama-7B outperforms the prior method with lower perplexity and higher zero-shot classification accuracy. We will release our code in the future.