Soft-Argmax: Differentiable Approximation to Argmax

• Soft-argmax is a differentiable approximation to the argmax operator that converts logits and heatmaps into continuous, trainable outputs.
• It enables gradient propagation through discrete decision stages, facilitating end-to-end optimization in localization, object detection, and control tasks.
• By adjusting a temperature parameter, soft-argmax balances between diffuse and nearly one-hot outputs, offering a practical trade-off between precision and uncertainty.

Soft-argmax is a differentiable approximation to the non-differentiable $\arg\max$ operator. In current usage, the term denotes two closely related constructions: a softmax with a temperature or scaling parameter applied to logits so that, as the parameter increases, the output approaches a one-hot vector; and the expectation, or “center of mass,” of a softmax-normalized spatial map, used to extract continuous coordinates from heatmaps or attention maps (Luo et al., 2024, Sauvalle et al., 2022). Its primary role is to replace hard discrete selection with smooth operators through which gradients can propagate, enabling end-to-end optimization in localization, object-centric representation learning, mixed-integer control, and related settings (Li et al., 2021).

1. Core mathematical forms

In the discrete-decision setting, soft-argmax is implemented as a softmax function with a temperature or scaling parameter $K$:

$$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$

Here $z=[z_1,\ldots,z_{n_b}]$ are logits, $i$ indexes the discrete options, and $K$ controls sharpness. As $K \rightarrow +\infty$, the output approaches a one-hot vector, making the operator a closer approximation to the non-differentiable $\arg\max$ (Luo et al., 2024).

In spatial localization, soft-argmax is the expected position under a softmax-normalized map:

$$\operatorname{softargmax}(\phi)=\sum_{p\in\mathbb{Z}^2}\operatorname{softmax}(\phi)(p)\,p,$$

with

$$\operatorname{softmax}(\phi)(p)=\frac{e^{\phi(p)}}{\sum_{q\in\mathbb{Z}^2}e^{\phi(q)}}.$$

This formulation turns a spatial score map into normalized coordinates by computing a weighted average over positions. In the multi-object segmentation setting, normalized coordinates are obtained from attention maps by

$K$0

where $K$1 is a spatial softmax of attention logits (Sauvalle et al., 2022).

These two forms are mathematically aligned. Both replace a hard maximizer with a differentiable weighted sum, and both interpolate between diffuse and nearly discrete behavior through the concentration of the underlying softmax distribution. A plausible implication is that the term “soft-argmax” names a family of relaxations rather than a single canonical operator.

2. Differentiability and optimization role

The principal motivation for soft-argmax is that it is fully differentiable with respect to both network parameters and inputs. In neural network optimizers, this permits standard gradient-based training with backpropagation:

$K$2

Because the soft-argmax layer is smooth, the chain rule propagates gradients through all layers, including the relaxed discrete-decision stage (Luo et al., 2024).

In spatial models, differentiability is not merely an implementation convenience; it determines how localization is learned. In unsupervised multi-object segmentation, attention weights are computed via spatial softmax, and coordinates are weighted sums of these attention maps. This allows the model to learn object locations end-to-end without direct supervision of coordinates. The same attention map can also be used to aggregate feature maps into per-object feature vectors, coupling “where” and “what” in a shared differentiable mechanism (Sauvalle et al., 2022).

The literature also gives a theoretical rationale for this construction. Under affine, translation-equivariant localization, the expected value or center of mass is presented as the natural operator for converting a normalized spatial map into coordinates. This situates soft-argmax not only as a heuristic relaxation, but as an operator compatible with translation equivariance in attention-based localization (Sauvalle et al., 2022).

3. Spatial localization and representation learning

Soft-argmax is commonly adopted in detection-based methods to localize targets from probability maps or heatmaps. A standard formulation is

$K$3

or, in heatmap regression for landmark detection,

$K$4

This pattern appears in 2D/3D human pose estimation, facial landmark localization, stereo matching, and object keypoint estimation (Li et al., 2021, Yang et al., 19 Aug 2025).

In object-centric representation learning, soft-argmax occupies a more structural role. An architecture for unsupervised multi-object detection and segmentation uses a translation-equivariant attention mechanism to predict coordinates of objects present in a scene, a transformer encoder to handle occlusions and redundant detections, and a convolutional autoencoder for background reconstruction. Each attention map is normalized by spatial softmax, then soft-argmax yields object coordinates and attention-weighted feature aggregation yields object-centric embeddings (Sauvalle et al., 2022).

This usage emphasizes several properties. First, hard argmax is avoided, so localization remains trainable end-to-end. Second, soft-argmax can express uncertainty through spread-out attention maps. Third, in multi-object settings each slot or query can have its own attention map and own soft-argmax localization, supporting parallel prediction of multiple objects (Sauvalle et al., 2022).

4. Mixed-integer control and relaxed discrete decisions

A different use of soft-argmax appears in mixed-integer optimization. In an online ecological gearshift strategy for a 2-speed electric vehicle, the control problem is reformulated as a mixed-integer model predictive control problem to minimize energy consumption. Outer convexification transforms integer variables into relaxed binary controls, and the neural network outputs logits for gear choices at each prediction step. Soft-argmax then maps these logits to relaxed binary variables $K$5, subject to $K$6 (Luo et al., 2024).

In this setting, soft-argmax serves two functions. It makes the relaxed decision layer differentiable, and with a larger $K$7 it pushes the relaxed binary variables close to $K$8 or $K$9. Because softmax outputs remain strictly between $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$0 and $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$1, an additional penalty term is introduced:

$$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$2

This discourages intermediate values during training, while deployment uses additional rounding to enforce $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$3 outputs (Luo et al., 2024).

The reported comparison with Bonmin illustrates the engineering trade-off. Relative to a rule-based method, the neural network optimizer with soft-argmax achieves energy savings of $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$4, while Bonmin achieves $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$5. The neural network’s average computation time is $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$6 ms per step on automotive dSPACE hardware, compared with $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$7 ms average and $$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$8 ms worst-case for Bonmin on a desktop CPU. The paper states that the neural network achieves almost the same binarized gearshift sequences as Bonmin, with only rare minor differences (Luo et al., 2024).

This use case places soft-argmax within a broader class of relaxations that bridge discrete optimization and gradient-based learning. The operator does not solve the original integer problem directly; rather, it provides a differentiable surrogate whose outputs can be sharpened and post-processed.

5. Sparse, structured, and generalized relaxations

A recurrent limitation of soft-argmax is that it typically produces dense outputs. In the probability-mapping literature, soft-argmax or softmax is described as returning a probability distribution with full support, while sparse alternatives can set some entries exactly to zero. The “sparsegen” framework unifies softmax, sparsemax, spherical softmax, and sum-normalization, and introduces sparsegen-lin and sparsehourglass to provide explicit control over the degree of sparsity (Laha et al., 2018).

Recent work has extended this line in several directions. “MultiMax” is a piece-wise differentiable function that modulates small and large entries differently, with the stated aim of suppressing irrelevant entries while preserving multimodality. The paper argues that standard SoftMax, as a smooth approximation to Argmax, distributes significant probability mass to residual entries, and that sparsity-oriented variants often do not preserve multi-modality (Zhou et al., 2024). A distinct approach is “Soft-Binary-Argmax,” which projects onto the hypersimplex

$$\mathrm{Softmax}(K z_i)=\frac{e^{K z_i}}{\sum_j e^{K z_j}}.$$9

thereby producing outputs that are close to binary and always sum to $z=[z_1,\ldots,z_{n_b}]$0 (Gomez et al., 26 Feb 2026).

On structured domains, soft-argmax has been generalized beyond categorical one-hot choices. Stochastic Softmax Tricks define a relaxation

$z=[z_1,\ldots,z_{n_b}]$1

where $z=[z_1,\ldots,z_{n_b}]$2 is the convex hull of a discrete structured set. This construction extends Gumbel-Softmax–style relaxations to subsets, spanning trees, arborescences, and other combinatorial spaces (Paulus et al., 2020). Related theoretical work on Perturb-Softmax and Perturb-Argmax analyzes completeness and minimality conditions, and positions soft-argmax-like perturbation models as universal parameterizations under suitable assumptions (Indelman et al., 2024).

Generalization has also entered reinforcement learning. The “f-softargmax” parameterization defines

$z=[z_1,\ldots,z_{n_b}]$3

so that the choice of $z=[z_1,\ldots,z_{n_b}]$4-divergence determines the geometry of the policy map. When paired with a regularizer induced by the same $z=[z_1,\ldots,z_{n_b}]$5-divergence, the resulting objective is stated to satisfy a Polyak-Lojasiewicz inequality, and Tsallis-based variants are reported to achieve polynomial sample complexity in contrast to the exponential complexity incurred by the standard softmax parameterization (Labbi et al., 18 Jan 2026).

6. Limitations, criticisms, and alternatives

A central criticism is that soft-argmax supervises the expected coordinate rather than the full shape of the probability map. In localization, the usual objective minimizes “the error of the expectation,” which constrains the mean but not the concentration or calibration of the map. Consequently, very different maps can share the same mean, and the network may produce unreliable or multi-modal heatmaps (Li et al., 2021).

This criticism becomes acute when inference uses hard argmax. In facial landmark detection, soft-argmax can yield the correct expected coordinate even if the heatmap is bimodal or multimodal and none of the peaks is at the ground truth. The expectation can be correct while the true $z=[z_1,\ldots,z_{n_b}]$6 prediction is wrong, creating a mismatch between training and inference. The same work argues that gradients can be small when peaks are far from the target or the distribution is overly spread, leading to slow convergence (Yang et al., 19 Aug 2025).

Several alternatives address these issues without relying on conventional soft-argmax training. “Sampling-argmax” minimizes “the expectation of the error” by introducing a continuous formulation of the output distribution and a differentiable sampling process. In reported experiments, sampling-argmax improved COCO Keypoint mAP from $z=[z_1,\ldots,z_{n_b}]$7 for soft-argmax to $z=[z_1,\ldots,z_{n_b}]$8 with a triangular basis, and improved calibration correlation from $z=[z_1,\ldots,z_{n_b}]$9 to $i$0 (Li et al., 2021). “Heatmap Regression without Soft-Argmax for Facial Landmark Detection” proposes a deep structured prediction framework and reports state-of-the-art performance on WFLW, COFW, and 300W, converging $i$1x faster during training while maintaining better/competitive accuracy (Yang et al., 19 Aug 2025).

Related objections arise in structured prediction and stereo matching. In structured NLP, soft-argmax or structured attention replaces hard argmax with a convex relaxation, but marginal inference can be intractable or $i$2-complete for many structures; SPIGOT instead preserves hard intermediate predictions and defines a proxy gradient via projection onto a relaxed feasible set (Peng et al., 2018). In stereo matching, soft-argmax-based disparity regression is described as prone to being multimodal due to absence of explicit constraint on the shape of the probability distribution, motivating “Sampling-Gaussian,” which the paper states improves accuracy on five baseline methods and two datasets without reduction in efficiency (Pan et al., 2024).

These criticisms do not negate the utility of soft-argmax. Rather, they delimit its operating regime. Soft-argmax remains attractive when differentiable coordinate extraction or differentiable relaxed decision-making is required, but multiple recent lines of work treat its dense outputs, weak shape supervision, or mismatch with hard inference as problems to be corrected rather than ignored.