[论文解读] Mechanistic Design and Scaling of Hybrid Architectures

Michael Poli, Armin W. Thomas+10 more|arXiv (Cornell University)|Mar 26, 2024

Architecture and Computational Design被引用 4

一句话总结

本文提出机械架构设计（MAD），一种快速、可扩展的流水线，通过合成的标记操作任务（如压缩和回忆）预测混合深度学习架构的缩放行为。通过在小模型上训练这些任务，MAD能够识别出性能优异的架构，在计算最优和过训练场景下均优于SOTA模型（如Transformer、Mamba、Hyena），且MAD得分与实际缩放性能之间具有强相关性。

ABSTRACT

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.

研究动机与目标

• 解决深度学习中架构原型设计成本高、效率低的问题，原因在于设计空间庞大且训练周期长。
• 开发一种系统化、可扩展的方法，在大规模训练前评估和优化新架构。
• 探究在精心设计的合成任务上的表现是否能预测新兴架构在真实世界中的缩放行为。
• 探索结合计算原语（如卷积、循环、MoE）的混合架构，以提升缩放效率和性能。
• 建立一种基于对模型能力机制性理解的快速、计算高效的架构搜索框架。

提出的方法

• 设计一套合成的标记操作任务——压缩、回忆、记忆——专门用于探测模型在上下文学习和序列保持等方面的能力。
• 应用MAD在小规模（如2层模型）下评估架构，使用归一化的状态维度以确保在循环与非循环原语之间的公平比较。
• 通过隔离的代理任务计算综合MAD得分，代表架构在多个单元测试中的能力概况。
• 在多种架构上训练超过500个语言模型（参数量从7000万到70亿），以执行计算最优和状态最优的缩放规律分析。
• 将MAD得分与The Pile上的计算最优困惑度进行相关性分析，以验证其预测能力，使用等FLOP组进行一致评估。
• 通过条纹化设计构建混合架构——交错排列不同计算原语的模块（如Mamba + 混合专家）以利用其专长能力。

实验结果

研究问题

• RQ1 合成的、小规模的标记操作任务能否作为大规模模型性能和缩放规律的可靠预测器？
• RQ2 结合专用计算原语（如门控卷积、循环）的混合架构是否在缩放效率上优于统一架构（如Transformer）？
• RQ3 MAD合成任务上的表现与真实世界语言建模中的实际计算最优困惑度之间是否存在强相关性？
• RQ4 MAD流水线能否通过识别哪些原语在特定子任务中表现优异，来指导架构的迭代改进？
• RQ5 状态最优缩放规律与计算最优规律有何不同？这对推理效率有何影响？

主要发现

• 综合MAD得分在所有FLOP预算下均与计算最优困惑度表现出强线性相关性，验证了MAD作为缩放行为预测代理的有效性。
• 混合架构——特别是结合Mamba与MoE或其他原语的条纹化变体——在缩放性能上优于SOTA模型（如Transformer++、Hyena、Mamba）。
• 通过MAD识别出的最佳架构在计算最优场景和过训练设置下均优于当前最先进模型。
• 记忆任务上的表现各架构间差异极小，且与缩放性能无关，表明该任务的区分能力有限。
• 状态最优缩放规律分析表明，固定状态块（如Mamba）在推理内存和延迟方面可实现更高效率，凸显了架构设计中的权衡。
• 基于MAD的迭代改进（如优化Hyena变体）在2×10¹⁹ FLOPs下直接转化为The Pile上困惑度的可测量提升，证实了该流水线的实际效用。