SoftTorch: Soft Differentiable Primitives

• SoftTorch is an open-source PyTorch library that provides soft surrogates for non-differentiable operations to enable gradient-based learning.
• It consolidates diverse relaxations including fuzzy logic, axiswise softening, and straight-through gradient estimation under a unified PyTorch-style API.
• The library mitigates issues of zero or undefined gradients in operators like argmax, sort, and thresholding, improving training dynamics.

Searching arXiv for the SoftTorch paper and closely related context. SoftTorch is an open-source, feature-complete PyTorch library for soft differentiable programming that provides drop-in soft replacements for hard, non-differentiable PyTorch primitives whose native gradients are zero, undefined, or otherwise uninformative for optimization. Introduced alongside SoftJAX in “SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients,” it targets the recurrent mismatch between discrete operators and gradient-based methods by consolidating soft relaxations, fuzzy-logic utilities, axiswise relaxations, and straight-through gradient estimation into a single library with PyTorch-style naming and APIs (Paulus et al., 9 Mar 2026).

1. Origin, scope, and motivating problem

SoftTorch is the PyTorch-facing counterpart of a broader library family for soft differentiable programming. Its stated purpose is to replace hard tensor-library operations such as thresholding, Boolean logic, discrete indexing, sorting, ranking, top-k selection, quantiles, and medians with soft surrogates that preserve useful gradients while remaining usable from ordinary PyTorch code (Paulus et al., 9 Mar 2026).

The central problem addressed by the library is that many hard PyTorch operators—including abs, sign, round, comparison operators, argmax, sort, rank, and topk—either have zero gradients almost everywhere, undefined gradients at discontinuities, or gradients that are not useful for learning. SoftTorch addresses this by supplying a standardized toolbox of operators that approximate the hard primitive in the limit of a softness parameter $\tau \to 0^+$, provide informative gradients during training, can be used as direct replacements for corresponding PyTorch functions, and support a straight-through mode for hard-forward/soft-backward execution (Paulus et al., 9 Mar 2026).

A defining design claim is consolidation. The paper frames existing soft relaxations as fragmented across papers and codebases, making them difficult to combine and compare. SoftTorch is positioned as a unifying abstraction layer that collects multiple relaxation families under consistent interfaces and exposes multiple softness and smoothness modes. This suggests that the library is not tied to a single surrogate mechanism, but instead organizes a family of alternatives under a PyTorch-native API surface (Paulus et al., 9 Mar 2026).

2. Operator families and PyTorch-facing interface

SoftTorch follows PyTorch naming conventions and adapts the operator set shared with SoftJAX to Torch-style APIs. The paper gives explicit name correspondences: clip becomes clamp, equal becomes eq, top_k becomes topk, take_along_axis becomes take_along_dim, dynamic_index_in_dim becomes index_select, and dynamic_slice_in_dim becomes narrow. It also notes that some JAX-specific utilities such as choose and dynamic_slice have no PyTorch equivalent and therefore are not offered in SoftTorch (Paulus et al., 9 Mar 2026).

The library organizes its soft replacements into several operator categories.

Elementwise operators include heaviside, sign, abs, round, relu, clip/clamp, and comparisons such as less, less_equal, equal, not_equal, greater, gtr_equal, and isclose (Paulus et al., 9 Mar 2026).

Boolean and index utilities extend softened comparisons into fuzzy-logic-style selection and control. Rather than treating conditionals and masks as hard discrete objects, SoftTorch uses probabilities or soft index distributions so that gradients can propagate through both branches of a selection rule (Paulus et al., 9 Mar 2026).

Axiswise operators constitute the most extensive class. The paper states that SoftTorch covers argmax, argsort, argquantile, argmedian, argtop_k, rank, sort, top_k, quantile, and median. These are treated through SoftIndices on the simplex, soft permutation matrices, or sorting-network-based differentiable compare-and-swap constructions (Paulus et al., 9 Mar 2026).

This operator coverage is significant because it spans both simple local nonlinearities and global, order-sensitive transformations. A plausible implication is that SoftTorch is intended not merely as a collection of isolated surrogates, but as a replacement layer for entire non-smooth subroutines in differentiable programs.

3. Elementwise softening, SoftBools, and selection semantics

The elementwise constructions are built from a softened Heaviside step function. The hard Heaviside function is defined in the paper as

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$

In smooth mode, SoftTorch uses the logistic sigmoid

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$

It also supports piecewise polynomial modes

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$

with explicit choices $g_{\mathtt{c0}}$, $g_{\mathtt{c1}}$, and $g_{\mathtt{c2}}$ giving $\mathcal{C}^0$, $\mathcal{C}^1$, and $\mathcal{C}^2$ smoothness, respectively:

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$0

From this basis, the paper derives

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$1

It further defines a smooth rounding operator by summing shifted softened bins:

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$2

where $$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$3 controls how many neighboring soft bins contribute (Paulus et al., 9 Mar 2026).

For ReLU, two constructions are given. One is an integration-style surrogate,

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$4

which becomes Softplus under the logistic choice. The other is a gating-style surrogate,

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$5

which becomes SiLU when $$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$6 is logistic. On that basis, clip/clamp is written as

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$7

These formulas clarify that SoftTorch systematically derives higher-level non-smooth primitives from a small number of softened building blocks rather than defining each operator independently (Paulus et al., 9 Mar 2026).

The library interprets soft comparison outputs as probabilities, called “SoftBools” in the paper’s exposition. These support fuzzy-logic-style composition:

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$8

with

$$H(x) \coloneqq \begin{cases} 0 & \text{if } x < 0,\ 0.5 & \text{if } x = 0,\ 1 & \text{else}. \end{cases}$$9

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$0

These SoftBools can be used in soft selection through expectation,

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$1

which the paper describes as being implemented through soft versions of where (Paulus et al., 9 Mar 2026).

4. Axiswise relaxations: SoftIndices, soft permutations, and algorithmic families

Axiswise operators are treated as substantially harder than elementwise ones because they return indices, ranks, or permutations rather than scalarwise transformed values. SoftTorch therefore replaces hard index outputs with probabilistic or matrix-valued relaxations (Paulus et al., 9 Mar 2026).

For soft argmax, the hard one-hot index is replaced by a probability vector $$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$2, where

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$3

This SoftIndex can then be used to compute expectations,

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$4

The same principle underlies soft replacements of argmax, max, and index selection (Paulus et al., 9 Mar 2026).

For sorting and ranking, the hard sorted output is expressed as multiplication by a permutation matrix $$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$5,

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$6

and the soft version relaxes this matrix to a bistochastic or row-stochastic matrix $$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$7. The paper distinguishes four axiswise method families:

1. Optimal transport (OT)-based methods
2. Unit-simplex projection / SoftSort / NeuralSort
3. Permutahedron projection-based methods, including FastSoftSort and SmoothSort
4. Sorting-network-based methods (Paulus et al., 9 Mar 2026)

For OT-based operators, SoftTorch solves the regularized transport problem

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$8

over

$$H_\tau(x) \coloneqq \sigma\!\left(\frac{x}{\tau}\right) = \frac{1}{1+\exp(-x/\tau)}.$$9

with squared-distance cost $$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$0. The library supports entropic, Euclidean, and $$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$1-norm regularization. From the transport-derived soft permutation, it defines

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$2

Analogous constructions are stated for top-k and quantiles by changing the anchor structure and reading off relevant rows of the transport plan (Paulus et al., 9 Mar 2026).

The simplex-projection family is introduced through

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$3

With entropic regularization, this becomes softmax; with Euclidean regularization, it yields simplex projection or SparseMax-like behavior; with $$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$4-norm regularization, it yields Bregman projections with controllable smoothness. NeuralSort is formulated using the matrix of soft absolute differences

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$5

and

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$6

The paper explicitly states that this framework covers not only sorting but also soft argmax, argmedian, and soft quantiles (Paulus et al., 9 Mar 2026).

Permutahedron-based methods are motivated by computational efficiency for large vectors. The permutahedron is defined as

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$7

and FastSoftSort-style operators are given as

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$8

SoftTorch further introduces a SmoothSort variant that adds entropic regularization to a dual permutahedron projection and smoothes the order-statistic bounds with a log-sum-exp relaxation,

$$H_\tau(x) \coloneqq \begin{cases} 0 & \text{if } x < -\tau,\ g(x/\tau) & \text{if } -\tau \le x \le \tau,\ 1 & \text{else}, \end{cases}$$9

The paper describes SmoothSort as $g_{\mathtt{c0}}$0, while the other permutahedron methods are at most $g_{\mathtt{c0}}$1 (Paulus et al., 9 Mar 2026).

Sorting-network-based methods soften compare-and-swap operations through

$g_{\mathtt{c0}}$2

These constructions can produce both sorted values and soft permutation matrices, enabling argsort and argmax in addition to sort, max, min, quantile, median, and top_k (Paulus et al., 9 Mar 2026).

5. Straight-through estimation, smoothness modes, and implementation choices

A major feature of SoftTorch is explicit support for straight-through gradient estimation (STE). The library is designed for cases in which the forward computation should remain hard while the backward pass should use a soft surrogate. The paper defines STE as

$g_{\mathtt{c0}}$3

where $g_{\mathtt{c0}}$4 is the stop-gradient operator. This yields hard forward behavior and soft backward gradients:

$g_{\mathtt{c0}}$5

SoftTorch exposes this through a decorator-style wrapper, exemplified in the paper as sj.st(sj.relu) in the shared library family (Paulus et al., 9 Mar 2026).

The paper also identifies an “STE pitfall.” If individually STE-wrapped functions are multiplied together, hard forward values can still gate gradients:

$g_{\mathtt{c0}}$6

The recommendation is to apply STE to the composite function instead,

$g_{\mathtt{c0}}$7

This is presented as a practical subtlety that can materially affect whether softened gradients remain informative in composed programs (Paulus et al., 9 Mar 2026).

SoftTorch supports multiple smoothness modes: smooth, c0, c1, and c2. The theoretical basis given in the paper is a smoothness analysis for $g_{\mathtt{c0}}$8-norm regularized projections. It states that with $g_{\mathtt{c0}}$9 the induced surrogate is $g_{\mathtt{c1}}$0, with $g_{\mathtt{c1}}$1 it is $g_{\mathtt{c1}}$2, with $g_{\mathtt{c1}}$3 it is $g_{\mathtt{c1}}$4, and entropic regularization yields $g_{\mathtt{c1}}$5. The analysis is said to apply to both simplex projection and OT projection under standard nondegeneracy conditions for OT (Paulus et al., 9 Mar 2026).

The implementation details reinforce the library’s role as an integration layer rather than a monolithic new AD framework. For OT-based methods, SoftTorch uses POT. For permutahedron projection methods, it uses Numba-JIT acceleration for the PAV-based projection routines. The paper also notes a standardization-and-squashing preprocessing step for axiswise operators: inputs are standardized and mapped into $g_{\mathtt{c1}}$6 with a sigmoid to improve numerical stability and make the softness parameter more scale-independent, and an inverse transform is applied afterward for value-returning outputs like sorted values. SoftSort is listed as an exception because its output can exceed the input range, so it skips this preprocessing (Paulus et al., 9 Mar 2026).

6. Benchmarking, case study, and nomenclature

The benchmark discussion emphasizes tradeoffs among speed, memory, smoothness, and exactness rather than a single dominant method. For sort in smooth mode on an NVIDIA RTX 3060, the paper reports that the sorting network is the fastest soft method, about 1.0 ms at $g_{\mathtt{c1}}$7, and roughly 3.8× slower than the hard baseline. It reports SoftSort at about 16 ms, NeuralSort at about 37 ms, and FastSoftSort as the most memory efficient at roughly 420 KB at $g_{\mathtt{c1}}$8 because it does not materialize the full permutation matrix. SmoothSort and OT-based methods are described as the slowest among the tested methods. The paper notes that the benchmarks are mostly reported for SoftJAX, but treats them as relevant to SoftTorch because the two libraries share the same operator design and algorithmic families (Paulus et al., 9 Mar 2026).

The practical case study appears in the appendix and concerns collision detection in MuJoCo XLA (MJX). A mesh-mesh collision subroutine selecting four polygon vertices through nested argmax and indexing operations is softened by replacing hard argmax with sj.argmax, hard indexing with sj.dynamic_index_in_dim, and hard abs with sj.abs. The softened version returns soft index distributions rather than hard indices, which requires corresponding downstream adjustments. The reported effect is that the softened algorithm produces smooth, non-zero gradients for all vertices, unlike the hard version where some vertices receive zero gradient. The paper adds that when small softness values are needed, gradient clipping may be helpful (Paulus et al., 9 Mar 2026).

A recurring source of confusion is nomenclature. A separate 2025 paper, “TorchSim: An efficient atomistic simulation engine in PyTorch,” introduces TorchSim, not SoftTorch; its nearest lexical connection is only the support for a soft-sphere classical potential, and it is otherwise unrelated in purpose and design (Cohen et al., 8 Aug 2025). SoftTorch, by contrast, is specifically a library of soft relaxations for non-differentiable PyTorch primitives (Paulus et al., 9 Mar 2026).

Taken together, the paper presents SoftTorch as a PyTorch-native replacement layer for discrete and non-smooth primitives. Its distinguishing features are the combination of elementwise soft surrogates, SoftBool-based fuzzy logic, axiswise relaxations spanning OT, simplex, permutahedron, and sorting-network constructions, and explicit support for STE. This suggests a conception of differentiable programming in which hard control, selection, and ordering operations are not removed, but systematically re-expressed so that optimization remains gradient-informative (Paulus et al., 9 Mar 2026).