SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

TL;DR

SoftJAX and SoftTorch are open-source libraries that provide soft, differentiable replacements for hard primitives in JAX and PyTorch, enabling more effective gradient-based optimization in complex operations.

Contribution

The paper introduces SoftJAX and SoftTorch, comprehensive libraries that unify and simplify the use of soft relaxations for hard primitives in differentiable programming.

Findings

Enhanced gradient information for discrete operations

Improved optimization performance demonstrated in benchmarks

Accessible implementation of various soft relaxations

Abstract

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.