混合架构的机械化设计与规模化
机器学习
2024-08-20 v2
摘要
深度学习架构的开发是一个资源密集型过程，由于设计空间的广阔、原型制作时间的延长，以及大规模模型训练和评估的高计算成本。我们通过建立以机械化架构设计（MAD）管道为基础的端到端方法来简化这一过程，该管道包含小规模能力单元测试，能够预测规模化规律。通过一套用于探索能力的合成 token 操作任务（如压缩和回忆），我们识别并测试了由各种计算原语构建的新型混合架构。我们通过大量计算最优和一种新的状态最优规模化法则分析对所得架构进行了实验验证，在 70M 至 7B 参数范围内训练了超过 500 个语言模型。意外地，我们发现 MAD 合成任务与计算最优困惑度之间存在相关性，能够通过孤立的代理任务准确评估新的架构。通过 MAD 发现的新架构，基于混合化和稀疏性等简单思想， 在计算最优预算和过度训练 regime 中，均优于 Transformer、卷积和循环体系结构（Transformer++、Hyena、Mamba）。总体而言，这些结果提供了证据表明，对精选合成任务的性能可以预测规模化规律，最佳架构应通过混合拓扑结构利用专用层。
引用
@article{arxiv.2403.17844,
title = {Mechanistic Design and Scaling of Hybrid Architectures},
author = {Michael Poli and Armin W Thomas and Eric Nguyen and Pragaash Ponnusamy and Björn Deiseroth and Kristian Kersting and Taiji Suzuki and Brian Hie and Stefano Ermon and Christopher Ré and Ce Zhang and Stefano Massaroli},
journal= {arXiv preprint arXiv:2403.17844},
year = {2024}
}