
1. 项目概述当随机计算遇上Vision Transformer在边缘AI和端侧部署的浪潮下我们这些搞硬件加速的工程师每天都在和功耗、面积、延迟这几个“硬骨头”较劲。传统的二进制计算虽然精度高但乘法器、加法器这些单元又大又耗电尤其是在处理像Vision TransformerViT这种参数和计算量都巨大的模型时硬件开销常常让人望而却步。这几年随机计算Stochastic Computing, SC作为一种“剑走偏锋”的计算范式重新回到了我们的视野。它的核心思想很巧妙不再用固定的0和1序列表示一个精确数值而是用一个随机比特流中“1”出现的概率来代表一个数值范围。比如一个8位的比特流中有6个1那么它代表的概率值就是6/80.75。这样一来复杂的乘法运算可以用一个简单的与门AND来实现加法运算则可以通过一个多路选择器MUX来完成硬件电路得到了极大的简化。听起来很美对吧但理想很丰满现实很骨感。当大家兴冲冲地想用随机计算去加速最新的ViT模型时却撞上了南墙。ViT模型里的两个关键非线性函数——GELU高斯误差线性单元激活函数和Softmax归一化指数函数——成了“拦路虎”。GELU函数在输入为负时是非单调的而传统的随机计算电路大多只能高效处理单调函数。Softmax则更麻烦它内部包含指数和除法运算这在随机计算里是实现起来极其昂贵且误差巨大的操作。直接套用为卷积神经网络CNN设计的随机计算方案要么精度崩盘要么硬件效率惨不忍睹。我最近深入研究了一篇来自北京大学团队的工作他们提出了一个名为ASCEND的框架算是真正捅破了这层窗户纸。ASCEND的核心思路是“协同设计”Co-design不再是简单地把训练好的ViT模型“硬塞”进随机计算电路里而是同时从电路设计和网络训练两个层面动手术让模型和硬件互相适应、互相优化。这篇文章我就结合自己的工程经验为你深度拆解ASCEND是如何解决这些挑战的并分享其中那些在论文图表之外、真正决定项目成败的设计细节和实操思考。2. 核心挑战拆解为什么ViT让随机计算“水土不服”在动手设计任何加速器之前搞清楚“敌人”的弱点和我们手中的“武器”的局限是成功的第一步。对于ASCEND要解决的难题我们可以从两个层面来看一是ViT模型本身对随机计算提出的新要求二是现有随机计算技术的能力边界。2.1 ViT模型中的“硬骨头”GELU与SoftmaxViT的核心是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电路可以看作一个“选择器”根据输入值从输入比特流中选择特定的几位直接作为输出。这必然导致输出是输入的单调函数。门辅助SIGate-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]变为1s[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。通过这种方式仅仅增加几个门电路我们就让原本只能单调递增的SI电路拥有了实现任意形状函数包括GELU的能力。而且由于它依然是并行、确定性的计算无需长比特流在面积和延迟上都有巨大优势。论文数据显示相比需要1024位BSL的伯恩斯坦多项式方案8位BSL的门辅助SI在面积-延迟积ADP上减少了5.29倍同时计算误差还降低了56.3%。3.2 迭代近似Softmax化指数和除法为乘加Softmax的挑战在于指数和除法。ASCEND采用了一个非常聪明的策略避开在随机计算域直接实现这两个棘手操作而是用一个迭代近似算法将它们转化为一系列乘法和加法。这个算法的精髓在于它将Softmax看作一个关于参数t的函数y(t) softmax(t * x)。当t0时y(0) 1/mm是向量维度所有元素相等。当t1时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_x2和B_x4的场景进行了详尽的探索分别找到了12个和21个帕累托Pareto最优设计点。什么是帕累托最优就是在这些设计点上你无法在不损害其中一个指标如精度的情况下去改善另一个指标如面积。这些点连成了一条“帕累托前沿”曲线。作为系统设计者我们的任务就是根据目标应用对精度和能效的要求在这条曲线上选择合适的点。5.2 硬件评估与配置选择论文给出了具体的评估数据非常具有参考价值GELU模块8位BSL的门辅助SI设计其ADP仅为1420 um²·ns而6项伯恩斯坦多项式设计需要1024位BSL的ADP高达7506 um²·ns前者是后者的约1/5误差还更低。Softmax模块当B_y8时其ADP为2.62e6 um²·ns比需要256位BSL的FSM方案ADP8.28e6 um²·ns降低了约3.2倍误差降低了25.7%。整体加速器Softmax模块在整个加速器中的面积占比很小。例如在一个配置为[B_y8, s132, s28, k3]的设计中Softmax模块面积约占加速器总面积的2%。即使将B_y增加到32以追求更高精度Softmax模块面积会激增30多倍导致整体加速器面积翻倍但精度提升仅约1.5%。给我的启示是在边缘设备上我们往往追求极致的能效比。论文中[B_y8, s132, s28, k3]这个配置是一个非常好的甜点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的核心计算单元并探索将其与更主流的线性注意力或动态稀疏注意力机制结合的可能性这或许是推动这项技术走向更广泛应用的关键。