[论文解读] Branches of a Tree: Taking Derivatives of Programs with Discrete and Branching Randomness in High Energy Physics

M. Kagan, L. Heinrich|arXiv (Cornell University)|Aug 31, 2023

Distributed and Parallel Computing Systems被引用 5

一句话总结

本文提出了一种梯度估计技术——特别是带有基线的评分函数估计器与随机自动微分（Stochastic AD）——用于对高能物理（HEP）中具有离散性和分支随机性的程序进行微分。该研究首次实现了完全可微分的分支程序，并证明了即使在小批量尺寸下，随机AD与带基线的评分函数估计器也能实现有效的基于梯度的探测器设计优化。

ABSTRACT

We propose to apply several gradient estimation techniques to enable the differentiation of programs with discrete randomness in High Energy Physics. Such programs are common in High Energy Physics due to the presence of branching processes and clustering-based analysis. Thus differentiating such programs can open the way for gradient based optimization in the context of detector design optimization, simulator tuning, or data analysis and reconstruction optimization. We discuss several possible gradient estimation strategies, including the recent Stochastic AD method, and compare them in simplified detector design experiments. In doing so we develop, to the best of our knowledge, the first fully differentiable branching program.

研究动机与目标

• 在高能物理（HEP）中实现基于梯度的优化，适用于具有离散性和分支随机性的程序，例如粒子簇射和探测器模拟。
• 解决标准自动微分（AD）在处理HEP软件中参数依赖的离散随机性时的局限性。
• 开发并验证一个完全可微分的分支程序，用于粒子簇射模拟，这是HEP可微分编程领域的一项新贡献。
• 在简化探测器设计优化任务中，对比梯度估计技术——评分函数、数值梯度与随机AD。
• 证明端到端可微分HEP模拟器在探测器设计、模拟器调优和重建优化等应用中的可行性。

提出的方法

• 应用带有控制变量（基线）的评分函数梯度估计器，对具有离散随机性的随机程序进行梯度估计。
• 实现并应用最近开发的随机AD（Arya et al., 2022）方法，用于对具有离散随机变量的程序进行微分，应用于HEP工作负载。
• 设计一个简化探测器模拟器，其中包含一个分支簇射模型，探测器半径为可调参数，损失函数由径向簇射深度定义。
• 使用Adam优化器，结合小批量梯度估计（批量大小 = 2），对探测器半径参数执行500次优化步骤。
• 通过重复优化（每种方法运行10次）对比梯度估计器，并基于优化步骤中的平均损失与标准差评估性能。
• 在原始程序与替代程序中重用随机变量，以降低随机AD估计器的方差，观察到当启用重用时，标准差降低1.5倍或以上。

(a) Single particle energy loss

实验结果

研究问题

• RQ1 带有基线的评分函数梯度估计器是否能有效应用于具有离散随机性的HEP探测器设计优化？
• RQ2 随机AD是否能成功应用于HEP中的分支过程（如部分子簇射），以实现可微分编程？
• RQ3 在低样本量的HEP优化设置下，不同梯度估计器（评分函数、数值梯度、随机AD）在方差与收敛速度方面如何比较？
• RQ4 在随机AD中重用随机变量对HEP应用中梯度估计器方差有何影响？
• RQ5 在梯度估计器方差较高的情况下，是否可行仅通过每次梯度步骤运行两次模拟来实现基于梯度的优化？

主要发现

• 随机AD与带基线的评分函数估计器在所有测试的梯度估计器中表现出最低的标准差，使得仅通过每次步骤两次模拟即可实现稳定优化。
• 不带基线的评分函数估计器与数值梯度表现出过高的方差，导致在给定的低样本量设置下优化不可行。
• 随机AD与带基线的评分函数估计器在向损失最小值收敛方面表现相似，表明两者均适用于未来的HEP可微分编程应用。
• 在随机AD的替代程序中重用随机变量，可使梯度估计器方差降低1.5倍或以上，凸显了耦合策略的重要性。
• 本研究首次展示了已知的可微分分支程序，能够通过HEP中的离散分支过程估计梯度，从而实现粒子簇射的端到端可微分模拟。
• 研究结果表明，即使数据极少，随机AD与带基线的评分函数估计器在基于梯度的探测器设计优化中均表现有效，显示出在更广泛HEP应用中的强大潜力。

(b) Particle shower with splitting