摘要
蛋白质语言模型为蛋白质结构预测、功能注释和药物发现提供了突破性的方法。广泛采用这些强大的模型的一个主要限制是与这些模型的训练和推理有关的高计算成本，特别是在较长的序列长度上。本文介绍了一个蛋白质设计和发现加速器ProSE（ProteinSystolicEngine）的结构、微架构和硬件实现。ProSE有一系列定制的异构脉动阵列和特殊函数，可以有效地处理迁移学习模型的推断。该架构将SIMD风格的计算与脉动阵列架构结合起来，优化跨模型层的粗粒度操作序列，在不牺牲通用性的情况下实现效率。与一个NVIDIAAGPU相比，ProSE执行ProteinBERT推断的速度提高了6.9倍，功率效率提高了48倍（性能/瓦特）。与TPUv3（TPUv2）相比，ProSE实现了高达5.5倍（12.7倍）的速度提升和倍（倍）的功率效率。
1、介绍
自然语言处理（NLP）领域的最新创新使强大的深度学习算法能够准确预测蛋白质结构和药物-靶标相互作用，两者都可以减少药物发现和验证的超常成本（每年亿美元）、失败率（90%）和实验室到市场的平均时间（12年）。这些生物重新利用的NLP模型，尤其是那些采用Bert式模型结构的模型，比wetlab实验更快地产生推论，比传统的计算生物学算法更准确，提高了在昂贵人体试验中候选药物的质量。不幸的是，最新的GPU和TPU系统等商品加速平台在执行Bert风格的模型时缺乏能效。此外，这些系统的性能主要针对非NLP模型或以短输入长度的人类语言为目标的NLP模型进行优化，而蛋白质工程和发现需要更长的输入长度。
在处理Transformer的架构中，很少有人能与基于张量核心的现代GPU或基于脉动阵列的TPU竞争，而且至今没有加速器尝试长长度输入。图1描述了输入序列长度对BERT式Transformer效率的影响。对于语言模型，一个标记是一个词或子词，一个句子通常是30-80个标记。对于蛋白质模型，一个氨基酸是一个标记，而一个蛋白质是一个通常为多个标记的输入序列。在这些较长的输入长度中，操作混合的比例发生了变化，使得GPU或TPU平台的效率远远低于执行短输入长度的相同任务。此外，Transformer偶尔需要矩阵乘法，其矩阵比TPU使用的更小（x脉动阵列），但比GPUTensorcore优化的矩阵更大（4x4x8张量），从而导致TPU未得到充分利用，GPU性能较差。
图1.Bert式模型的推断功效
本文设计并评估了一个针对BERT式模型的收缩引擎，该引擎针对蛋白质大小的输入进行了优化，称为ProSE。ProSE是一个多线程的异构软硬件编码系统，包含异构脉动阵列和特殊函数，每种脉动阵列都能以数据流方式执行BERT式模型的操作序列，如矩阵乘法后的矩阵加法或GELU功能。ProSE使用一个输出稳定的脉动阵列对流式输入数据进行操作，没有提供专门的中间数据存储。相反，ProSE使用每个乘法-累加单元内的累加器寄存器作为中间存储，以减少大型人工智能模型的内存需求。
本文有以下贡献：
1.一个名为ProSE的性能、功率和面积高效的蛋白质发现系统引擎。使用具有不同大小和功能的异质流脉动阵列的集合，ProSE展示了一个能够有效执行矩阵乘法、SIMDALU操作和特殊激活功能的架构，与最新的GPU和TPU系统相比，提供了两到三个数量级的效率。
2.ProSE的设计过程包括对生物重新利用的BERT式模型的计算模式分析，详细的设计空间探索，以及微架构优化，以消除昂贵的中间存储，减少多余的数据移动，近似的特殊功能，并通过线程间的任务并行和线程内的管道并行提高效率。
3.对ProSE和TPU进行了详细的微结构比较，突出了三个创新点。1）使用异构脉动阵列，随着输入序列长度的增加，效率明显提高；2）通过本地数据流利用流式脉动阵列，而不是TPU的全局数据流，后者采用大型、耗电的统一缓冲器；3）新型的可左旋脉动阵列，与SIMD单元紧密结合，执行元素的SIMDALU操作或特殊功能，而无需与本地缓冲器通信。
2背景和动机
2.1人类语言的BERT模型vs.蛋白质BERT模型
在谷歌AI开发的人类语言BERT模型已经扩展为各种BERT风格的模型。如图2(a)所示，无论模型参数如何，这些模型将一个或多个句子作为输入，通过单词对输入进行标记（每个标记表示为一个矩形边界框，如"猫"），对BERT模型进行推理，并将下游模型用于所需的语言任务。
图2.(a)自然语言处理应用的BERT模型的输入和下游任务的总结
我们使用的蛋白质BERT模型在结构上与人类语言BERT模型相同。该模型接受一个蛋白质序列，以氨基酸字母表示，将序列标记为每个标记的单个字符，对蛋白质BERT模型进行推理，然后使用下游模型来执行所需的蛋白质设计任务，如图2（b）所示。为了执行抗体药物开发任务，蛋白质BERT模型可以使用学习蛋白质表征的特定预训练参数。然后对这些模型进行下游微调任务的训练，如荧光（某些蛋白质在合适的生物条件下会发出荧光）、稳定性（即预测一个蛋白质是否会保持其原始折叠构象或展开并改变构象）、结合亲和力（即两个蛋白质之间相互作用的强度）或结构预测（即从其氨基酸序列中预测一个蛋白质的三维结构，如AlphaFold）。在下面的实验中将在结合亲和力预测任务上训练一个下游模型。
图2.(b)对蛋白质设计应用的蛋白质BERT模型的输入和下游任务的总结
尽管人类语言的BERT风格模型和我们的蛋白质BERT模型之间没有结构上的差异（抛开模型参数不谈），但由于输入领域的不同（语言与蛋白质序列），这两种类型的模型的典型标记长度有很大的不同。对于人类语言模型，典型的输入序列长度为30至80个标记。然而，对于蛋白质BERT模型来说，蛋白质输入需要的最小输入序列长度在到个标记之间。
输入序列的长度极大地改变了BERT式模型的运行行为，带来了可扩展性的挑战，如在某些操作中，计算时间和内存占用都会随着输入序列长度的增加而二次增长。由于TPU或GPU等商品平台只对非BERT模型或短输入BERT式模型进行了优化，这个长度的问题进一步推动了对更长的序列长度进行更有效的专门设计。
2.2软件蛋白质结合的评估
预先训练的蛋白质语言模型的迁移学习预测能力已经被验证。本文通过对抗体Trastuzumab（又名为赫赛汀）和类似的抗体BH1进行下游模型的训练，对药物开发、结合亲和力预测的一般任务进行了硅质验证。一般来说，抗体和病原体之间的结合亲和力越高，抗体就越有可能与蛋白质结合并中和它，从而成为有希望的候选药物。在实验中，赫赛汀和BH1抗体都会结合HER2蛋白，这是乳腺癌中的一个关键靶点。通过在PyTorch中构建一个蛋白质BERT和下游任务，来评估软件模型对药物开发的可行性和适用性。蛋白质BERT模型的输入是赫赛汀和BH1变体的Fab子序列。Fab子序列是蛋白质结合最重要的部分（长度为个氨基酸）。
下游模型通过TAPE的ProteinBERT模型进行特征提取，并在39个变异赫赛汀Fab序列上拟合正则化线性回归模型，其中35个BH1Fab序列作为独立测试集。我们使用秩相关来衡量测试集的准确性，它可以告知哪种抗体被预测为最有可能与HER2蛋白结合。这个软件实验达到了0.的秩相关。此外，BERT式蛋白质设计软件的模块化使我们的工作流程有能力随着更大和更强大的蛋白质BERT式模型的开发而自动改进（无需人工工程），这将有可能缩小这一准确性差距。
2.3蛋白质BERT模型的分析
为了设计一个高效的蛋白质发现推理加速器，本文对一个在PyTorch中实现的蛋白质BERT模型进行了分析，该模型在输入序列长度为32到的范围内进行推理。这个蛋白质BERT模型的输入是合成蛋白质字符串。图3显示了表1中描述的机器配置在AGPU上的执行情况。我们对输入长度为32、64、、、、和的数据分别使用了、、、、、和64的批处理量，因为这可以利用GPU有限的40GiB内存获得最佳推理时间。我们观察到矩阵乘法（MatrixMultiply）的比例下降，而其他操作，特别是元素间的操作（MatrixAdd/Div），以及特殊函数（如Softmax）的比例上升。然而，在所有输入长度的分析中，矩阵乘法（批量和非批量）占总运行时间的35%-52%，这表明设计用于有效执行矩阵操作的加速器的重要性。
从蛋白质BERT模型中观察到。1)矩阵乘法（MatMuls）和分批矩阵乘法（BMMs）大部分时间用于执行比TPU规定的小得多的矩阵（即TPUv2和v3的×脉动阵列），但比A张量核心（即4×4×8张量计算）所提供的要大得多，并且2）从元素到SIMD的操作，包括特殊函数通常依赖于MatMul或BMM的结果。这些观察结果推动了设计决策，即使用较小的和不同大小的脉动阵列来执行较小的输入矩阵尺寸，以提高利用率，并构建ProSE，以使用新的左旋转能力的脉动阵列来执行MatMuls，然后在单个数据流中执行依赖的元素SIMD操作或特殊函数而不需要存储和重新获取中间数据。
图3.蛋白质BERT操作的运行时间与输入序列长度的关系
表1：用于BERT分析的AGPU平台的配置
3、ProSE系统
3.1软硬件协同设计
ProSE计算模式。为了利用ProSE中常见的计算模式，我们分析了蛋白质BERT模型的执行，并将大约90%的操作分为三个主要的操作序列，数据流1、2和3，如图5所示。每个模式都可以通过流水线数据流式链接完全在加速器上执行，以消除不必要的数据移动和主机-加速器通信。
图4.三个数据流
ProSE的脉动阵列。图5(a)描述了一个包含n×n处理单元（PE）的脉动阵列。为了支持三个主要的数据流，所有这些数据流都包含矩阵乘法和SIMDALU操作，我们将脉动阵列设计为两种操作模式：matmul和simd。在matmul模式下，脉动阵列通过从上到下、从左到右的PE移动数据来执行MatMuls，并在每个PE上执行乘积，如图5（b）所示。在simd模式下，脉动阵列作为大型的左列旋转器，通过从右到左移动数据穿过PE，执行SIMD式ALU操作，使用来自脉动阵列最左列的输入和存储流输入的向量寄存器执行SIMDALU操作，如图5（c）所示。为了在操作之间实现无停滞的数据流管道，simd模式从右向左移动数据，这样SIMDALU就可以在收缩阵列的最左列完成其矩阵乘法后立即开始执行，而不必等待其他列完成操作。这两种模式协同工作，通过将中间数据保留在PE累加器，实现无停顿交错的MatMul和SIMDALU操作。
图5该图在（a）中显示了尺寸为n×n的传统脉动阵列设计，在（b）中显示了脉动阵列在matmul模式下执行矩阵乘法，在（c）中显示了脉动阵列在simd模式下执行SIMDALU操作。
脉动阵列类型。为了支持图4中所示的数据流，每个数据流都具有不同的特殊功能，三个明显的设计选择就足够了：1）一个同质的脉动阵列集合，每个都有执行GELU和Exp的能力；2）一个同质的脉动阵列集合，其中一个子集有执行GELU的能力，另一个子集有执行Exp的能力；或3）一个异质的脉动阵列集合，其中一个子集有执行GELU的能力，另一个子集有执行Exp的能力。由于三个主要的操作序列是相互排斥的，并且执行模型的不同部分，我们选择将脉动阵列分为三种类型，每种类型能够执行一个操作序列，以最大限度地提高系统效率。ProSE中包含三种类型的脉动阵列：M型支持MatMuls和SIMDALU操作，G型支持MatMuls、SIMDALU操作和GELU特殊函数，E型支持MatMuls、SIMDALU操作和指数函数（Exps）。
脉动阵列的映射和尺寸。如图6所示，该模型可分为注意、中间和输出子层，以及数据流映射。例如，注意层计算以及输出子层使用最大的矩阵进行计算，通常矩阵尺寸m=，k=，n=。这些对应于执行数据流1的操作序列和M型收缩阵列。
图6.蛋白质BERT模型层：注意层、中间层和输出子层，以及它们的数据流映射
在计算不同大小的MatMul时，如果输入矩阵较小，较小的脉动阵列与较大的脉动阵列相比，启动和消耗的成本较低，但它们完成MatMuls的时间要长得多。使用图6(c)所示的设计计算SIMDALU操作，如矩阵加法或除法（通过乘以倒数常数），脉动阵列越小，可用的SIMDALU计算单元与PE的比例越大。例如，对于一个n×n的收缩阵列来说，有n个SIMDALU单元，有n个PE，比例为1：n。因此，较小的脉动阵列在每个PE上会有更多的SIMD单元，使任何SIMD密集型的操作序列表现得更好。
图3中的分批矩阵乘法由注意子层中的点乘组成，使用最小的矩阵，通常尺寸为m=，k=64，n=。这一点被数据流3和E型收缩阵列所涵盖，它们将矩阵除法和Exp与MatMul交织在一起，产生最SIMD密集的操作序列。为了最大限度地提高利用率，我们选择将E型收缩阵列实现为较小的收缩阵列，即32×32或16×16的收缩阵列。支持交织GELU和MatMul的Dataflow2和G-Type脉动阵列也是出于类似的原因，用中等或小的脉动阵列实现的。
异构脉动阵列的多线程执行。图7显示了蛋白质BERT模型中典型的数据依赖数据流序列（数据流1→3→1→2→1）以及软件如何被映射到三种类型的脉动阵列上。利用任务的并行性，使多个脉动阵列并行执行，大大增加了系统的吞吐量。本文展示了一个单线程、一个2线程、一个4线程和一个32线程的ProSE的协调和调度。每当一种类型的脉动阵列被映射到一个活跃的软件线程时，就会发生一次主机-加速器数据传输。我们实现了三个I/O缓冲区，每种类型的脉动阵列都有一个，并通过互斥锁处理线程争夺。随着线程的增加线程争夺的开销也会增加，通过实验为ProSE选择了32个线程。
图7.ProSE使用一种编排和调度机制，将软件线程映射到脉动阵列。多线程可以实现并行执行脉动阵列，显著提高了系统的吞吐量和效率
3.2ProSE体系结构和微架构
系统概述。ProSE是一个异构的输出平稳的流式脉动阵列集合，封装在一块加速器卡上，通过高带宽的外部接口如NVLink2.0与主机CPU通信。图8显示了ProSE的一个配置实例。对于我们的ProSE系统，我们设想了一个能够支持四个NVLinks的主机CPU，类似于最新的NVIDIAGraceCPU的能力，每个NVLink连接到一个ProSE实例，每个系统总共有四个ProSE实例。
图8.ProSE配置实例
ProSE操作。为ProSE设计的脉动阵列支持数据类型bfloat16和从ProteinBERT模型的特征中得到的以下五个原始操作的组合：
MatMul执行矩阵乘法C=A×B，其中A和B为输入矩阵，C为输出矩阵。
MulAdd执行C=αA+βB，其中α和β是标量常数。
MatDiv执行元素对等的倒数乘法C=A×1/α其中α是一个标量常数。
Exp使用一个查找表执行一个元素级指数函数exp(x)，其中x是输入矩阵A的一个元素，而C=exp
(A)。该操作用于支持与主机CPU配合的softmax激活函数。在脉动阵列上完成指数函数后，在CPU上执行softmax激活的求和和划分。
GELU执行元素级高斯误差线性激活函数
其中x是输入矩阵A的一个元素，C=GELU(A)。
ProSE处理元件。图9(a)显示了一个16×16的脉动阵列的示意图，该阵列使用一个8-deep的流媒体缓冲器来容纳来自主机的不间断流媒体，用于两个输入矩阵。我们验证了8-deep的流缓冲器足以覆盖所有类型和所有尺寸的ProSE脉动阵列所提供的NVLink带宽的延时。这些流缓冲器被合成为寄存器。图9(b)显示了我们的收缩阵列处理元件（PE）的特写微架构框图。MACs使用bfloat16数据类型执行，并使用类似TPU的32位累加器进行累加，以防止精度损失。每个PE中的32位累加器被用来作为中间存储。
图9.(a)每个ProSE收缩阵列配备了两个8-Deep流缓冲区和延迟槽，以允许流输入
(b)每个PE都是用一个16位乘法器、加法器和一个32位累加器来实现的
在TPUv2与ProSE上执行的MatMul。图10显示了在TPUv2和ProSE中如何使用脉动阵列进行矩阵乘法。为了说明这一点，我们展示了使用2×2脉动阵列对两个4×4的输入矩阵进行的MatMul。图10(a)是在TPUv2上运行所需的操作，图10(b)是在ProSE上运行所需的操作，图10(c)是执行MatMul的前两个步骤。A4×4×B4×4MatMul被分解成四个A2×4×B4×2MatMuls，在2×2脉动阵列上执行。
图10.MatMul在(a)TPUv2微结构和(b)ProSE微结构上执行
图10.(c)一个较大的MatMul分解成较小的MatMuls，在一个较小的脉动阵列上执行
(d)实现了输入部分缓冲区，允许在有限带宽场景下进行数据重用和提高性能
为了利用数据重用并减轻主机到加速器的流媒体带宽要求，我们实现了ProSE的配置，提供了一个输入部分缓冲区，如图10（d）所示。输入部分缓冲区有足够的空间来容纳输入矩阵A的一个步骤（在这个例子中，是一个2×4的矩阵），让步骤2重新使用矩阵A的上半部分，而不需要从主机上流。然后通过图10(d)中的操作执行我们所说的本地数据流，并重复进行，直到存储在部分缓冲区的输入数据必须被替换。请注意，输入部分缓冲区只存储来自主机的部分输入流，而不是中间数据，因为这些数据都存储在累加器中。
ProSE的特殊函数。ProSE支持两个特殊函数GELU和Exp，这些特殊函数在simd模式下执行，如图11所示。虽然TPUv2没有GELU激活单元，将被迫使用精度较低的RELU激活或GELU的近似扩展（如泰勒级数扩展），但ProSE将这些特殊函数查找表作为SIMDALU的一部分来实现，大大提高了执行这些特殊功能的性能和效率。GELU和Exp都是使用两级索引的查找表来实现的。
图11.每个SIMDALU单元实现一个GELU查找表，以提供快速的激活函数的近似值
4、PROSE评估
为了评估硬件资源对ProSE性能和效率的影响，我们在固定的90%可实现的NVLink2.0带宽GB/s的条件下，进行了PE数量从8K到24K的DSEs扫频。图12左边显示了资源与性能的关系，在右边显示了每个资源配置的BestPerf和MostEfficient配置的资源与功率效率的关系。我们观察到，在硬件资源配置为16KPEs(ProSE)或20KPEs(ProSE+)的情况下，设计是最平衡的，每个ProSE实例都具有相当高的功耗和性能。
图12.资源和性能关系
对于20KPEs的资源状况，我们探讨了更大的ProSE实例受带宽限制的可能性，并以更大的通信链路带宽GB/s（AGPU采用的90%可实现的NVLink3.0带宽）执行另一个DSE。
我们使用AGPU的一个实例来测量GPU系统的性能，并使用适合输入长度为tokens的批量大小来最大化性能。对于TPU系统的测量，我们使用谷歌云的TPUv3的一个实例。同样地使用适当的批处理规模来最大化推理性能。对于ProSE的性能和功率效率评估，我们使用输入序列长度为tokens，批处理大小为。
图13显示了ProSE的速度，左边是一个NVIDIAAGPU，右边是一个TPUv3实例（4个芯片/8个内核，硬件资源共计KPE），主机-加速器通信带宽不同。BestPerf和MostEfficient设计比A的速度提高了3.9-4.7倍，比使用NVLink2.0的TPUv3的速度提高了3.1-3.8倍。BestPerf+和MostEfficient+设计需要更快的链接，因为它们有更多的计算资源，并且在带宽达到GB/s之前不会成为计算障碍。然而，同构设计即使在无限的带宽下也不能提供理想的性能水平，因为它在执行较小的矩阵操作时受到启动和耗尽大型收缩阵列的巨大影响。此外，同质设计没有足够的SIMDALU和特殊功能单元，在对具有长输入长度的BERT式模型进行推理时，无法与异质设计竞争。
图13.在不同的链路带宽评估时，ProSE加速超过NVIDIAAGPU和TPUv3
ProSE设计的功率效率遵循图14所示的类似趋势，比GPU和TPU系统的效率提高了一到两个等级。即使我们选择了流式收缩阵列的设计，因为中间数据大多适合在主机上的大型L3缓存中，并且可以存储和检索，而不需要频繁地去主机端内存，DRAM大多是在对缓存的冷错过期间被访问，提供了比GPU和TPU平台更高的功率效率。
图14.在不同链路带宽下评估ProSE超过NVIDIAAGPU和TPUv3的功率效率
5、结论
随着人工智能可实现效益的增加，当务之急是提高迁移学习的效率，不仅是自然语言处理，还有药物设计和发现。为了提高生产更准确的候选药物的性能、功率和面积效率，我们提出了ProSE，一个用于蛋白质发现的收缩引擎。通过更换被加速的转化器模型权重（例如，增加语言翻译的解码器层）或增加不同的微调下游模型，ProSE很容易适用于其他众多的蛋白质和NLP相关的任务。
通过ProSE，我们提出了一个功率和面积高效的设计，即不同大小和数量的输出静止流异质脉动阵列的集合。这些脉动阵列配备了执行SIMDALU操作、SIMDGELU函数和SIMDExp函数的能力。ProSE执行ProteinBERT推断的速度比一个NVIDIAAGPU提高了6.9倍，比最新的GPU和TPU平台提高了两个数量级的功率效率。我们的工作表明，从跨模型层的共同计算模式中找出优化机会，可以在不牺牲通用性的情况下实现专业化的效率提升。