SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

• URL: http://arxiv.org/abs/2603.08824v1
• Date: Mon, 09 Mar 2026 18:35:51 GMT
• Title: SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients
• Abstract summary: This work introduces SoftJAX and SoftTorch, open-source, feature-complete libraries for soft differentiable programming.<n>These libraries provide a variety of soft functions as drop-in replacements for their hard JAX and PyTorch counterparts.<n>Overall, SoftJAX and SoftTorch make the toolbox of soft relaxations easily accessible to differentiable programming.
• Score: 27.151039291881677
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Beyond Softmax: A Natural Parameterization for Categorical Random Variables [61.709831225296305]
  We introduce the $textitcatnat$ function, a function composed of a sequence of hierarchical binary splits.<n>A rich set of experiments show that the proposed function improves the learning efficiency and yields models characterized by consistently higher test performance.
  arXiv  Detail & Related papers  (2025-09-29T12:55:50Z)
• A projection-based framework for gradient-free and parallel learning [50.96641619247761]
  We introduce PJAX, a JAX-based software framework that enables this paradigm.<n>PJAX composes projection operators for elementary operations, automatically deriving the solution operators for the feasibility problems.<n>We train diverse architectures (MLPs, CNNs, RNNs) on standard benchmarks using PJAX, demonstrating its generality.
  arXiv  Detail & Related papers  (2025-06-06T08:44:56Z)
• SoftCoT++: Test-Time Scaling with Soft Chain-of-Thought Reasoning [48.28847964704554]
  Test-Time Scaling (TTS) refers to approaches that improve reasoning performance by allocating extra computation during inference.<n>Recent studies in Coconut and SoftCoT have demonstrated that thinking in the continuous latent space can further enhance the reasoning performance.<n>We introduce SoftCoT++ to extend SoftCoT to the Test-Time Scaling paradigm by enabling diverse exploration of thinking paths.
  arXiv  Detail & Related papers  (2025-05-16T17:47:50Z)
• Geometric-Averaged Preference Optimization for Soft Preference Labels [78.2746007085333]
  Many algorithms for aligning LLMs with human preferences assume that human preferences are binary and deterministic.<n>In this work, we introduce the distributional soft preference labels and improve Direct Preference Optimization (DPO) with a weighted geometric average of the LLM output likelihood in the loss function.
  arXiv  Detail & Related papers  (2024-09-10T17:54:28Z)
• Flexible and Efficient Surrogate Gradient Modeling with Forward Gradient Injection [2.0577627277681887]
  PyTorch, for example, allows the custom specification of the backward pass of an operation by overriding its backward method. While these methods are common practice and usually work well, they also have several disadvantages such as limited flexibility, additional source code overhead, poor usability, or a potentially strong negative impact on the effectiveness of automatic model optimization procedures. In this paper, an alternative way to formulate surrogate gradients is presented, namely, forward gradient injection (FGI)
  arXiv  Detail & Related papers  (2024-05-31T20:14:35Z)
• QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Acceleration [3.8494302715990845]
  QDax is an open-source library with a streamlined and modular API for Quality-Diversity (QD) optimization algorithms in Jax. The library serves as a versatile tool for optimization purposes, ranging from black-box optimization to continuous control.
  arXiv  Detail & Related papers  (2023-08-07T15:29:44Z)
• SynJax: Structured Probability Distributions for JAX [3.4447129363520337]
  SynJax provides efficient vectorized implementation of inference algorithms for structured distributions. We can build large-scale differentiable models that explicitly model structure in the data.
  arXiv  Detail & Related papers  (2023-08-07T04:20:38Z)
• Learning to Rasterize Differentiably [10.790255179159459]
  Differentiableization changes the standard formulation of primitiveization, by enabling flow from a pixel to its underlying triangles. Previous work has analyzed and compared several combinations of softening. In this work, we take it a step further and, instead of making a choice of softening operations, parameterize the continuous space of common softening operations.
  arXiv  Detail & Related papers  (2022-11-23T22:28:29Z)
• Theseus: A Library for Differentiable Nonlinear Optimization [21.993680737841476]
  Theseus is an efficient application-agnostic library for differentiable nonlinear least squares (DNLS) optimization built on PyTorch. Theseus provides a common framework for end-to-end structured learning in robotics and vision.
  arXiv  Detail & Related papers  (2022-07-19T17:57:40Z)
• Gradient Estimation with Stochastic Softmax Tricks [84.68686389163153]
  We introduce softmax tricks, which generalize the Gumbel-Softmax trick to spaces. We find that softmax tricks can be used to train latent variable models that perform better and discover more latent structure.
  arXiv  Detail & Related papers  (2020-06-15T00:43:44Z)
• Torch-Struct: Deep Structured Prediction Library [138.5262350501951]
  We introduce Torch-Struct, a library for structured prediction. Torch-Struct includes a broad collection of probabilistic structures accessed through a simple and flexible distribution-based API.
  arXiv  Detail & Related papers  (2020-02-03T16:43:02Z)
• Variance Reduction with Sparse Gradients [82.41780420431205]
  Variance reduction methods such as SVRG and SpiderBoost use a mixture of large and small batch gradients. We introduce a new sparsity operator: The random-top-k operator. Our algorithm consistently outperforms SpiderBoost on various tasks including image classification, natural language processing, and sparse matrix factorization.
  arXiv  Detail & Related papers  (2020-01-27T08:23:58Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.