---
title: 'SoftTorch: Soft Differentiable Primitives'
url: https://www.emergentmind.com/topics/softtorch
type: topic
---

# SoftTorch: Soft Differentiable Primitives

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 [2603.08824].

## 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 [2603.08824].

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 [2603.08824].

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 [2603.08824].

## 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 [2603.08824].

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` [2603.08824].

**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 [2603.08824].

**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 [2603.08824].

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:
$$
g_{\mathtt{c0}}(s)=\tfrac12+\tfrac{s}{2},\qquad
g_{\mathtt{c1}}(s)=\tfrac12+\tfrac{3s}{4}-\tfrac{s^3}{4},\qquad
g_{\mathtt{c2}}(s)=\tfrac12+\tfrac{15s}{16}-\tfrac{5s^3}{8}+\tfrac{3s^5}{16}.
$$
From this basis, the paper derives
$$
\operatorname{sign}_\tau(x)=2H_\tau(x)-1,\qquad
\operatorname{abs}_\tau(x)=\operatorname{sign}_\tau(x)\cdot x.
$$
It further defines a smooth rounding operator by summing shifted softened bins:
$$
\operatorname{round}_\tau(x) =
\sum_{k=\lfloor x \rfloor - K}^{\lfloor x \rfloor + K}
k\Bigl[H_\tau\!\left(x-k+\tfrac12\right)-H_\tau\!\left(x-k-\tfrac12\right)\Bigr],
$$
where $K$ controls how many neighboring soft bins contribute [2603.08824].

For ReLU, two constructions are given. One is an integration-style surrogate,
$$
\operatorname{relu}_\tau(x)=\int_0^x H_\tau(t)\,dt,
$$
which becomes Softplus under the logistic choice. The other is a gating-style surrogate,
$$
\operatorname{relu}_\tau(x)=x\cdot H_\tau(x),
$$
which becomes SiLU when $H_\tau$ is logistic. On that basis, `clip`/`clamp` is written as
$$
\operatorname{clip}_\tau(x,a,b)=a+\operatorname{relu}_\tau(x-a)-\operatorname{relu}_\tau(x-b).
$$
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 [2603.08824].

The library interprets soft comparison outputs as probabilities, called “SoftBools” in the paper’s exposition. These support fuzzy-logic-style composition:
$$
\operatorname{all}(p_1,\dots,p_n)=\prod_{j=1}^n p_j,\qquad
\operatorname{not}(p)=1-p,
$$
with
$$
\operatorname{any}(p_1,\ldots,p_n)=\neg \operatorname{all}(\neg p_1,\ldots,\neg p_n),
$$
$$
\operatorname{and}(p,q)=\operatorname{all}(p,q),\quad
\operatorname{or}(p,q)=\operatorname{any}(p,q),\quad
\operatorname{xor}(p,q)=\operatorname{or}(\operatorname{and}(p,\neg q),\operatorname{and}(\neg p,q)).
$$
These SoftBools can be used in soft selection through expectation,
$$
z_i=p_i x_i + (1-p_i)y_i,
$$
which the paper describes as being implemented through soft versions of `where` [2603.08824].

## 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 [2603.08824].

For soft argmax, the hard one-hot index is replaced by a probability vector $\mathbf{p}\in\Delta_n$, where
$$
\Delta_n = \{\mathbf{p}\in[0,1]^n \mid \sum_j p_j = 1\}.
$$
This SoftIndex can then be used to compute expectations,
$$
\mathbf{p}\cdot \mathbf{x} = \sum_j p_j x_j = \mathbb{E}_{j\sim \mathbf{p}}[x_j].
$$
The same principle underlies soft replacements of `argmax`, `max`, and index selection [2603.08824].

For sorting and ranking, the hard sorted output is expressed as multiplication by a permutation matrix $P^\star$,
$$
\operatorname{sort}(\mathbf{x}) = P^\star \mathbf{x},
$$
and the soft version relaxes this matrix to a bistochastic or row-stochastic matrix $P_\tau$. 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** [2603.08824]

For OT-based operators, SoftTorch solves the regularized transport problem
$$
\Gamma_\tau^\star \coloneqq \arg\min_{\Gamma\in U(\mathbf{a},\mathbf{b})} \langle \Gamma, C \rangle + \tau R(\Gamma),
$$
over
$$
U(\mathbf{a},\mathbf{b}) = \{\Gamma\in\mathbb{R}_+^{n\times m}\mid \Gamma\mathbf{1}_m=\mathbf{a},\ \Gamma^\top\mathbf{1}_n=\mathbf{b}\},
$$
with squared-distance cost $C_{ij}=(x_i-y_j)^2$. The library supports entropic, Euclidean, and $p$-norm regularization. From the transport-derived soft permutation, it defines
$$
\argsort_\tau(\mathbf{x}) = P_\tau^\star,\qquad
\sort_\tau(\mathbf{x}) = P_\tau^\star \mathbf{x},\qquad
\rank_\tau(\mathbf{x}) = (P_\tau^\star)^\top [n,\ldots,1]^\top.
$$
Analogous constructions are stated for top-k and quantiles by changing the anchor structure and reading off relevant rows of the transport plan [2603.08824].

The simplex-projection family is introduced through
$$
\Pi_\tau(\mathbf{x}) = \arg\max_{\mathbf{p}\in\Delta_n} \langle \mathbf{x}, \mathbf{p}\rangle - \tau R(\mathbf{p}).
$$
With entropic regularization, this becomes softmax; with Euclidean regularization, it yields simplex projection or SparseMax-like behavior; with $p$-norm regularization, it yields Bregman projections with controllable smoothness. NeuralSort is formulated using the matrix of soft absolute differences
$$
(A_\mathbf{x}^\tau)_{ij} = \operatorname{abs}_\tau(x_i-x_j),
$$
and
$$
\argsort_\tau(\mathbf{x})_i = \Pi_\tau\Bigl((2i-n-1)\mathbf{x} - A_\mathbf{x}^\tau \mathbf{1}_n\Bigr).
$$
The paper explicitly states that this framework covers not only sorting but also soft argmax, argmedian, and soft quantiles [2603.08824].

Permutahedron-based methods are motivated by computational efficiency for large vectors. The permutahedron is defined as
$$
\mathcal{P}(\mathbf{z}) = \mathrm{conv}\{\mathbf{z}_\sigma\mid \sigma\in\Sigma\},
$$
and FastSoftSort-style operators are given as
$$
\sort_\tau(\mathbf{x}) = \operatorname{Proj}_\tau([1,\ldots,n],\mathbf{x}),\qquad
\rank_\tau(\mathbf{x}) = \operatorname{Proj}_\tau(-\mathbf{x},[1,\ldots,n]).
$$
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,
$$
\widetilde{b}_k = \tau \log \sum_{|S|=k} \exp\Bigl(\frac{1}{\tau}\sum_{i\in S} z_i\Bigr) = \tau \log e_k(e^{\mathbf{z}/\tau}).
$$
The paper describes SmoothSort as $\mathcal{C}^\infty$, while the other permutahedron methods are at most $\mathcal{C}^2$ [2603.08824].

Sorting-network-based methods soften compare-and-swap operations through
$$
\sigma = H_\tau(a-b),\qquad
\operatorname{soft\_min}(a,b)=\sigma b+(1-\sigma)a,\qquad
\operatorname{soft\_max}(a,b)=\sigma a+(1-\sigma)b.
$$
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` [2603.08824].

## 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
$$
f_{\mathrm{STE}}(x) = \operatorname{sg}(f(x)) + f_\tau(x) - \operatorname{sg}(f_\tau(x)),
$$
where $\operatorname{sg}$ is the stop-gradient operator. This yields hard forward behavior and soft backward gradients:
$$
f_{\mathrm{STE}} = f,\qquad \nabla f_{\mathrm{STE}} = \nabla f_\tau.
$$
SoftTorch exposes this through a decorator-style wrapper, exemplified in the paper as `sj.st(sj.relu)` in the shared library family [2603.08824].

The paper also identifies an “STE pitfall.” If individually STE-wrapped functions are multiplied together, hard forward values can still gate gradients:
$$
\nabla(f_{\mathrm{STE}}\cdot g_{\mathrm{STE}})=\nabla f_\tau \cdot g + f \cdot \nabla g_\tau.
$$
The recommendation is to apply STE to the composite function instead,
$$
\nabla(f\cdot g)_{\mathrm{STE}}=\nabla(f_\tau \cdot g_\tau)=\nabla f_\tau \cdot g_\tau + f_\tau \cdot \nabla g_\tau.
$$
This is presented as a practical subtlety that can materially affect whether softened gradients remain informative in composed programs [2603.08824].

SoftTorch supports multiple smoothness modes: `smooth`, `c0`, `c1`, and `c2`. The theoretical basis given in the paper is a smoothness analysis for $p$-norm regularized projections. It states that with $p=2$ the induced surrogate is $\mathcal{C}^0$, with $p=3/2$ it is $\mathcal{C}^1$, with $p=4/3$ it is $\mathcal{C}^2$, and entropic regularization yields $\mathcal{C}^\infty$. The analysis is said to apply to both simplex projection and OT projection under standard nondegeneracy conditions for OT [2603.08824].

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 $(0,1)$ 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 [2603.08824].

## 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 $n=4096$, 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 $n=4096$ 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 [2603.08824].

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 [2603.08824].

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 [2508.06628]. SoftTorch, by contrast, is specifically a library of soft relaxations for non-differentiable PyTorch primitives [2603.08824].

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 [2603.08824].

Source: https://www.emergentmind.com/topics/softtorch