---
title: 'Soft-Argmax: Differentiable Approximation to Argmax'
url: https://www.emergentmind.com/topics/soft-argmax
type: topic
---

# Soft-Argmax: Differentiable Approximation to Argmax

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 [2402.18076, 2205.13271]. 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 [2110.08825].

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

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
$$
x_k^0=\sum_{i=1}^{w^*}\sum_{j=1}^{h^*}\mathcal{A}_k(i,j)x(i),\qquad
y_k^0=\sum_{i=1}^{w^*}\sum_{j=1}^{h^*}\mathcal{A}_k(i,j)y(j),
$$
where $\mathcal{A}_k$ is a spatial softmax of attention logits [2205.13271].

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:
$$
\Theta^{(n+1)}=\Theta^{(n)}-\eta \nabla_{\Theta^{(n)}}\mathcal{L}.
$$
Because the soft-argmax layer is smooth, the chain rule propagates gradients through all layers, including the relaxed discrete-decision stage [2402.18076].

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

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

## 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
$$
\hat{y}=\mathrm{soft\mbox{-}argmax}(\pi)=\sum_i \pi_{y_i} y_i,
$$
or, in heatmap regression for landmark detection,
$$
\tilde{\mathbf{y}}_n=\texttt{Soft-argmax}(\hat{h}_n)\triangleq \sum_{\mathbf{y}'_n}\mathbf{y}'_n\cdot \mathrm{Softmax}(\hat{h}_n)[\mathbf{y}'_n].
$$
This pattern appears in 2D/3D human pose estimation, facial landmark localization, stereo matching, and object keypoint estimation [2110.08825, 2508.14929].

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

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

## 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 $b_{i,k}\in[0,1]$, subject to $\sum_i b_{i,k}=1$ [2402.18076].

In this setting, soft-argmax serves two functions. It makes the relaxed decision layer differentiable, and with a larger $K$ it pushes the relaxed binary variables close to $0$ or $1$. Because softmax outputs remain strictly between $0$ and $1$, an additional penalty term is introduced:
$$
\xi(B_j)=B_j(1-B_j).
$$
This discourages intermediate values during training, while deployment uses additional rounding to enforce $\{0,1\}$ outputs [2402.18076].

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 $6.02\%$, while Bonmin achieves $6.57\%$. The neural network’s average computation time is $0.045$ ms per step on automotive dSPACE hardware, compared with $229$ ms average and $2315$ 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 [2402.18076].

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

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 [2406.01189]. A distinct approach is “Soft-Binary-Argmax,” which projects onto the hypersimplex
$$
\Delta_k^n=\{\mathbf{y}\in[0,1]^n\mid \sum_{i=1}^n y_i=k\},
$$
thereby producing outputs that are close to binary and always sum to $k$ [2602.23336].

On structured domains, soft-argmax has been generalized beyond categorical one-hot choices. Stochastic Softmax Tricks define a relaxation
$$
X_t=\arg\max_{x\in P}\{U^\top x-t f(x)\},
$$
where $P=\operatorname{conv}(\mathcal{X})$ 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 [2006.08063]. 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 [2406.02180].

Generalization has also entered reinforcement learning. The “f-softargmax” parameterization defines
$$
\softargmax_f(x,q)=\arg\max_{p\in\Delta}\{\langle p,x\rangle-D_f(p\|q)\},
$$
so that the choice of $f$-divergence determines the geometry of the policy map. When paired with a regularizer induced by the same $f$-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 [2601.12604].

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

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 $\arg\max$ 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 [2508.14929].

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 $64.5$ for soft-argmax to $69.8$ with a triangular basis, and improved calibration correlation from $0.233$ to $0.432$ [2110.08825]. “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 $2.2$x faster during training while maintaining better/competitive accuracy [2508.14929].

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 $\#P$-complete for many structures; SPIGOT instead preserves hard intermediate predictions and defines a proxy gradient via projection onto a relaxed feasible set [1805.04658]. 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 [2410.06527].

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.

Source: https://www.emergentmind.com/topics/soft-argmax