在 7 nm、5 nm 乃至 3 nm 的晶体管森林里，声子的平均自由程动辄与器件尺寸同阶；热量不再像水一样“渗”过去，而是沿着特定的方向直直的“冲”过去——傅里叶定律说好的局部热平衡就此崩塌。想要看清声子们走哪条“高速路”，声子玻尔兹曼输运方程（BTE） 是目前最靠谱的“高德地图”。

Fig. 1 | Results of 1D cross-plane phonon transport.

可惜，这张地图太大：空间 × 频率 × 方向，多维卷在一起就是亿级未知数；传统 CPU上求解，单次迭代也得排队好几个小时。更别提我们还想问：如果薄膜再薄一点温度能降几度？如果热源换位置能不能让芯片凉快一点？——这些问题都需要做“反向推理”，而传统的BTE数值求解器并不能提供梯度信息，难住了优化算法。

Fig. 2 | Results of 3D FinFET simulations with diffusive reflecting boundary conditions of all faces except the bottom surface, which is set at 300 K as an isothermal boundary condition.

美国圣母大学尚文杰、周嘉航等同学不服，他们在导师罗腾飞、康奈尔王建勋教授的指导下下，开发出了JAX‑BTE，一口气解决了以下三个痛点。

Python接口 + JIT 编译：用户只需十几行 Python 即可设定材料、网格与边界条件；JAX 的 just-in-time 编译把脚本自动转成高度优化的 GPU 内核，解决了 Python “慢”的顾虑。

FVM → GNN 的算法重构：团队首次把有限体积法 (FVM) 的计算流抽象为图神经网络 (Graph Neural Network)的信息传递（message passing）过程：把每个网格单元视作图节点，守恒方程算出来的结果则化为消息；再利用 JAX 的向量化广播在 GPU 上同时推送到所有节点，实现亿级自由度的并行更新。

可微分+ 端到端反演：JAX 原生自动微分让温度场对“薄膜厚度”“热源强度”等设计参数可直接求梯度，通过手中的数据，可以直接反向求解找到适用于当前数据的最优参数，首度实现高效反演。

Fig. 3 | Temperaturedistribution for the linear transistor array simulation.

大规模基准验证显示：在1D-3D的各种复杂情况下，JAX‑BTE 计算的温度场与解析解或 GiftBTE 等 C++求解器保持 <1% 误差；在单张 NVIDIAA100 GPU 上，正向模拟速度比 64 核AMDEPYC CPU 提升 10倍。依托端到端梯度，薄膜厚度、热源强度等反演任务往往几十步梯度下降即可收敛，比传统传统的无梯度优化算法带来巨大的性能优势。JAX-BTE 既保持 GiftBTE 等 C++ 求解器的精度，又引入 AI 友好的梯度，极大推动半导体器件多尺度热设计、材料反演与混合物理-数据建模。

因此，AX‑BTE 让 BTE 不仅“跑得快”，还“会思考”——正向仿真 + 反向设计，一套代码全搞定。该文近期发表于npj Computational Materials11:129 (2025)。

Fig. 4 | Results of inverse learning for the heat source intensity of the 2D square domain.

论文链接：

https://www.nature.com/articles/s41524-025-01635-0