用于半导体热机械可靠性的递归Transformer

# 递归权重共享Transformer用于半导体热机械可靠性
来源：https://arxiv.org/html/2607.27251
###### 摘要

基于Transformer的代理模型越来越多地被用于替代工程设计中的昂贵的第一性原理仿真。然而，传统的Transformer架构通常针对工程设计空间中典型的小规模、低维数据集过参数化，因为大规模仿真数据生成成本高昂。在这种情况下，过多的参数容量会导致过拟合而非提高精度，同时还会带来不必要的内存和计算开销——这促使我们转向关注额外计算而非额外可学习参数的架构。本文对用于先进封装热机械代理分析的三种递归Transformer范式进行了硬件感知评估：a) Tiny Recursive模型，b) 我们提出的Depth Recursive Transformer，c) 以及简单递归Transformer。我们系统地比较了它们的预测性能（Recall@K、平均倒数排名）、参数数量、计算复杂度（FLOPs），为在资源受限场景下选择递归Transformer架构提供了实用设计指南。我们在两个低维工程预测任务上验证了这一原则：1) 先进半导体封装的热机械可靠性分析，其中热循环产生的应力和翘曲必须在昂贵的有限元分析（FEA）设计实验扫描中反复评估。2) 用于电容场的Laplace PDE迭代数值求解器。总体而言，递归权重共享Transformer为小数据工程代理建模提供了一种有效且可泛化的预测精度、参数效率和计算成本之间的权衡——这一场景传统大参数Transformer并不适用——正如先进封装可靠性预测和电容器静电场建模所证明的那样。

## I 引言

Transformer模型在广泛的工程应用中展示了强大的预测性能，包括电子设计自动化（EDA）[7](https://arxiv.org/html/2607.27251#bib.bib8)、科学机器学习[16](https://arxiv.org/html/2607.27251#bib.bib17)和数字孪生[25](https://arxiv.org/html/2607.27251#bib.bib25)。然而，它们日益增长的计算和内存需求对在硬件受限环境中的部署提出了重大挑战，例如片上推理、边缘加速器和实时设计优化[22](https://arxiv.org/html/2607.27251#bib.bib23)、[2](https://arxiv.org/html/2607.27251#bib.bib3)。因此，大量研究集中于通过算法、架构和系统层面的改进来减少Transformer计算[3](https://arxiv.org/html/2607.27251#bib.bib5)、[8](https://arxiv.org/html/2607.27251#bib.bib9)、[5](https://arxiv.org/html/2607.27251#bib.bib6)、[11](https://arxiv.org/html/2607.27251#bib.bib12)。热机械仿真仍计算成本高昂，这促使使用基于Transformer的代理模型进行快速可靠性预测。然而，在实际EDA工作流中部署此类模型需要低计算复杂度、紧凑内存占用和高推理吞吐量的架构。这些需求催生了广泛的Transformer计算缩减技术，可根据其优化的计算方面大致分为七个互补范式，包括高效注意力[14](https://arxiv.org/html/2607.27251#bib.bib15)、[3](https://arxiv.org/html/2607.27251#bib.bib5)，Token缩减[21](https://arxiv.org/html/2607.27251#bib.bib21)、[1](https://arxiv.org/html/2607.27251#bib.bib2)，条件计算[8](https://arxiv.org/html/2607.27251#bib.bib9)、[6](https://arxiv.org/html/2607.27251#bib.bib7)，架构优化[5](https://arxiv.org/html/2607.27251#bib.bib6)、[15](https://arxiv.org/html/2607.27251#bib.bib16)、[24](https://arxiv.org/html/2607.27251#bib.bib24)、[23](https://arxiv.org/html/2607.27251#bib.bib22)，模型压缩[11](https://arxiv.org/html/2607.27251#bib.bib12)、[12](https://arxiv.org/html/2607.27251#bib.bib13)、[13](https://arxiv.org/html/2607.27251#bib.bib14)，硬件优化[4](https://arxiv.org/html/2607.27251#bib.bib4)、[20](https://arxiv.org/html/2607.27251#bib.bib20)，以及状态空间模型[10](https://arxiv.org/html/2607.27251#bib.bib11)、[9](https://arxiv.org/html/2607.27251#bib.bib10)。尽管这些方法追求提高计算效率的共同目标，但它们针对不同的计算瓶颈，从降低注意力复杂度和序列长度到重新设计网络架构或完全替换基于注意力的模型。在这些范式中，本文聚焦于架构优化，即通过重新设计Transformer本身来实现计算效率。特别是，递归权重共享能够在参数不按比例增加的情况下实现更深的计算[5](https://arxiv.org/html/2607.27251#bib.bib6)、[24](https://arxiv.org/html/2607.27251#bib.bib24)、[23](https://arxiv.org/html/2607.27251#bib.bib22)，使其对资源受限的EDA应用特别有吸引力。尽管已有许多减少Transformer计算的方法被提出，但它们的适用性取决于目标应用。高效注意力机制主要惠及自注意力主导计算成本的长序列模型[14](https://arxiv.org/html/2607.27251#bib.bib15)、[3](https://arxiv.org/html/2607.27251#bib.bib5)，而token缩减技术依赖大token集来实现有意义的计算节省[21](https://arxiv.org/html/2607.27251#bib.bib21)、[1](https://arxiv.org/html/2607.27251#bib.bib2)。条件计算和混合专家架构提高了大型基础模型的可扩展性，但引入了额外的路由复杂性和内存开销[8](https://arxiv.org/html/2607.27251#bib.bib9)、[6](https://arxiv.org/html/2607.27251#bib.bib7)、[19](https://arxiv.org/html/2607.27251#bib.bib1)。模型压缩和硬件特定优化通常在模型设计后应用，或针对特定执行平台[11](https://arxiv.org/html/2607.27251#bib.bib12)、[13](https://arxiv.org/html/2607.27251#bib.bib14)、[4](https://arxiv.org/html/2607.27251#bib.bib4)。相比之下，我们的应用涉及在严格内存和计算约束下对短输入序列进行操作的紧凑Transformer代理。因此，架构优化提供了最合适的设计策略，能够通过递归权重共享从本质上减少计算，同时保持紧凑、硬件友好的架构[5](https://arxiv.org/html/2607.27251#bib.bib6)、[24](https://arxiv.org/html/2607.27251#bib.bib24)、[23](https://arxiv.org/html/2607.27251#bib.bib22)。

参见图1：Tiny Recursive模型 参见图2：Depth Recursive模型 参见图3：Simple Recursive模型
## II 递归权重共享Transformer

递归Transformer的主要概念是将Transformer视为一个可重用的模板（例如 RECUR\\operatorname{RECUR}）贯穿整个模型。模型中使用RECUR的多个副本，但它们共享同一组权重，该权重在每次训练迭代中更新一次。RECUR\\operatorname{RECUR}应独立于我们选择部署的模型。RECUR\\operatorname{RECUR}最简单的配置是方程1中的单块Transformer，由自注意力ATTN⁡(⋅)\\operatorname{ATTN}(\\cdot)、归一化层NORM⁡(⋅)\\operatorname{NORM}(\\cdot)、前馈网络（FFN）组成。反之，如果我们想要更深的Transformer，可以使用例如两个Transformer块来定义RECUR\\operatorname{RECUR}。

RECUR=\\displaystyle \\operatorname{RECUR}=\{\}ATTN⁡(⋅∣WQ,WK,WV,WO)\\displaystyle \\operatorname{ATTN}(\\cdot\\mid W_{Q},W_{K},W_{V},W_{O})(1)→ADD1+NORM1(⋅∣γ1,β1)\\displaystyle \\rightarrow\\operatorname{ADD}_{1}+\\operatorname{NORM}_{1}(\\cdot\\mid\\gamma_{1},\\beta_{1})→FFN(⋅∣W1,W2)\\displaystyle \\rightarrow\\operatorname{FFN}(\\cdot\\mid W_{1},W_{2})→ADD2+NORM2(⋅∣γ2,β2).\\displaystyle \\rightarrow\\operatorname{ADD}_{2}+\\operatorname{NORM}_{2}(\\cdot\\mid\\gamma_{2},\\beta_{2}).
### II-A Tiny Recursive Model

Tiny Recursive Model（TRM）使用权重共享的RECUR来限制参数数量，同时在图1所示的数独和谜题任务上执行AI推理。在每个迭代中，RECURZ\\operatorname{RECUR}_{Z}通过馈入输入嵌入E∈RL×dtoken=[e1,...,eL]\\mathbf{E}\\in\\mathbb{R}^{L\\times d_{token}}=[\\mathbf{e}_{1},\\ldots,\\mathbf{e}_{L}]（其中LL是token数量，dtokend_{token}是嵌入维度）以及先前的潜在特征z(n−1)\\mathbf{z}^{(n-1)}和先前的输出特征y(n−1)\\mathbf{y}^{(n-1)}来更新潜在特征z(n)\\mathbf{z}^{(n)}。更新的潜在特征z(n)\\mathbf{z}^{(n)}随后传递给RECURY\\operatorname{RECUR}_{Y}，它更新输出表示y(n)\\mathbf{y}^{(n)}，同时用零掩码输入嵌入。重复该耦合更新NN次，定义了方程2中的LATENT\\operatorname{LATENT}操作。这会逐步细化z(n)\\mathbf{z}^{(n)}和y(n)\\mathbf{y}^{(n)}。

LATENT≜(RECURY∘RECURZ)N\\displaystyle \\operatorname{LATENT}\\triangleq\\left(\\operatorname{RECUR}_{Y}\\circ\\operatorname{RECUR}_{Z}\\right)^{N}(2)z(n)=RECURZ([e1,...,eL,z(n−1),y(n−1)]),\\displaystyle \\mathbf{z}^{(n)}=\\operatorname{RECUR_{Z}}\\!\\left([\\mathbf{e}_{1},\\ldots,\\mathbf{e}_{L},\\mathbf{z}^{(n-1)},\\mathbf{y}^{(n-1)}]\\right),y(n)=RECURY([0,...,0,z(n),y(n−1)]).\\displaystyle \\mathbf{y}^{(n)}=\\operatorname{RECUR_{Y}}\\!\\left([\\mathbf{0},\\ldots,\\mathbf{0},\\mathbf{z}^{(n)},\\mathbf{y}^{(n-1)}]\\right).最后，为计算TRM（方程3），LATENT\\operatorname{LATENT}执行TT次，并馈入另一个RECUR\\operatorname{RECUR}。增加内部NN或外部TT循环会增加计算成本（FLOPs），而不会通过权重共享增加参数。

TRM(E,zini,yini)≜RECUR∘(LATENT)T=RECUR(E,[(RECURY(0,RECURZ(E,z,y),y))N]T).\\begin\{split\}&\\mathrm\{TRM\}\\!\\left(\\mathbf\{E\},\\mathbf\{z\}^\{ini\},\\mathbf\{y\}^\{ini\}\\right\)\\triangleq\\operatorname\{RECUR\}\\circ\\left\(\\operatorname\{LATENT\}\\right\)^\{T\}\\\\ &=\\operatorname\{RECUR\}\\!\\left\(\\mathbf\{E\},\\left\[\\left\(\\operatorname\{RECUR\}\_\{Y\}\\!\\left\(\\mathbf\{0\},\\operatorname\{RECUR\}\_\{Z\}\\!\\left\(\\mathbf\{E\},\\mathbf\{z\},\\mathbf\{y\}\\right\),\\mathbf\{y\}\\right\)\\right\)^\{N\}\\right\]^\{T\}\\right\)\.\\end\{split\}(3)

### II-B Depth Recursive Model

根据我们在先进封装领域的领域专业知识，我们观察到对于每个固定设计案例，数据集呈现出逐渐变化的深度式结构。受此观察启发，我们重新利用RECUR，将顺序深度视为状态输入，同时完全保留图2中的权重共享目标。我们称这种方法为深度递归模型（DEPTH）。受循环神经网络（RNNs）启发，DEPTH将隐状态的概念调整为RECUR内依赖权重的顺序深度状态。使用RNN的顺序建模术语，在DEPTH（方程4）中，我们将顺序状态定义为S(t)S^{(t)}，初始隐状态定义为嵌入输入E=H(0)\\mathbf{E}=\\mathbf{H}^{(0)}，对于深度t∈{1,...,T}t\\in\\{1,\\dots,T\\}，顺序（或递归）输出为H(t+1)H^{(t+1)}。

DEPTH(ESt⊈E,S(1),...,S(T))≜(RECUR∘(⋅+S(t)Ws))T=RECUR(E+S(1)Ws)→H(1)→RECUR(H(1)+S(2)Ws)→H(2)⋯→RECUR(H(T−1)+S(T)Ws)→H(T).\\begin\{array\}\[\}\]\{l\}\\mathrm\{DEPTH\}\\!\\Big\(\\mathbf\{E\}\_\{S^\{t\}\\nsubseteq E\},\\mathbf\{S\}^\{\(1\)\},\\ldots,\\mathbf\{S\}^\{\(T\)\}\\Big\)\\triangleq\\left\(\\operatorname\{RECUR\}\\circ\\left\(\\cdot\+\\mathbf\{S\}^\{\(t\)\}\\mathbf\{W\}\_\{s\}\\right\)\\right\)^\{T\}\\\\\[5\.69054pt\] =\\operatorname\{RECUR\}\\!\\left\(\\mathbf\{E\}+\\mathbf\{S\}^\{\(1\)\}\\mathbf\{W\}\_\{s\}\\right\)\\rightarrow\\mathbf\{H\}^\{\(1\)\}\\\\ \\rightarrow\\operatorname\{RECUR\}\\!\\left\(\\mathbf\{H\}^\{\(1\)\}+\\mathbf\{S\}^\{\(2\)\}\\mathbf\{W\}\_\{s\}\\right\)\\rightarrow\\mathbf\{H\}^\{\(2\)\}\\\\ \\cdots\\rightarrow\\operatorname\{RECUR\}\\!\\left\(\\mathbf\{H\}^\{\(T\-1\)\}+\\mathbf\{S\}^\{\(T\)\}\\mathbf\{W\}\_\{s\}\\right\)\\rightarrow\\mathbf\{H\}^\{\(T\)\}\.\\end\{array\}(4)

### II-C Simple Recursive model

图3中的简单递归模型（SIMPLE）被看作重复应用相同的递归变换。使用（方程5），SIMPLE对RECUR应用TT次递归迭代以生成输出H(T)H^{(T)}。

SIMPLE(E)\\displaystyle \\operatorname{SIMPLE}(\\mathbf{E})≜(RECUR)T\\displaystyle \\triangleq\\left(\\operatorname{RECUR}\\right)^{T}(5)=RECUR(E)→H(1)\\displaystyle =\\operatorname{RECUR}(\\mathbf{E})\\rightarrow\\mathbf{H}^{(1)}→RECUR(H(1))→H(2)\\displaystyle \\rightarrow\\operatorname{RECUR}\\left(\\mathbf{H}^{(1)}\\right)\\rightarrow\\mathbf{H}^{(2)}...→RECUR(H(T−1))→H(T)\\displaystyle \\dots\\rightarrow\\operatorname{RECUR}\\left(\\mathbf{H}^{(T-1)}\\right)\\rightarrow\\mathbf{H}^{(T)}

## III 实验结果

### III-A 基线和提出的模型

为了全面评估跟踪和检索性能，我们将我们的方法与跨四个不同架构范式的七个基线配置进行比较：

- •VANILLA：包含M1，作为基线1块层归一化（LN）Transformer架构。
- •SIMPLE：包含M2和M3，用于简单递归模型，表示一种2块均方根（RMS）递归结构，递归步数分别为T=1和T=3。
- •TRM：代表Tiny Recursive Model，其中M4表示包含潜在变量（ZZ）的1块RMS变体，M5表示在联合潜在配置（Z,YZ,Y）上操作的1块LN变体。使用了T=5和N=3。
- •DEPTH：包含深度递归模型，其中M6指定1块RMS设计，M7代表2块RMS架构。深度自适应为T=16，适用于Stress10K和Warpage10K两个数据集。

### III-B 数据集描述

我们在三个数据集上评估递归模型，其中包含部分DOE因子设计案例，N=因子^水平——Stress10k和Warpage10k数据集（N=5^4×16=10,000对）由FEA仿真生成。合成电容器静电场数据集由物理信息神经网络（PINN）生成（N=22^2×15=7,265对）。

Stress10k、Warpage10k[18](https://arxiv.org/html/2607.27251#bib.bib19)：先进封装设计需要理解材料属性和几何配置如何影响t