Binary Gumbel-Softmax: Differentiable Binary Sampling

Updated 17 March 2026

Binary Gumbel-Softmax is a continuous relaxation of a Bernoulli variable using a logistic sigmoid applied to reparameterized Gumbel noise.
It introduces a temperature parameter τ that balances bias and variance, enabling smooth approximations for discrete stochastic processes.
Practical applications include variational autoencoders, selective computation, and combinatorial optimization in modern neural network architectures.

A binary Gumbel-Softmax, also known as Binary Concrete, is a continuous relaxation of a Bernoulli (binary) random variable that enables differentiable, sample-based stochasticity suitable for backpropagation in neural networks with discrete latent variables or binary control gates. The approach generalizes the Gumbel-Softmax trick—originally developed for multi-class categorical variables—to the binary (two-class) setting, producing a smooth, temperature-controlled approximation to a Bernoulli sample. This relaxation is central to enabling gradient-based learning in architectures where discrete 0/1 choices are required, such as stochastic binary networks, variational autoencoders with discrete latents, and selective neural networks. At its core, the binary Gumbel-Softmax combines Gumbel reparameterization with a logistic sigmoid relaxation, exposing a single temperature parameter τ by which the tradeoff between bias and variance in the gradient estimator is governed.

1. Mathematical Formulation and Sampling Procedure

Given a Bernoulli variable $x \sim \mathrm{Bernoulli}(p)$ with scalar logit $n = \log\left(\frac{p}{1-p}\right)$ , the exact binary sample can be re-expressed via the Gumbel-Max trick as

$x = 1[n + (g_1 - g_0) \ge 0], \qquad g_i \sim \mathrm{Gumbel}(0,1).$

To obtain a differentiable relaxation, a temperature parameter $\tau > 0$ is introduced and the indicator is replaced by a sigmoid:

$\tilde{x} = \sigma\left(\frac{n + (g_1 - g_0)}{\tau}\right) = \frac{1}{1 + \exp\left(-\frac{n + g_1 - g_0}{\tau}\right)}.$

Here, $g_0, g_1$ are drawn as $g_i = -\log(-\log U_i)$ with $U_i \sim \mathrm{Uniform}(0,1)$ , so $g_1 - g_0 \sim \mathrm{Logistic}(0,1)$ . This formula yields a relaxed, continuous value $\tilde{x} \in (0,1)$ that converges in distribution to the discrete Bernoulli variable as $\tau \to 0^+$ . The stochasticity is fully reparameterized, allowing gradients to propagate through $n$ and, by extension, through any parameters upon which $p$ depends (Shekhovtsov, 2021, Salem et al., 2022, Jang et al., 2016, Joo et al., 2020).

2. Gradient Estimator: Pathwise Differentiation and Properties

The binary Gumbel-Softmax estimator leverages pathwise (reparameterization) derivatives. For a scalar function $f$ (differentiable on $[0,1]$ ), the single-sample gradient with respect to $n$ is:

$\hat{g}_{GS} = \frac{\partial}{\partial n} f(\tilde{x}) = f'(\tilde{x}) \cdot \frac{1}{\tau} \tilde{x}(1 - \tilde{x}).$

This estimator is easily implemented by forward-propagating a realization of $\tilde{x}$ and back-propagating through the sigmoid relaxation. Conditioning on the sampled Gumbel noise, the derivative with respect to the Bernoulli probability $p$ is:

$\frac{\partial \tilde{x}}{\partial p} = \frac{1}{p(1-p)} \frac{\partial \tilde{x}}{\partial n}.$

The same pathwise construction holds for any downstream computational graph fed by $\tilde{x}$ . For vector-valued logit input or minibatches, all computations naturally vectorize (Shekhovtsov, 2021, Salem et al., 2022, Jang et al., 2016, Joo et al., 2020).

In the "straight-through" Gumbel-Softmax variant (ST-GS), the forward pass uses the thresholded (discrete) $x_{hard}=1[n+z\geq0]$ sample but backpropagation still uses the Jacobian of the soft relaxation, retaining pathwise gradients (Shekhovtsov, 2021, Salem et al., 2022).

3. Bias–Variance Tradeoff, Asymptotic Limits, and Theoretical Analysis

A central feature of the binary Gumbel-Softmax gradient estimator is its temperature-controlled bias–variance tradeoff:

Bias: As proven in [(Shekhovtsov, 2021), Prop. 1], the estimator is biased except for linear $f$ , with bias $B(\tau) = O(\tau)$ as $\tau \to 0$ . For linear $f$ , bias reduces to $O(\tau^2)$ . Explicit expansions show $E[\hat{g}_{GS}] = f(1) - f(0) + O(\tau)$ .
Variance: The estimator's variance diverges as $\mathrm{Var}[\hat{g}_{GS}] = O(1/\tau)$ for small temperature, a result of rare but large gradient spikes. For fixed $\epsilon > 0$ , the tail probability $P(|\hat{g}_{GS}|\geq\epsilon) = O(\tau \log(1/\epsilon))$ as $\tau \to 0$ [(Shekhovtsov, 2021), Prop. 2/4]. As a consequence, annealing the temperature too low leads to intolerable estimator variance and instability in deep or wide architectures.

This bias–variance structure dictates that practical usage requires $\tau$ to be kept high enough (typically $\tau \approx 0.5-1$ ), sacrificing some bias for manageable variance (Shekhovtsov, 2021, Salem et al., 2022, Jang et al., 2016).

4. Algorithmic Implementations and Pseudocode

The binary Gumbel-Softmax is implemented in standard autodiff libraries (PyTorch, TensorFlow, JAX) via minimal primitives. The canonical pseudocode for a single binary variable is:

logit = np.log(p) - np.log(1 - p)
u0 = np.random.uniform(0,1)
u1 = np.random.uniform(0,1)
g0 = -np.log(-np.log(u0))
g1 = -np.log(-np.log(u1))
z = g1 - g0
x_tilde = 1 / (1 + np.exp(-(logit + z)/tau))
L = f(x_tilde)

For the ST-GS variant:

1	x_hard = float(logit + z >= 0) # forward uses x_hard, backward uses x_tilde

Batch and vectorized forms are trivial extensions; PyTorch's torch.nn.functional.gumbel_softmax and TensorFlow Probability's RelaxedBernoulli provide direct support (Salem et al., 2022). Typical temperature schedules involve exponential or multi-step annealing (e.g., $\tau_{t+1} = \tau_t \times r$ with $r\approx0.98$ , stopping at final $\tau \approx 0.1 - 0.5$ ) (Salem et al., 2022, Jang et al., 2016, Joo et al., 2020).

5. Comparative Analysis with Alternative Estimators

Multiple estimator variants have been benchmarked against the binary Gumbel-Softmax:

Estimator	Bias	Variance	Key Properties
Binary Gumbel-Softmax	$O(\tau)$	$O(1/\tau)$	Biased; reparameterizable; variance explodes $\tau \rightarrow 0$
ST-Gumbel-Softmax	$O(\tau)$	$O(1/\tau)$	Hard forward, soft backward; still fundamentally biased
Straight-Through (ST)	Small, fixed	Moderate	Simple, non-reparameterized, often outperforms GS in SBNs
IR/DARN	Zero ( $f$ quadratic)	Large at $p\to0,1$	Unbiased for quadratic $f$ , high variance near 0/1 probabilities
FouST variants	Model-dependent	Model-dependent	Importance reweighting, Taylor/Fourier disc., surrogate-bias

BayesBiNN, while nominally using binary Gumbel-Softmax with extremely low $\tau\approx 10^{-10}$ , effectively degenerates to a deterministic ST update with weight decay, thus eliminating variance but incurring higher bias than stochastic ST (Shekhovtsov, 2021). The IR/DARN estimator achieves unbiasedness only for quadratic $f$ , but its variance can diverge for extreme $p$ (Shekhovtsov, 2021).

Empirically, plain Gumbel-Softmax provides low-variance gradients and outperforms older methods at moderate $\tau$ , but in deep binary networks the ST estimator is often more reliable due to GS's variance pathologies at small $\tau$ (Shekhovtsov, 2021, Salem et al., 2022, Jang et al., 2016).

6. Practical Considerations, Usage, and Limitations

Binary Gumbel-Softmax is widely used whenever a model requires end-to-end differentiable binary decisions: selective prediction (coverage-risk tradeoff), learned binary masks, conditional computation gates, variational inference in SBNs, and combinatorial optimization tasks (Salem et al., 2022, Li et al., 2020).

However, practitioners must balance bias and variance via $\tau$ tuning. Aggressive annealing toward $0$ is not safe in deep or large models, as estimator variance becomes unacceptably large, causing gradients to vanish or explode and optimization to stall (Shekhovtsov, 2021). Straight-through Gumbel-Softmax reduces forward-pass bias but cannot eliminate the underlying bias-variance tradeoff governed by $\tau$ . Effective schedules typically start with $\tau \approx 1$ and anneal toward $0.1$–$0.5$ (Jang et al., 2016, Joo et al., 2020).

Contemporary work has produced improved relaxations with lower bias, such as piecewise-linear relaxations and improved continuous relaxation (ICR) estimators. These match the variance profile of the standard estimator while achieving unbiasedness or reduced bias, leading to improved empirical performance in discrete VAEs and combinatorial problems (Andriyash et al., 2018).

7. Empirical Results and Application Domains

Empirical studies benchmarked binary Gumbel-Softmax and its variants on structured output prediction, variational autoencoders, selective networks, and combinatorial graph problems:

On structured-output prediction and SBNs, plain Gumbel-Softmax ( $\tau=1$ ) outperformed competing estimators like Score-Function, NVIL, DARN, and MuProp (Jang et al., 2016).
For selective networks, the relaxed selection mechanism enabled coverage-constrained, risk-minimizing operation in end-to-end differentiable frameworks (Salem et al., 2022).
In combinatorial optimization (e.g., maximum independent set, minimum vertex cover), the binary Gumbel-Softmax provided competitive or superior solutions compared to simulated annealing, greedy methods, and reinforcement-learned policies, with substantially reduced compute time on large-scale problems (Li et al., 2020).
On standard VAEs with binary latents, Gumbel-Softmax and improved relaxations achieved state-of-the-art or near–state-of-the-art log-likelihoods, contingent on careful $\tau$ tuning and sometimes outperformed by unbiased or lower-bias relaxations (Andriyash et al., 2018).

Key limitations remain the estimator's inherent bias (vanishing only as $O(\tau)$ ) and its variance blow-up at low temperatures, which precludes aggressive annealing in practical deep learning contexts (Shekhovtsov, 2021). This explains why, despite theoretical appeal, the binary Gumbel-Softmax rarely supersedes the straight-through estimator as the practical default for discrete stochastic networks.