arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2403.17844cs.LG
收藏

混合架构的机制设计与扩展(Scaling)

Mechanistic Design and Scaling of Hybrid Architectures

Together AI;Stanford University(斯坦福大学);Hessian AI;RIKEN(理化学研究所);Arc Institute;CZ Biohub;Liquid AI
Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, Ce Zh… 展开作者

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

AI总结:

本研究提出端到端机制架构设计(MAD)流程,基于小规模合成任务测试预测扩展定律,发现的混合架构在70M-7B参数规模下优于Transformer++等主流架构,验证了合成任务对架构扩展性能的预测价值。

摘要:

深度学习架构的开发是一个资源密集型过程,原因在于设计空间庞大、原型制作周期长,且大规模模型训练与评估相关的计算成本极高。我们旨在通过一套端到端的机制架构设计(Mechanistic Architecture Design,MAD)流程简化这一过程,该流程包含可预测扩展定律(scaling laws)的小规模能力单元测试。通过一系列旨在探测能力的合成词元操作任务(如压缩与召回),我们识别并测试了由多种计算基元构建的新型混合架构。我们通过广泛的计算最优扩展定律分析和全新的状态最优扩展定律分析,对得到的架构进行了实验验证,训练了500余个参数量在70M到7B之间的语言模型。令人意外的是,我们发现MAD合成任务的表现与计算最优困惑度(perplexity)具有相关性,这使得通过孤立的代理任务准确评估新架构成为可能。通过MAD发现的新架构基于混合化和稀疏性等简单思路,在计算最优预算和过训练场景下的扩展表现均优于当前最优的Transformer、卷积和循环架构(Transformer++、Hyena、Mamba)。总体而言,这些结果证明,精心设计的合成任务上的性能可以预测扩展定律,且最优架构应通过混合拓扑结构利用专用层。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑