让深度学习框架支持更智能的梯度计算,解决离散操作无法优化的问题。
SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients
- 用模糊逻辑和最优传输方法实现软化版的阈值、排序等操作
- 提供可直接替换原生库的软函数,支持零梯度场景下的有效优化
- 适合需要处理离散变量的科研与工程应用,如神经网络结构搜索
自动微分框架如 JAX 和 PyTorch 在众多科学领域中推动了基于梯度的优化。然而,这些库中的许多“硬”操作(如阈值、布尔逻辑、离散索引、排序)会产生零或未定义梯度,无法用于优化。尽管已有多种“软”松弛方法能提供有用梯度,但实现分散在不同项目中,难以整合与比较。本文提出 SoftJAX 与 SoftTorch,两个开源、功能完整的软可微编程库,提供一系列软函数作为 JAX 与 PyTorch 对应硬函数的即插即用替代品。包括:(i) 元素级操作如 clip、abs;(ii) 通过模糊逻辑操纵布尔值与索引的工具方法;(iii) 基于最优传输或排列多面体投影的轴向操作如 sort、rank;(iv) 完整支持直通梯度估计。实验与案例研究证明,该工具箱显著提升了软松弛方法的可用性。代码开源地址:github.com/a-paulus/softjax 与 github.com/a-paulus/softtorch。
原文摘要 · Abstract (English)
Automatic differentiation (AD) frameworks such as JAX and PyTorch have enabled gradient-based optimization for a wide range of scientific fields. Yet, many "hard" primitives in these libraries such as thresholding, Boolean logic, discrete indexing, and sorting operations yield zero or undefined gradients that are not useful for optimization. While numerous "soft" relaxations have been proposed that provide informative gradients, the respective implementations are fragmented across projects, making them difficult to combine and compare. This work introduces SoftJAX and SoftTorch, open-source, feature-complete libraries for soft differentiable programming. These libraries provide a variety of soft functions as drop-in replacements for their hard JAX and PyTorch counterparts. This includes (i) elementwise operators such as clip or abs, (ii) utility methods for manipulating Booleans and indices via fuzzy logic, (iii) axiswise operators such as sort or rank -- based on optimal transport or permutahedron projections, and (iv) offer full support for straight-through gradient estimation. Overall, SoftJAX and SoftTorch make the toolbox of soft relaxations easily accessible to differentiable programming, as demonstrated through benchmarking and a practical case study. Code is available at github.com/a-paulus/softjax and github.com/a-paulus/softtorch.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。