STAR(Synthesis of Tailored Architectures):深度学习的合成定制化架构

论文STAR: Synthesis of Tailored Architectures(《STAR: 合成定制化架构》)提出了一种名为STAR的创新方法,用于优化深度学习模型架构。STAR采用了一种结合新的搜索空间和进化算法(evolutionary algorithm)的方法,使得在多个目标(例如模型质量、参数量、推理缓存大小等)上优化模型成为可能。

论文作者为Armin W. Thomas, Rom Parnichkun, Alexander Amini, Stefano Massaroli, Michael Poli,均来自Liquid AI公司。

一、引言

在引言中,作者指出深度学习领域中的模型架构设计存在两个主要的优化途径:自动化和手动设计。虽然手动设计可以利用人类的直觉和经验,产生有效的模型架构,但这种方法通常需要大量的资源和时间。而自动化架构搜索方法,尽管在某些特定任务中表现良好(例如卷积神经网络的优化),但由于搜索空间的局限性和过于简单的模式,常常无法显著超越现有的手工设计的标准模型。这也体现在Transformer时代,模型架构的同质化明显,几乎所有大型语言模型都采用了类似的架构。

为了解决现有方法的局限性,作者提出了STAR(Synthesis of Tailored Architectures)方法,这种方法通过引入新的搜索空间以及与进化算法的结合,致力于创建适应多样化需求的模型架构,以满足不同任务在效率和质量方面的要求。

二、搜索空间的基础

STAR方法的核心在于它的搜索空间设计。传统自动化架构设计面临的一个主要挑战是如何定义一个既稳定又全面的搜索空间。在这种背景下,作者提出了基于线性输入变动系统(LIV,Linear Input-Varying Systems)的搜索空间。LIV为深度学习中的各类计算单元(例如注意力变体、卷积、线性回归等)提供了一个广义化的框架。

2.1 线性输入变动系统(LIV)的定义

LIV是一类可用来描述输入序列的系统,其中输入是一个向量序列,例如 {x0,x1,…,x},其中每个向量都可以看作是一个token,每个token有多个channel组成(类似于一个多维特征向量)。在LIV框架下,计算由一个线性操作符控制,且该操作符的行为由输入自身决定。例如,在经典的注意力机制中,操作符的权重由输入通过某些线性或非线性变换计算得到,从而形成了LIV的一个典型实例。

2.2 LIV的层次化描述

LIV系统可以在三个层次上描述:

  1. 特征化(Featurization):定义输入上下文如何调制线性计算。这包括使用输入计算注意力权重的过程,例如通过线性变换得到query和key的特征。
  2. 操作符结构(Operator Structure):定义token和channel的混合结构,即如何在输入的不同token和不同channel之间进行线性组合。
  3. 骨架(Backbone):定义多个LIV如何组成一个完整的架构。在STAR中,这种骨架结构可以通过进化算法进行优化,从而找到最优的LIV组合形式。
三、STAR基因组

在STAR方法中,架构通过基因组(Genome)进行编码,称为“STAR基因组”。基因组的结构是层次化的,分为三个主要部分:特征化、操作符结构和骨架。每一个部分都用整数编码,使得它们可以通过进化算法进行优化。

3.1 骨架基因组

骨架基因组描述了LIV之间的连接方式。每个LIV都被表示为一个5位整数序列,其中包括:

  1. LIV类别:表示当前LIV的类型,例如是卷积、注意力还是某种线性变换。
  2. 特征化权重共享:确定在骨架中不同深度的LIV之间是否共享特征化权重。
  3. 特征化共享策略:定义如何实现特征化的共享,例如仅共享计算某些特征组的权重。
  4. 特征组共享:LIV之间可以直接共享特征组,这意味着某些计算可以重复使用,从而提高效率。
  5. 特征组共享策略:描述在所有特征组中,哪些特征组可以共享。

这种编码方式可以使得进化算法在不同的层次上对架构进行优化。例如,特征组共享策略的变化可以产生新的模型架构,而不需要改变整体的架构骨架。

3.2 操作符和特征化基因组

操作符基因组定义了LIV的具体形式,包括特征化方法、token混合结构等。特征化基因组则进一步描述了每个特征组的具体结构,例如channel的扩展因子、重复因子等。

四、通过进化算法合成架构

STAR通过进化算法对基因组进行优化,以找到最优的模型架构。整个进化过程包括以下步骤:

  1. 评估(Assessment):首先,对初始种群中的每个基因组进行评估。这包括通过训练模型并根据性能指标(例如困惑度PPL)对其进行打分。除了质量目标外,还可以进行静态分析,如参数数量和推理缓存大小,以评估效率。
  2. 配对(Pairing):然后,选择表现最好的个体进行配对,采用锦标赛选择(Tournament Selection),即随机选择一些基因组并挑选表现最好的个体。
  3. 重组(Recombination)和变异(Mutation):接着,通过交换遗传物质生成新的候选基因组,并引入随机变异以保持种群的多样性。重组通过k点交叉(k-point crossover)完成,即在基因组中随机选择k个点进行交换,从而生成新的候选基因组。
  4. 修复(Repair):如果重组或变异导致基因组配置无效(例如共享策略不兼容),则通过重新采样有效的基因组值来修复这些无效连接。

通过这些步骤,STAR逐渐演化出优化的架构,最终使得合成的模型在多个质量和效率目标上达到最佳。作者指出,通过结合LIV的搜索空间、基因组编码以及针对变异和重组的指导原则,STAR可以有效地获得稳定且性能优良的候选架构。

五、实验结果与分析

在实验中,作者将STAR应用于自回归语言建模,重点评估了STAR在多个质量和效率目标上的优化能力,包括:

  • 质量(困惑度PPL):这是衡量语言模型预测质量的主要指标。
  • 模型大小(参数数量):优化参数数量以减少模型的内存占用。
  • 推理缓存大小:这对于推理效率至关重要,尤其是在处理长序列输入时。
5.1 质量优化实验

实验显示,经过STAR演化的架构在RedPajama数据集上的评估困惑度显著低于手动设计的Transformer++和Stripe混合模型。在多个下游任务上,如HellaSwag、ARC-Easy、Winogrande等,STAR的表现超越了标准的混合架构,显示了更高的泛化能力。

5.2 参数高效模型的合成

在同时优化质量和模型大小的实验中,STAR成功减少了参数数量,同时保持了优异的质量表现。与手动设计的基准模型相比,STAR合成的模型参数数量减少了约13%,而性能仍然达到或超过基准模型。

5.3 高效推理缓存模型的合成

为了减少推理缓存的大小,STAR还优化了困惑度和缓存大小。实验结果表明,在相同的序列长度下,STAR合成的模型缓存大小比标准Transformer++减少了90%,并且在多个下游任务上保持了高质量。

5.4 发现架构的合成模式

通过进化算法,STAR还能够自动发现架构中有助于性能提升的模式(称为“架构图案”)。这些图案包括了对特定LIV的偏好,例如短卷积、分组查询注意力等。这些架构图案在多个任务上都表现出了较好的性能,进一步证明了STAR方法在架构合成中的灵活性和有效性。

六、总结与未来展望

论文总结了STAR在架构合成中的优势,并展示了它在优化模型质量和效率上的显著提升。相比于现有的手动或自动化架构设计方法,STAR结合了层次化的设计空间和进化算法,使得在复杂任务上的模型架构优化成为可能。此外,作者指出STAR具备良好的可扩展性,可用于更大规模的模型或更复杂的任务,例如将LIV架构与传统的缩放法则结合,进一步提升模型在大规模训练中的表现。


关于Liquid AI公司

Liquid AI是一家由麻省理工学院(MIT)计算机科学与人工智能实验室(CSAIL)前研究人员共同创立的初创公司,专注于开发新型人工智能模型,特别是液态神经网络(Liquid Neural Networks,LNN)。Liquid AI的创始团队包括MIT CSAIL主任Daniela Rus教授,以及博士后研究员Ramin Hasani、Mathias Lechner和博士生Alexander Amini。

液态神经网络(LNN)是一种新型的神经网络架构,其特点是结构动态可变,能够根据输入数据实时调整自身参数,从而提高对新环境的适应性。 与传统的深度学习模型相比,LNN在处理时间序列数据和需要实时响应的任务中表现出色。

2024年9月,Liquid AI推出了首批多模态AI模型,称为液态基础模型(Liquid Foundation Models,LFM)。 这些模型并非基于传统的Transformer架构,而是采用了基于动态系统理论的新型计算单元。 LFM在多个规模上实现了先进的性能,同时保持较小的内存占用和更高的推理效率。

Liquid AI的模型在多个基准测试中表现优异,特别是在处理长序列输入时,能够高效地利用上下文长度,显著减少内存占用。 例如,LFM-3B模型仅需16GB内存,而同等规模的其他模型可能需要超过48GB内存。

Liquid AI的技术适用于多种数据模态,包括文本、音频、图像和视频,具有广泛的应用潜力。 其高效的内存使用和推理能力,使其在资源受限的环境中也能发挥作用,如移动设备和边缘计算场景。

2023年12月,Liquid AI宣布完成3750万美元的种子轮融资,估值达到3.03亿美元。 投资者包括GitHub联合创始人Tom Preston-Werner、Shopify联合创始人Tobias Lütke和Red Hat联合创始人Bob Young等知名人士。

发表评论

您的邮箱地址不会被公开。 必填项已用 * 标注