随着有关顶级人工智能公司在开发更新、更强大的大型语言模型 (LLM) 方面面临困难的谣言和报道不断流传,人们的注意力越来越多地转向“Transformer”的替代架构。近日,麻省理工学院诞生的初创公司 Liquid AI 报告了在自动化神经网络架构设计和定制方面的进展。他们开发了基于进化算法的定制架构合成(STAR)算法,该算法基于一种新的设计理论,用于为模型架构提供数值表示。
STAR 自动化了架构发现和优化的过程,将其转变为端到端的过程。通过这些方法,我们能够根据自定义的任务、指标和硬件来定制架构。Liquid AI 使用 STAR 合成了数百种不同的设计,这些设计在质量上优于强大的 Transformer 和混合架构,同时具有较小的缓存和参数数量。
模型架构设计是人工智能的一个基本支柱,它塑造了从扩展能力及效率到预训练、对齐和推理等基础的一切。架构设计中的一个关键挑战是平衡质量和硬件约束,特别是延迟和内存成本,以确保人工智能系统可以在不同的环境中有效部署。
即使将搜索限制在仅在 GPU 上快速运行的优化模型上,设计高性能架构仍然是一个非常重要的组合问题。这种复杂性通常会导致人工智能实验室和公司在早期就致力于特定的设计,因为手动启发式方法在预测性能权衡方面是有限制的。此外,不同应用领域的要求可能存在很大的差异:例如,为边缘用例设计的语言模型在内存占用、能耗方面应该很轻,并且能在特定的目标硬件上表现出良好的性能,例如在 CPU 上的快速预填充,而为云使用设计的语言模式通常优先考虑质量和 GPU 延迟,而不是内存占用。在实践中,需求变得更加微妙且复杂。
架构性能的基础是由构建它的计算单元以及这些计算单元之间的相互连接方式决定的。当前的大多数深度学习架构都是通过将注意力算子和门控线性单元进行顺序交错构建而成的。这些层是一个更大类别的计算单元的简单示例, Liquid AI 称之为线性输入可变系统(LIV)。LIV 是结构化运算符,其动作由输入本身逐点调制,并提供抽象来泛化不同类别的计算单元,如注意力变体、线性注意力、(门控)卷积、具有线性状态转换的(门控)递归、状态空间层和(门控)线性单元。
已经有证据表明“超越 Transformer”架构是具有潜力的。将简单的 LIV(如门控卷积和递归)与条纹混合模式中的自注意力相结合的模型已经显示了适度的质量改进、对更长序列的更有效的扩缩以及更快的推理。
虽然新的计算单元和互连策略在模型性能方面有了新的突破,但由于可能的设计数量众多,它们也对架构设计提出了关键的挑战。Liquid AI 利用的是针对 LIV 量身定制的进化算法,而不是依赖于应用于计算单元特定子类(如注意力和卷积)或互连策略(权重共享、KV 共享、并行互连)的手动优化和启发式方法。
STAR 的核心创新之一是将模型架构表示为称为 STAR 基因组的分层数值序列, Liquid AI 使用进化优化的原理对其进行进化。这个过程是迭代的:将基因组编译成具体的架构,对其进行评估,然后选择并重组性能最佳的架构来创建下一代。
重要的是,进化过程可以由静态和动态目标来指导:静态目标由架构的特定配置给出,例如它的参数数量或缓存大小。另一方面,动态目标需要评估架构,例如,通过测量其在给定数据集上训练后的困惑度或其在目标硬件上的延迟。
为了确保架构候选者是新颖且高性能的,基因组编码依赖于 Liquid AI 的设计理论,该理论为神经网络引入了一类新的通用计算单元:线性输入可变系统(LIV)。Liquid AI 已经确定并围绕基本机制构建了抽象,这些机制控制着深度学习中的现代计算单元如何根据输入上下文调节其计算。该框架以张量网络和系统理论为基础,通过两个关键方面来表征 LIV:它们的结构(算子的令牌和通道混合结构)和特征化(算子中输入依赖的函数形式)。为了能够表示复杂的架构设计, Liquid AI 的框架将运算符组合视为首要关注点,为构建超越层顺序堆叠的架构开辟了新的途径。
STAR 基因组使 Liquid AI 能够将 LIV 设计空间映射到适合进化优化的分层数值编码。它定义了编码架构所采用的每个计算单元的特征,以及这些单元是如何相互连接的。
Liquid AI 首先在改进的语言建模架构的设计中评估 STAR,针对三个组合目标进行优化:i.)质量(训练后的困惑度),ii.)质量和参数效率,以及 iii.)质量和缓存效率。
在经过短短的两到三轮进化之后,大多数架构的表现都优于 Transformer 和强混合基线之类的主流产品,并且随着更多轮次的执行,其表现持续改善。特别是,当仅优化质量时, Liquid AI 发现所有经过评估的 STAR 进化架构在下游评估基准上的表现都优于注意力循环混合架构,在基准上的改进是 Transformer 混合架构的两倍。Liquid AI 发现的这一结果有力地证明了进化搜索在设计空间中的有效性。混合架构是经过手动设计和改进的,需要投入大量的资源,而 STAR 可以在不到一天的时间内生成架构,命中率超过 90%。
该搜索可以支持多目标优化。当同时优化质量和模型大小时,进化的架构始终优于 Transformer 和条带混合型架构,同时减少了参数数量,使 Liquid AI 能够针对边缘和资源受限的环境压缩模型。Liquid AI 尝试了不同的方法来跨尺度传递进化和评估的结果,通常会发现对薄而深的架构候选者进行优化,比在目标宽度上优化架构主题能产生更好的传递效果。
同样, Liquid AI 使用 STAR 来平衡目标硬件的质量、模型大小和延迟,这些是通过直接在推理堆栈上进行分析获得的。这是可能的,因为 STAR 不需要度量的梯度:它与使用从基因组编译的架构计算的静态和动态度量的混合兼容,包括详细的分析过程,以最大限度地减少延迟和通信开销。
除了针对特定目标优化架构之外,STAR 还提供了一种分析工具,用于识别进化过程中出现的重复架构主题,从而推动观察到的性能提升。有趣的是,之前提出的手动互连模式,如 KV 共享和某些形式的权重共享,也自然而然地出现了,同时还出现了全新的模式。
Liquid AI 用 STAR 展示的功能只是在暗示它的全部潜力。得益于优化任何指标组合的能力,再加上 LIV 的多功能性,合成设计的多样性和质量得以不断提升。随着对哪些模式和目标同时出现的理解的不断加深, Liquid AI 现在希望进一步完善 STAR 的进化算法和初始种群,使每一代都比上一代更好,包括在基因组层次结构的最低级别进行优化。Liquid AI 表示也有兴趣将类似的方法应用于可以构建模块化设计空间的其他领域。
有关更多的详细信息,请参阅论文:https://arxiv.org/abs/2411.17800
相关链接:
https://www.liquid.ai/research/automated-architecture-synthesis-via-targeted-evolution
评论