TRAM:为低功耗 AI 加速器训练近似乘法器结构

arXiv cs.LG 论文

摘要

本文介绍了 TRAM,一种联合优化近似乘法器结构和 AI 模型参数的方法,旨在降低 AI 加速器的功耗,同时保持模型精度。

arXiv:2605.08231v1 公告类型:新提交 摘要:降低 AI 加速器的功耗日益重要。近似计算可以在保持精度损失较小的同时降低功耗。由于乘法器是 AI 模型中的高功耗组件,本文致力于综合设计低功耗近似乘法器(AxM)。与以往将 AxM 设计与 AI 模型训练分离的工作不同,我们提出了 TRAM,该方法通过联合优化 AxM 结构和 AI 模型参数,在精度损失较小的情况下降低功耗。实验表明,与最先进的 AxM 相比,TRAM 在 CIFAR-10 上的卷积神经网络(CNN)中实现了高达 25.05% 的 AxM 功耗降低,并在 ImageNet 上的视觉变换器(Vision Transformers)中降低了高达 27.09% 的功耗。
查看原文
查看缓存全文

缓存时间: 2026/05/12 07:09

# TRAM:为低功耗 AI 加速器训练近似乘法器结构

来源:https://arxiv.org/html/2605.08231
,Hanyu Wang 加州大学洛杉矶分校 洛杉矶 美国 hanyuwang@g\.ucla\.edu (https://arxiv.org/html/2605.08231v1/mailto:[email protected]), Yuyang Ye 香港中文大学 香港 中国 yuyangye@cuhk\.edu\.hk (https://arxiv.org/html/2605.08231v1/mailto:[email protected]), Mingfei Yu 洛桑联邦理工学院 洛桑 瑞士 mingfei\.yu@epfl\.ch (https://arxiv.org/html/2605.08231v1/mailto:[email protected]), Wayne Burleson 马萨诸塞大学阿默斯特分校 阿默斯特 美国 burleson@umass\.edu (https://arxiv.org/html/2605.08231v1/mailto:[email protected]) 和 Giovanni De Micheli 洛桑联邦理工学院 洛桑 瑞士 giovanni\.demicheli@epfl\.ch (https://arxiv.org/html/2605.08231v1/mailto:[email protected])

###### 摘要\.

降低 AI 加速器的功耗日益重要。近似计算可以在保持精度损失较小的同时降低功耗。由于乘法器是 AI 模型中的高功耗组件,本文专注于合成低功耗近似乘法器(AxMs)。与以往将 AxM 设计与 AI 模型训练分开进行的工作不同,我们提出了 TRAM,它联合优化 AxM 结构和 AI 模型参数,以在较小的精度损失下降低功耗。实验表明,与最先进的 AxM 相比,TRAM 在 CIFAR-10 上的卷积神经网络(CNN)中实现了高达 25.05% 的 AxM 功耗降低,并在 ImageNet 上的视觉变压器(Vision Transformers)中将功耗降低了高达 27.09%。

近似乘法器,软硬件协同优化,低功耗,AI 加速器

††copyright:acmcopyright††conference:; ;## 1\. 引言

AI 加速器的广泛部署引发了对功耗的关注,并产生了对低功耗计算方案的迫切需求(Schwartz 等人,2020 (https://arxiv.org/html/2605.08231#bib.bib49))。近似计算通过允许计算中的不准确性来降低功耗,成为解决这些问题的有前景的方法(Leone 等人,2025 (https://arxiv.org/html/2605.08231#bib.bib27))。由于乘法器是 AI 加速器中功耗最高的组件之一(Armeniakos 等人,2022 (https://arxiv.org/html/2605.08231#bib.bib51)),本文研究了低功耗近似乘法器(AxMs)的自动综合。

许多研究已经调查了 AxMs 的自动综合和手动设计(Wu 等人,2024 (https://arxiv.org/html/2605.08231#bib.bib35))。例如,Mrazke 等人(Mrazke 等人,2017 (https://arxiv.org/html/2605.08231#bib.bib16))提出了一种基于遗传编程的方法来综合 AxM,随后将其扩展到卷积神经网络(CNNs)(Mrazke 等人,2020 (https://arxiv.org/html/2605.08231#bib.bib33))。Xiao 等人(Xiao 等人,2022 (https://arxiv.org/html/2605.08231#bib.bib4))将 AxM 综合表述为整数规划问题,并生成了低成本 AxM。Hu 等人(Hu 等人,2024 (https://arxiv.org/html/2605.08231#bib.bib11))使用部分积推测手动设计了用于 CNN 的 AxM。此外,近似逻辑综合工具,如(Wang 等人,2023 (https://arxiv.org/html/2605.08231#bib.bib23); Ma 等人,2021 (https://arxiv.org/html/2605.08231#bib.bib30); Men 等人,2026 (https://arxiv.org/html/2605.08231#bib.bib28))中提出的工具,也可以综合 AxM。

然而,上述所有方法都忽视了 AI 模型的具体上下文,这可能导致在加速器中部署 AxM 时结果次优。首先,许多现有方法不考虑 AI 模型使用的数据分布。例如,Xiao 等人(Xiao 等人,2022 (https://arxiv.org/html/2605.08231#bib.bib4))假设输入分布是均匀的,而实际数据分布因层而异。忽略这种可变性可能导致设计次优。其次,大多数现有工作使用局部误差指标(如误差率或误差距离)来设计或综合 AxM。然而,较小的局部误差并不总是转化为 AI 模型中较小的最终精度损失。

Refer to captionFigure 1\. 4 位无符号阵列乘法器。红色叉号表示可能被近似的候选信号。结构参数 θi\\theta\_{i} 控制第 ii 累加列的近似程度。我们假设最多 P=3P=3 列可以被近似。Refer to captionFigure 2\. TRAM 框架概述。为了解决这些问题,我们提出了 TRAM,这是一个软硬件协同优化框架,用于为低功耗 AI 加速器 **T**rain **R**esource-efficient **A**pproximate **M**ultiplier 结构。TRAM 将 AxM 综合表述为一个联合优化问题,在训练期间同时更新 AxM 结构和 AI 模型参数。通过使用真实训练数据,TRAM 捕获每个乘法器看到的统计信息,并根据最终精度损失优化乘法器结构和模型参数。我们的贡献总结如下:

- • 我们引入了 AxM 结构的参数化,其中压缩树中的每一列都被分配一个连续的结构参数,以控制该列的近似程度。这些参数使用梯度下降法进行优化。
- • 我们设计了一个解析功耗模型,从结构参数估计乘法器功耗,并在训练期间提供有用的硬件感知指导。
- • 我们提出了一种高效的映射方法,将优化后的结构参数转换为具体的 AxM 设计。

实验结果表明,与最先进的 AxM 设计相比,TRAM 在保持相同精度水平的情况下,将 CIFAR-10 上 CNN 的 AxM 功耗降低了高达 25.05%,将 ImageNet 上视觉变压器的功耗降低了 27.09%。由于 TRAM 允许不同模型层使用不同的结构参数,它自然支持在不同层应用不同的 AxM。与最先进的逐层 AxM 探索方法相比,TRAM 将 AxM 能耗降低了 40.86%。我们的工作开源,可在 https://github.com/changmg/TRAM 获取。

本文的其余部分组织如下。第 2 节 (https://arxiv.org/html/2605.08231#S2) 描述预备知识。第 3-5 节 (https://arxiv.org/html/2605.08231#S3)–(https://arxiv.org/html/2605.08231#S5) 详细介绍 TRAM 框架。第 6 节 (https://arxiv.org/html/2605.08231#S6) 讨论实验结果。第 7 节 (https://arxiv.org/html/2605.08231#S7) 总结本文。

## 2\. 整数乘法器预备知识

本文重点关注在 AI 加速器中广泛使用的*无符号整数乘法器*(Simone 等人,2021 (https://arxiv.org/html/2605.08231#bib.bib9); Jain 等人,2022 (https://arxiv.org/html/2605.08231#bib.bib10); Men 等人,2025a (https://arxiv.org/html/2605.08231#bib.bib12); Zhen 等人,2022 (https://arxiv.org/html/2605.08231#bib.bib3))。此后,我们将无符号整数乘法器简称为乘法器。AxMs 通常是通过修改精确乘法器(AccMuls)获得的。一个 BB 位 AccMul 计算两个无符号整数输入 WW 和 XX 的精确乘积,它们以二进制表示为 W=wB−1wB−2...w0W=w\_{B-1}w\_{B-2}\\ldots w\_{0} 和 X=xB−1xB−2...x0X=x\_{B-1}x\_{B-2}\\ldots x\_{0}。

乘法器包含 2B2B 个累加列。第 cc 列累加部分积为 Sc=∑i=0cppi,c−iS\_{c}=\\sum\_{i=0}^{c}pp\_{i,c-i},其中 ppi,j=wi⋅xjpp\_{i,j}=w\_{i}\\cdot x\_{j} 是 wiw\_{i} 和 xjx\_{j} 的部分积,0≤c≤2B−10\\leq c\\leq 2B-1 是列索引。最终乘积通过对所有列的加权累加结果求和得到,即 Y=∑c=02B−1Sc⋅2cY=\\sum\_{c=0}^{2B-1}S\_{c}\\cdot 2^{c}。例如,图 1 (https://arxiv.org/html/2605.08231#S1.F1) 显示了一个具有 8 个累加列的 4 位阵列乘法器。每一列生成分部积并使用半加器(HAs)和全加器(FAs)对其进行累加。可以将近似引入这些列中的部分积或半加器和全加器的求和及进位输出信号中,如图 1 (https://arxiv.org/html/2605.08231#S1.F1) 中的红色叉号所示。

为了评估 BB 位 AxM 的精度,常见的误差指标包括误差率(ER)、归一化平均误差距离(NMED)和最大误差距离(MaxED)(Jian 等人,2020 (https://arxiv.org/html/2605.08231#bib.bib46)),定义为

ER=∑1≤i≤22B:Y\(i\)≠Yacc\(i\)pi,NMED=∑i=122B\|Y\(i\)−Yacc\(i\)\|⋅pi22B−1,\{\\textit{ER}\}=\\sum\_{1\\leq i\\leq 2^{2B}:Y^{(i)}\\neq Y_{\\textit{acc}}^{(i)}}{p_{i}},\\quad\{\\textit{NMED}\}=\\sum\_{i=1}^{2^{2B}}{\\frac{\\left\\|Y^{(i)}-Y_{\\textit{acc}}^{(i)}\\right\\|\\cdot p_{i}}{2^{2B}-1}},\\quad MaxED=max1≤i≤22B⁡\|Y\(i\)−Yacc\(i\)\|\.\{\\textit{MaxED}\}=\\max\_{1\\leq i\\leq 2^{2B}}{\\left\\|Y^{(i)}-Y_{\\textit{acc}}^{(i)}\\right\\|}\.

其中 Y\(i\)Y^{(i)} 和 Yacc\(i\)Y_{\\textit{acc}}^{(i)} 分别是 AxM 和 AccMul 在第 ii 输入组合下的输出,pip_{i} 是第 ii 输入组合的概率,22B2^{2B} 是输入组合的总数。

## 3\. TRAM 概述与乘法器结构参数化

Refer to captionFigure 3\. 计算公式 (2 (https://arxiv.org/html/2605.08231#S4.E2)) 中目标函数的数据流。上半部分计算功耗损失 Lpower\\mathcal{L}\_{\\textit{power}}(第 4.2 节 (https://arxiv.org/html/2605.08231#S4.SS2)),下半部分计算 AI 模型损失 LAI\_model\\mathcal{L}\_{\\textit{AI\\_model}}(第 4.3 节 (https://arxiv.org/html/2605.08231#S4.SS3))。### 3.1\. TRAM 框架概述

TRAM 旨在为 AI 加速器生成低功耗 AxM。TRAM 的整体流程如图 2 (https://arxiv.org/html/2605.08231#S1.F2) 所示。它从一个预训练的浮点 AI 模型开始,然后将其量化为整数模型。为了进一步降低功耗,AxMs 替换量化模型中的 AccMuls。为了探索 AxM 设计空间,我们使用结构参数 Θ\\Theta 来表示 AxM 结构。更改 AxM 结构对应于更新 Θ\\Theta。Θ\\Theta 的详细参数化在第 3.2 节 (https://arxiv.org/html/2605.08231#S3.SS2) 中介绍。基于此参数化,我们提出了一种三阶段方法来为高精度 AI 模型生成低功耗 AxM:

第一阶段。设计空间探索(详情见第 4 节 (https://arxiv.org/html/2605.08231#S4))。此阶段通过模型重新训练探索由 Θ\\Theta 定义的 AxM 设计空间,并平衡功耗与精度。

第二阶段。AxM 结构映射(详情见第 5 节 (https://arxiv.org/html/2605.08231#S5))。此阶段将第一阶段中优化后的连续结构参数 Θ∗\\Theta^{\*} 映射到模型每一层的具体 AxM 结构。

第三阶段。精度恢复。在第二阶段将结构参数映射到 AxM 结构后,我们将这些 AxM 应用于 AI 模型并重新训练以恢复精度。

### 3.2\. AxM 结构参数化

我们使用连续结构参数对 AxM 结构进行参数化,从而通过模型重新训练实现基于梯度的 AxM 结构优化。令 Θ={Θ\(l\)\}(1≤l≤L)\\Theta{=}\\{\\Theta^{(l)}\\}(1{\\leq}l{\\leq}L) 表示模型中所有 LL 层结构参数的集合。我们假设每一层对该层中的所有乘法使用一个 AxM 结构,由 Θ\(l\)\\Theta^{(l)} 描述。Θ\(l\)\\Theta^{(l)} 有 PP 个参数:Θ\(l\)=[θ0\(l\),θ1\(l\),...,θP−1\(l\)]\\Theta^{(l)}{=} [\\theta^{(l)}\_0, \\theta^{(l)}\_1, \\ldots, \\theta^{(l)}\_{P-1}](见图 1 (https://arxiv.org/html/2605.08231#S1.F1)),其中 θc\(l\)\\theta^{(l)}\_c (0≤c≤P−10{\\leq}c{\\leq}P{-}1) 描述列 cc 的近似程度。这里,PP 是用户定义的可近似列的最大数量。

每个 θc\(l\)∈[0,1]\\theta^{(l)}\_c\\in[0,1] 是一个连续结构参数,控制第 ll 层 AxM 中列 cc 的近似程度。值 0 表示该列保持完全精确,而值 1 表示该列完全移除。中间值 0<θc\(l\)<10{<}\\theta^{(l)}\_c{<}1 表示部分近似,其中仅移除列 cc 中部分积或压缩器的子集。

接下来,我们解释结构参数如何控制 AxM 的功能行为。对于第 ll 层,第 cc 个累加列的近似误差定义为 Ec=θc\(l\)⋅Sc=θc\(l\)⋅∑i=0cppi,c−iE\_{c}=\\theta^{(l)}\_c\\cdot S\_{c}=\\theta^{(l)}\_c\\cdot\\sum\_{i=0}^{c}pp\_{i,c-i},其中 ScS\_{c} 是列 cc 的精确累加结果。所有 PP 个近似列上的总近似误差通过将列误差乘以其权重 2c2^{c} 并求和来计算,即 Etotal=∑c=0P−1Ec⋅2cE\_{total}=\\sum\_{c=0}^{P-1}E\_{c}\\cdot 2^{c}。然后通过从精确乘积中减去该误差来获得 AxM 输出:

\(1\)Y=WX−Etotal=WX−∑c=0P−1θc\(l\)⋅Sc⋅2c,Y=WX-E\_{total}=WX-\\sum\_{c=0}^{P-1}\\theta^{(l)}\_c\\cdot S\_{c}\\cdot 2^{c},

其中 WWWW 是 WW 和 XX 的精确乘积。公式 (1 (https://arxiv.org/html/2605.08231#S3.E1)) 允许通过在范围 [0,1][0,1] 内变化结构参数 θc\(l\)\\theta^{(l)}\_c 来平滑调整每一列 cc 的近似程度。较大的 θc\(l\)\\theta^{(l)}\_c 会导致较大的近似误差并降低功耗,而较小的 θc\(l\)\\theta^{(l)}\_c 会减少误差并增加功耗。此公式可以扩展到阵列乘法器之外。

## 4\. 第一阶段:通过 AI 模型重新训练进行设计空间探索

### 4.1\. 问题表述

结构参数 Θ\\Theta 定义 AxM 设计空间。Θ\\Theta 的不同选择会导致具有不同功耗和不同 AI 模型精度的 AxM。为了平衡功耗和精度,我们制定以下优化问题:

\(2\)minΘ,W⁡(Lpower\(Θ\)⋅λ+LAI\_model\(Θ,W,X\))\.\\min\_{\\Theta,\\mathbf{W}}\\left(\\mathcal{L}\_{\\textit{power}}(\\Theta)\\cdot\\lambda+\\mathcal{L}\_{\\textit{AI\\_model}}(\\Theta,\\mathbf{W},\\mathbf{X})\\right)\.

公式 (2 (https://arxiv.org/html/2605.08231#S4.E2)) 包含两个损失项:功耗损失 Lpower\\mathcal{L}\_{\\textit{power}} 和 AI 模型损失 LAI\_model\\mathcal{L}\_{\\textit{AI\\_model}}。Lpower\\mathcal{L}\_{\\textit{power}} 将结构参数 Θ\\Theta 映射到 AI 加速器中所有 AxM 消耗的总功耗。LAI\_model\\mathcal{L}\_{\\textit{AI\\_model}} 是原始模型损失(例如,分类的交叉熵损失),取决于结构参数 Θ\\Theta、模型权重 W\\mathbf{W} 和输入 X\\mathbf{X}。引入权衡参数 λ\\lambda 以平衡这两个损失。增加 λ\\lambda 会给 Lpower\\mathcal{L}\_{\\textit{power}} 赋予更多权重,从而降低功耗但增加模型损失(即,模型精度降低)。通过调整 λ\\lambda,我们可以探索不同的功耗-精度权衡。

我们通过模型重新训练来解决公式 (2 (https://arxiv.org/html/2605.08231#S4.E2)) 中的优化问题。在重新训练期间,计算公式 (2 (https://arxiv.org/html/2605.08231#S4.E2)) 中目标函数的数据流如图 3 (https://arxiv.org/html/2605.08231#S3.F3) 所示。上半部分计算 Lpower\\mathcal{L}\_{\\textit{power}},下半部分计算 LAI\_model\\mathcal{L}\_{\\textit{AI\\_model}}。第 4.2 节 (https://arxiv.org/html/2605.08231#

相似文章

Ada-MK:基于自动化 DAG 搜索的 LLM 推理自适应 MegaKernel 优化

arXiv cs.CL

本文介绍了 Ada-MK,一种利用自动化基于有向无环图(DAG)的搜索来消除运行时分支并减少大语言模型(LLM)推理共享内存使用的自适应 MegaKernel 优化方法。通过集成到 TensorRT-LLM 中,该方法在 NVIDIA Ada GPU 上展示了显著的吞吐量提升,在商业广告系统中相比原生 TensorRT-LLM 性能最高提升 23.6%。