Mechanistic Design and Scaling of Hybrid Architectures
作者: Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, Ce Zhang, Stefano Massaroli
分类: cs.LG
发布日期: 2024-03-26 (更新: 2024-08-19)
💡 一句话要点
提出机械化架构设计以简化深度学习模型开发
🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)
关键词: 深度学习 架构设计 混合架构 扩展规律 合成任务 性能评估 计算优化
📋 核心要点
- 现有深度学习架构开发过程复杂，资源消耗大，设计和评估效率低下。
- 提出了一种机械化架构设计（MAD）管道，通过小规模测试和合成任务简化架构开发过程。
- 实验结果表明，基于MAD的架构在性能上超过了当前最先进的模型，且能有效预测扩展规律。
📝 摘要（中文）
深度学习架构的开发过程资源需求高，设计空间广泛，原型时间长，以及大规模模型训练和评估的计算成本高。为简化这一过程，本文提出了一种基于端到端机械化架构设计（MAD）管道的方法，涵盖了小规模能力单元测试以预测扩展规律。通过一系列合成的令牌操作任务，识别并测试了由多种计算原语构建的新型混合架构。实验验证显示，基于MAD的合成任务与计算最优困惑度相关，能够通过孤立的代理任务准确评估新架构。基于简单理念如混合和稀疏的新架构在计算最优预算和过拟合情况下均超越了现有的Transformer、卷积和递归架构。
🔬 方法详解
问题定义：本文旨在解决深度学习架构开发过程中的高资源消耗和低效率问题。现有方法在设计和评估阶段面临长时间的原型制作和高计算成本的挑战。
核心思路：论文提出的机械化架构设计（MAD）管道，通过小规模能力单元测试来预测架构的扩展规律，从而简化设计过程。该方法强调通过合成任务来评估新架构的能力。
技术框架：MAD管道包括多个阶段：首先进行小规模的能力单元测试，然后通过合成的令牌操作任务（如压缩和回忆）来评估架构性能，最后进行计算最优和状态最优的扩展规律分析。
关键创新：最重要的创新在于通过合成任务与计算最优困惑度的相关性，提供了一种新的评估新架构的方式。这与传统方法的评估方式有本质区别。
关键设计：在设计中，采用了混合和稀疏的理念，构建了新型混合架构，并在训练过程中优化了参数设置和损失函数，以确保在不同规模下的性能提升。具体细节包括对超过500个语言模型的训练，参数范围从7000万到70亿。
🖼️ 关键图片
📊 实验亮点
实验结果显示，基于MAD的混合架构在计算最优预算和过拟合情况下均优于现有的Transformer、卷积和递归架构，具体性能提升幅度达到显著水平，验证了合成任务在预测扩展规律中的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、计算机视觉和其他需要高效模型设计的领域。通过提供一种简化的架构设计流程，研究成果可以加速新模型的开发和部署，提升实际应用中的效率和效果。
📄 摘要（原文）
The development of deep learning architectures is a resource-demanding process, due to a vast design space, long prototyping times, and high compute costs associated with at-scale model training and evaluation. We set out to simplify this process by grounding it in an end-to-end mechanistic architecture design (MAD) pipeline, encompassing small-scale capability unit tests predictive of scaling laws. Through a suite of synthetic token manipulation tasks such as compression and recall, designed to probe capabilities, we identify and test new hybrid architectures constructed from a variety of computational primitives. We experimentally validate the resulting architectures via an extensive compute-optimal and a new state-optimal scaling law analysis, training over 500 language models between 70M to 7B parameters. Surprisingly, we find MAD synthetics to correlate with compute-optimal perplexity, enabling accurate evaluation of new architectures via isolated proxy tasks. The new architectures found via MAD, based on simple ideas such as hybridization and sparsity, outperform state-of-the-art Transformer, convolutional, and recurrent architectures (Transformer++, Hyena, Mamba) in scaling, both at compute-optimal budgets and in overtrained regimes. Overall, these results provide evidence that performance on curated synthetic tasks can be predictive of scaling laws, and that an optimal architecture should leverage specialized layers via a hybrid topology.