1. 项目概述:当随机计算遇上Vision Transformer
在边缘AI和端侧部署的浪潮下,我们这些搞硬件加速的工程师,每天都在和功耗、面积、延迟这几个“硬骨头”较劲。传统的二进制计算虽然精度高,但乘法器、加法器这些单元又大又耗电,尤其是在处理像Vision Transformer(ViT)这种参数和计算量都巨大的模型时,硬件开销常常让人望而却步。这几年,随机计算(Stochastic Computing, SC)作为一种“剑走偏锋”的计算范式,重新回到了我们的视野。它的核心思想很巧妙:不再用固定的0和1序列表示一个精确数值,而是用一个随机比特流中“1”出现的概率来代表一个数值范围。比如,一个8位的比特流中有6个1,那么它代表的概率值就是6/8=0.75。这样一来,复杂的乘法运算可以用一个简单的与门(AND)来实现,加法运算则可以通过一个多路选择器(MUX)来完成,硬件电路得到了极大的简化。
听起来很美,对吧?但理想很丰满,现实很骨感。当大家兴冲冲地想用随机计算去加速最新的ViT模型时,却撞上了南墙。ViT模型里的两个关键非线性函数——GELU(高斯误差线性单元)激活函数和Softmax(归一化指数函数)——成了“拦路虎”。GELU函数在输入为负时是非单调的,而传统的随机计算电路大多只能高效处理单调函数。Softmax则更麻烦,它内部包含指数和除法运算,这在随机计算里是实现起来极其昂贵且误差巨大的操作。直接套用为卷积神经网络(CNN)设计的随机计算方案,要么精度崩盘,要么硬件效率惨不忍睹。
我最近深入研究了一篇来自北京大学团队的工作,他们提出了一个名为ASCEND的框架,算是真正捅破了这层窗户纸。ASCEND的核心思路是“协同设计”(Co-design):不再是简单地把训练好的ViT模型“硬塞”进随机计算电路里,而是同时从电路设计和网络训练两个层面动手术,让模型和硬件互相适应、互相优化。这篇文章,我就结合自己的工程经验,为你深度拆解ASCEND是如何解决这些挑战的,并分享其中那些在论文图表之外、真正决定项目成败的设计细节和实操思考。
2. 核心挑战拆解:为什么ViT让随机计算“水土不服”?
在动手设计任何加速器之前,搞清楚“敌人”的弱点和我们手中的“武器”的局限,是成功的第一步。对于ASCEND要解决的难题,我们可以从两个层面来看:一是ViT模型本身对随机计算提出的新要求,二是现有随机计算技术的能力边界。
2.1 ViT模型中的“硬骨头”:GELU与Softmax
ViT的核心是Transformer编码器块,其中两个关键组件带来了计算上的特殊性:
GELU激活函数:这是ViT中MLP(多层感知机)块使用的激活函数。它的公式是
GELU(x) = x * Φ(x),其中Φ(x)是标准高斯分布的累积分布函数。它的曲线特点是:当x为较大的负数时,输出接近0;当x为0时,输出也为0;当x为正数时,输出近似于x。关键在于,它在x为负的小区间内(大约-1到0之间)有一个先下降后上升的“非单调”区域。现有的随机计算电路,如基于有限状态机(FSM)或选择性互连(SI)的设计,要么无法处理非单调性,要么需要极长的比特流来逼近,导致效率低下。Softmax函数:这是ViT中多头自注意力(MSA)块的核心,用于计算注意力权重。公式为
Softmax(x_i) = exp(x_i) / Σ exp(x_j)。它带来了双重挑战:- 指数运算(exp):指数函数对输入值的变化极其敏感,需要很高的数值表示精度。在随机计算中,精度直接关联到比特流长度(BSL)。要精确表示指数运算的输出范围,所需的BSL长度会呈指数级增长,严重拖累硬件效率。
- 除法运算:在随机计算中实现精确除法是出了名的困难。常见的方法如
p_x / (p_x + p_y)只能实现近似,且需要复杂的电路(如带JK触发器的随机数生成器或比较器),不仅面积大,还会引入显著的随机波动误差。
2.2 随机计算的传统武器库及其局限
在ASCEND之前,社区为非线性函数设计随机计算电路主要有三种思路,但面对GELU和Softmax都力不从心:
- 基于有限状态机(FSM)的设计:通过设计状态转移条件来逼近函数。问题是,它是串行处理比特流的,为了降低随机波动误差,需要非常长的比特流(论文中提到需要1024位以上),导致计算延迟巨大,完全不适合对延迟敏感的ViT推理。
- 基于伯恩斯坦多项式的设计:用多项式来拟合非线性函数。但为了拟合像GELU这样复杂的函数,需要高阶多项式,这意味着需要大量的计算项和随机数生成器,硬件开销激增。同时,它也无法避免随机波动。
- 基于选择性互连(SI)的设计:这是目前用于确定性温度计编码的最高效方案,它能并行处理整个比特流,精度高且延迟低。但它有一个致命缺点:只能实现单调函数。因为它的工作原理是,输出比特流中“1”的数量会随着输入比特流中“1”的数量增加而单调增加。这对于ReLU、Sigmoid没问题,但对非单调的GELU就无能为力了。
实操心得:理解编码格式是关键这里必须提一下随机计算中不同的“编码格式”,它直接决定了底层电路如何构建。温度计编码(Thermometer Encoding)是ASCEND采用的核心技术。假设我们用4位比特流表示数值,那么数值“-2”可能编码为
0011(前两位是1),数值“0”编码为1100,数值“+2”编码为1111。它的好处是确定性的,没有随机波动,并且乘法和加法可以用查找表或排序网络高效实现。理解你所用编码格式的数学本质和硬件代价,是设计高效SC电路的第一步。
3. 电路层创新:为ViT量身定制的随机计算引擎
ASCEND在电路层面的贡献是决定性的,它没有沿用旧思路,而是针对GELU和Softmax的特性,设计了全新的计算单元。
3.1 门辅助选择性互连:攻克非单调GELU
面对SI电路只能处理单调函数的限制,ASCEND的解决方案非常巧妙,我称之为“给SI电路加一个智能开关”。传统的SI电路可以看作一个“选择器”:根据输入值,从输入比特流中选择特定的几位,直接作为输出。这必然导致输出是输入的单调函数。
门辅助SI(Gate-Assisted SI)的核心思想是:我们不直接输出选中的比特,而是让选中的比特经过一个简单的组合逻辑门(如与门、非门)后再输出。通过精心设计这个“辅助逻辑”和选择信号,我们就能让输出比特流中“1”的数量呈现先减后增的非单调变化,从而精确匹配GELU函数的形状。
具体是如何实现的呢?我们以一个简化的“三元GELU”为例(输出只有-1, 0, +1三个值)。假设输入是8位温度计编码比特流x[7:0],输出是2位比特流y[1:0]。
- 电路结构:核心是一个“选择器”,它根据输入x的值,生成3个选择信号
s[2:0]。这3个信号分别连接到输入比特流的某一位上(例如,s[2]连x[7], s[1]连x[5], s[0]连x[3])。当输入x很小时,所有x[i]都是0,因此s[2:0]也都是0。 - 辅助逻辑:输出
y[1]被设计为!(s[2]) & s[1](即“非s[2]与s[1]”),y[0]直接等于s[0]。 - 工作过程:
- 当x很小(负值)时,
s[2:0]=000,则y[1:0] = (!0 & 0, 0) = (1 & 0, 0) = 00,对应输出值-1(在双极编码下)。 - 随着x增大,
x[7]先变为1,导致s[2]=1。此时y[1:0] = (!1 & s[1], 0) = (0 & s[1], 0) = 00,输出仍为-1。 - x继续增大,
x[5]变为1,s[1]=1。此时y[1:0] = (!1 & 1, 0) = (0 & 1, 0) = 00,输出还是-1?等等,这里需要看具体GELU的量化映射表。实际上,设计者会根据目标函数,精心选择s信号连接的输入位和辅助逻辑,使得在s[2]=1, s[1]=0时,输出恰好为0(例如y[1:0]=10)。而当s[2]=1, s[1]=1时,y[1]又变回0,输出可能回到-1或0,以此模拟下降沿。 - 最终,当x为正且较大时,
s[0]也变为1,使得y[0]=1,输出变为+1(例如y[1:0]=11)。
- 当x很小(负值)时,
通过这种方式,仅仅增加几个门电路,我们就让原本只能单调递增的SI电路,拥有了实现任意形状函数(包括GELU)的能力。而且,由于它依然是并行、确定性的计算,无需长比特流,在面积和延迟上都有巨大优势。论文数据显示,相比需要1024位BSL的伯恩斯坦多项式方案,8位BSL的门辅助SI在面积-延迟积(ADP)上减少了5.29倍,同时计算误差还降低了56.3%。
3.2 迭代近似Softmax:化指数和除法为乘加
Softmax的挑战在于指数和除法。ASCEND采用了一个非常聪明的策略:避开在随机计算域直接实现这两个棘手操作,而是用一个迭代近似算法将它们转化为一系列乘法和加法。
这个算法的精髓在于,它将Softmax看作一个关于参数t的函数y(t) = softmax(t * x)。当t=0时,y(0) = 1/m(m是向量维度,所有元素相等)。当t=1时,y(1)就是我们想要的softmax(x)。那么,y(1)可以通过对y(t)从0到1积分来近似。通过数学推导(具体可参考论文),这个积分可以转化为一个迭代更新公式:
对于第j次迭代(总共k次):
- 计算中间向量
z = x * y^(j-1)(逐元素乘法) - 计算标量
sum_z = sum(z)(求和) - 更新输出向量
y^(j) = y^(j-1) + (z - y^(j-1) * sum_z) / k
这个公式的妙处在哪里?
- 消除指数和除法:原始Softmax的
exp(x_i)和除以Σexp(x_j)的操作不见了。 - 仅剩乘加和常数除法:迭代中只剩下向量点乘、标量乘向量、向量加法,以及除以一个常数k。在随机计算中,乘法和加法是高效操作,而除以常数k只需要调整输出结果的缩放因子,无需任何实际电路操作!
- SC友好:整个计算流程完全由乘法和加法构成,完美契合随机计算中利用与门、多路选择器和排序网络实现乘加的操作模式。
基于这个算法,ASCEND设计了对应的并行电路块。它包含m个相同的处理单元(对应m维向量),每个单元负责计算一个y_i。电路核心是两个乘法器阵列和一个全局的比特onic排序网络(BSN,用于高效实现加法)。通过k次迭代后,即可得到近似的Softmax结果。通过调整迭代次数k和内部比特流长度等参数,可以在计算精度和硬件开销之间进行灵活的权衡。
4. 网络层协同:训练出“硬件友好”的低精度ViT
优秀的电路设计解决了“算得了”的问题,但要“算得好”、“算得准”,还需要模型本身的配合。这就是ASCEND“协同设计”思想的另一半:通过改进训练过程,让ViT模型主动适应随机计算硬件的特性。
4.1 两阶段训练流水线
直接将一个全精度(FP32)的ViT模型量化到极低精度(如权重和激活都用2位比特流表示),精度会暴跌。ASCEND提出了一个循序渐进的两阶段训练流程:
第一阶段:渐进式量化(Progressive Quantization)这一步的目标是让模型平稳地“瘦身”到低精度,避免精度骤降。
- 起点:从一个训练好的全精度ViT模型开始。
- 第一步量化:先将权重(W)和激活(A)量化到16位,残差连接(R)保持16位(记为W16-A16-R16)。用知识蒸馏(KD),以原全精度模型为教师,训练这个中等精度模型。
- 第二步量化:在第一步模型的基础上,将激活进一步量化到2位(W16-A2-R16)。此时,教师模型换为第一步得到的W16-A16-R16模型。因为教师和学生的“能力差距”变小了,蒸馏效果更好。
- 第三步量化:最后,将权重也量化到2位,得到目标低精度模型(W2-A2-R16)。继续使用W16-A16-R16作为教师进行蒸馏。
这种“小步快跑”的量化策略,给了模型足够的调整空间,论文中显示,相比直接量化到W2-A2-R16,渐进式量化在CIFAR10和CIFAR100数据集上分别带来了32.99%和21.4%的精度提升。
第二阶段:近似Softmax感知微调(Approximate Softmax Aware Fine-Tuning)第一阶段得到的低精度模型,其内部使用的仍然是标准的、精确的Softmax函数。但我们的硬件实际实现的是迭代近似的Softmax,这之间存在一个“语义鸿沟”。直接替换会导致精度损失。
因此,在第二阶段,ASCEND做了一个关键操作:在训练好的低精度模型中,用我们设计的迭代近似Softmax电路的前向计算逻辑,替换掉原来的精确Softmax函数。然后,以这个“带近似计算模块的模型”为基础,进行一个短周期的微调(Fine-tuning)。微调的目标是让模型的所有参数去适应这个近似的、有误差的Softmax计算过程。
这个过程可以理解为“硬件在环训练”的一个简化版。模型在训练阶段就“见识”并学会了补偿硬件计算带来的误差,从而在部署到真实的ASCEND硬件上时,能保持更高的精度。实验表明,这一阶段能进一步挽回约1.5%的精度损失。
4.2 层归一化的替代:一项容易被忽略的优化
在ViT的原始设计中,每个块后使用的是层归一化(Layer Normalization, LN)。LN需要对同一层内所有特征点的均值和方差进行计算,这在随机计算中实现起来比较麻烦。ASCEND论文中提到,他们将LN替换为了批归一化(Batch Normalization, BN)。
为什么这么做?BN在推理时是固定的缩放和平移操作,可以融合到前一层的线性变换中,或者用非常简单的电路实现。而LN是动态计算的,依赖于当前输入的数据统计。在硬件上,BN比LN更简单、更高效。当然,直接替换通常会导致精度下降。ASCEND同样借助知识蒸馏,在训练低精度模型时就用BN替代LN,并通过蒸馏损失来最小化这一改变带来的影响,最终实现了精度损失小于0.1%,这是一个非常成功的工程折中。
注意事项:协同设计的核心是联合优化很多硬件加速项目失败的原因在于“先训模型,再设计硬件”的割裂流程。ASCEND的成功启示我们,必须将硬件约束(如支持的操作类型、数值精度、非线性函数近似方式)作为先验知识,反馈到模型训练或微调阶段。在设计初期,硬件工程师和算法工程师就需要紧密沟通,定义好“硬件-软件”接口(即模型中可以定制的近似计算模块),才能实现整体最优。
5. 设计空间探索与硬件评估权衡
有了创新的电路和训练方法,下一步就是如何将它们组合成一个高效的加速器,并在精度、面积、延迟之间做出最佳权衡。ASCEND在这方面提供了一个很好的范例。
5.1 软硬件协同的设计空间探索
ASCEND的迭代近似Softmax电路有多个可调参数,形成了一个丰富的设计空间:
B_x: 输入x的比特流长度。B_y: 内部及输出y的比特流长度。k: 迭代次数。s1,s2: 求和路径中的子采样率(用于进一步优化面积)。
这些参数共同决定了电路的计算精度(MAE)和硬件成本(面积、延迟,以及二者的乘积ADP)。论文中对B_x=2和B_x=4的场景进行了详尽的探索,分别找到了12个和21个帕累托(Pareto)最优设计点。
什么是帕累托最优?就是在这些设计点上,你无法在不损害其中一个指标(如精度)的情况下,去改善另一个指标(如面积)。这些点连成了一条“帕累托前沿”曲线。作为系统设计者,我们的任务就是根据目标应用对精度和能效的要求,在这条曲线上选择合适的点。
5.2 硬件评估与配置选择
论文给出了具体的评估数据,非常具有参考价值:
- GELU模块:8位BSL的门辅助SI设计,其ADP仅为1420 um²·ns,而6项伯恩斯坦多项式设计(需要1024位BSL)的ADP高达7506 um²·ns,前者是后者的约1/5,误差还更低。
- Softmax模块:当
B_y=8时,其ADP为2.62e6 um²·ns,比需要256位BSL的FSM方案(ADP=8.28e6 um²·ns)降低了约3.2倍,误差降低了25.7%。 - 整体加速器:Softmax模块在整个加速器中的面积占比很小。例如,在一个配置为
[B_y=8, s1=32, s2=8, k=3]的设计中,Softmax模块面积约占加速器总面积的2%。即使将B_y增加到32以追求更高精度,Softmax模块面积会激增30多倍,导致整体加速器面积翻倍,但精度提升仅约1.5%。
给我的启示是:在边缘设备上,我们往往追求极致的能效比。论文中[B_y=8, s1=32, s2=8, k=3]这个配置是一个非常好的甜点(Sweet Spot)。它在CIFAR10上实现了超过90%的精度,同时硬件开销增加非常有限。在实际项目中,我们需要用类似的评估方法,针对自己的数据集(如ImageNet)和精度要求,快速遍历设计空间,找到最适合自己芯片面积和功耗预算的配置。
6. 工程落地思考与潜在挑战
ASCEND为我们展示了一条可行的ViT随机计算加速路径,但在实际工程化过程中,还有一些问题需要仔细考量。
6.1 精度与通用性的权衡
ASCEND在CIFAR10/100这类相对较小的数据集上证明了有效性,但ViT的强大能力往往体现在ImageNet等大规模数据集上。低精度(如W2-A2)模型在大数据集上的精度保持能力,仍需更广泛的验证。此外,渐进式量化和近似感知微调的训练流程相对复杂,训练时间和计算资源成本会高于标准训练。
实操建议:在项目启动时,不要急于追求最低的2位精度。可以从较高的精度(如W4-A4)开始,作为基线验证整个协同设计流程的可行性。同时,可以探索混合精度策略,例如对注意力层中的Q、K、V矩阵使用稍高精度,而对其他部分使用低精度。
6.2 系统级集成与数据流优化
ASCEND论文主要聚焦于核心计算单元的设计。一个完整的加速器还需要考虑:
- 片上存储(SRAM):低精度模型虽然参数小,但ViT的注意力机制会产生大量的中间激活(特别是Key和Value矩阵)。如何高效地组织数据在存储层次(DRAM -> 全局缓存 -> 局部寄存器)中的流动,避免成为性能瓶颈,是关键。
- 控制与调度:需要设计一个高效的调度器,来管理多个处理单元(PE)并行计算多个注意力头,以及协调MSA块和MLP块之间的数据依赖。
- 随机数生成器(SNG)开销:虽然温度计编码本身是确定性的,但在一些接口或需要生成随机比特流的地方,高质量SNG的面积和功耗开销不容忽视。
6.3 与其它高效化技术的结合
随机计算不是孤立的,它可以与其他模型压缩和加速技术结合:
- 稀疏化:ViT的注意力矩阵通常是稀疏的。可以结合注意力剪枝,跳过对零值的计算。在随机计算中,零值有特定的比特流模式,检测并跳过这些操作可以节省功耗。
- 动态比特流长度:不同层、不同张量对精度的敏感度不同。可以设计一种动态机制,为不敏感的操作分配更短的比特流,进一步节省计算资源和时间。
ASCEND的工作为Vision Transformer的硬件高效部署打开了一扇新的大门。它深刻地体现了“算法-硬件协同设计”的威力:不是让硬件去将就算法,也不是让算法去将就硬件,而是让双方在同一个设计框架下互相妥协、互相优化,最终达到全局最优。对于从事AI芯片、边缘计算加速的工程师来说,理解并掌握这种协同设计的方法论,其价值可能比任何一个具体的电路设计技巧都更为重要。下一步,我计划在开源硬件平台上,尝试复现ASCEND的核心计算单元,并探索将其与更主流的线性注意力或动态稀疏注意力机制结合的可能性,这或许是推动这项技术走向更广泛应用的关键。