---
title: Ensemble Gumbel-Softmax (EGS) in NAS
url: https://www.emergentmind.com/topics/ensemble-gumbel-softmax-egs
type: topic
---

# Ensemble Gumbel-Softmax (EGS) in NAS

The Ensemble Gumbel-Softmax (EGS) estimator is a technique developed to improve the effectiveness and efficiency of differentiable neural architecture search (NAS) by enabling the simultaneous optimization of network architecture and parameters in a fully differentiable, end-to-end manner. EGS builds upon the standard Gumbel-Softmax reparameterization by aggregating several independent relaxed samples, thereby supporting the selection of multiple operations per decision point and providing a low-variance, differentiable gradient estimator suited for gradient-based NAS frameworks [1905.01786].

## 1. Standard Gumbel-Softmax and Its Relaxation

The foundation of EGS lies in the Gumbel-Softmax, also known as the Concrete distribution. Given a categorical distribution over $K$ classes parameterized by unnormalized scores $\alpha_1,\ldots,\alpha_K$, sampling is achieved by the Gumbel-Max trick:
$$
c = \arg\max_k ( \log\alpha_k + G_k )
$$
with $G_k \sim \text{Gumbel}(0,1)$ for each $k$. Since $\arg\max$ is non-differentiable, Gumbel-Softmax introduces a softmax relaxation in the log-domain, yielding a differentiable probability vector
$$
y_k = \frac{\exp((\log\alpha_k + G_k)/\tau)}{\sum_j \exp((\log\alpha_j + G_j)/\tau)}
$$
where $G_k = -\log(-\log U_k)$, $U_k \sim \text{Uniform}(0,1)$ and $\tau>0$ is a temperature hyperparameter. As $\tau \to 0^+$, the sample approaches a one-hot vector; as $\tau \to \infty$, it resembles a uniform distribution. Differentiability with respect to $\log\alpha_k$ allows gradients to flow through the architecture selection process.

## 2. Ensemble Sampling Mechanism

While standard Gumbel-Softmax ensures single-operation sampling per edge, EGS enables multi-operation selection essential for richer architectural search spaces. EGS produces $M$ independent Gumbel-Softmax samples $y^{(1)},\ldots,y^{(M)} \in \Delta^{K-1}$ and aggregates them into a binary decision vector $\hat{y} \in \{0,1\}^K$ by
$$
\hat{y}_k = \max_{m=1..M} [ \text{hard\_onehot}(y^{(m)}) ]_k
$$
Alternatively, an ensemble-averaged soft assignment is constructed
$$
\bar{y} = \frac{1}{M} \sum_{m=1}^M y^{(m)}
$$
with discretization performed by retaining the top $M$ entries or via thresholding, thus allowing up to $M$ operations to be selected per edge.

## 3. Differentiability, Gradient Estimation, and Variance Reduction

EGS targets minimization of the expected loss over sampled architectures:
$$
J(\alpha, w) = E_{A \sim p_\alpha}[ L(w; A) ]
$$
where $A$ is a discrete architecture sampled by EGS, $w$ are network weights, and $\alpha$ the architecture logits. The reparameterization trick ensures that each Gumbel draw is a deterministic function of uniform randomness, rendering $y^{(m)}(\alpha, U^{(m)})$ differentiable with respect to $\alpha$. Instead of using a single sample gradient, EGS averages over $M$ samples to compute
$$
\bar{g} = \frac{1}{M} \sum_m g^{(m)}
$$
with $g^{(m)} = \partial L / \partial y^{(m)}$, and gradients are backpropagated through $\bar{y}$. The sample mean reduces variance by a factor of $M$, leveraging the independence of $y^{(m)}$ conditional on $\alpha$.

## 4. Algorithmic Integration into Differentiable NAS

The EGS procedure can be integrated into standard NAS workflows that involve joint or bi-level optimization of weights $w$ and architecture parameters $\alpha$. The NAS search loop proceeds as follows:

```python
for epoch in range(1, T+1):
    for batch in D_train:
        for edge in cell:
            # EGS sampling
            for m in range(1, M+1):
                G_k = -log(-log U_k)
                y_e^{(m)}[k] = softmax((log α_e[k] + G_k) / τ)
            ȳ_e = (1/M) sum_m y_e^{(m)}
            o_e(x) = sum_k ȳ_e[k] * O_k(x)
        output = model_w(x, {ȳ_e})
        loss = ℓ(output, y)
        # Backpropagation through ȳ_e
        w ← w – η_w ∇_w loss
        α ← α – η_α ∇_α loss
    # Optionally adjust τ schedule
# Discretize architecture by final α and select top-M per edge.
```
At search termination, discrete architectures are fixed by selecting the top-$M$ operations per edge according to logits $\alpha$.

## 5. Hyperparameters, Computational Complexity, and Theoretical Guarantees

EGS introduces key hyperparameters: ensemble size $M$ and temperature $\tau$.

- **Ensemble size $M$**: Dictates the number of GS samples per edge, controlling search space expressivity and gradient variance. Proposition 3 confirms EGS can represent all binary patterns with up to $M$ active operations per edge, with search capacity scaling as $\mathcal{O}(K^M)$.
- **Temperature $\tau$**: Lower $\tau$ produces near one-hot outputs, favoring discrete architectures; higher $\tau$ yields softer selections and easier optimization. Annealing $\tau$ from approximately $1.0$ to $0.1$ during search is common.
- **Computational cost**: Sampling and combining across $M$ GS samples incurs cost $\mathcal{O}(M \cdot K)$ per edge, or $\mathcal{O}(E \cdot M \cdot K)$ for a cell with $E$ edges. Empirically, $M = 4 \ldots 8$ provides sufficient coverage at a modest overhead.
- **Theoretical properties**: Proposition 2 states EGS is a monotonic increasing set function of the selection probabilities, and variance of the estimator scales inversely with $M$.

## 6. Empirical Evaluation Across Tasks

EGS has demonstrated competitive or superior performance in canonical NAS benchmarks at low computational cost. For CIFAR-10, DARTS-EGS (with $M=7$) achieves $2.79\%$ error in $1$ GPU-day, improving upon DARTS and SNAS. On ImageNet (mobile setting), DARTS-EGS attains $24.9\%$ top-1 accuracy in $1.5$ days, outperforming DARTS and NASNet at orders-of-magnitude lower search cost. In PTB language modeling, DARTS-EGS ($M=7$) reaches $57.1 / 55.3$ perplexity in $0.5$ days, while transfer to WT2 yields $66.5 / 64.2$ perplexity. Ablation on $M$ reveals monotonically improved results as ensemble size grows:
| M | CIFAR-10 Error (%) |
|---|-------------------|
| 1 | 3.38              |
| 4 | 3.05              |
| 7 | 2.79              |
| 9 | 2.73              |

Consistently, larger $M$ enhances final model accuracy and reduces gradient variance at acceptable computational cost.

## 7. Context, Significance, and Integration with NAS Frameworks

EGS generalizes prior gradient-based NAS searchers by supporting multi-operation selection and stabilizing gradients via ensembling, thereby extending the utility of continuous relaxation frameworks (DARTS-style) where joint optimization of discrete architecture and network weights is required. The capacity to sample exponentially more architectures and reduce estimator variance positions EGS as a principal technique for efficient, scalable architecture search. EGS is compatible with both joint and bi-level optimization schemes, and naturally integrates into feed-forward search spaces with arbitrary connectivity.

Ensemble Gumbel-Softmax thus extends the classical Gumbel-Softmax by facilitating diverse, gradient-efficient search over architectural building blocks, yielding state-of-the-art performance on both image and language tasks at practical search cost [1905.01786].

Source: https://www.emergentmind.com/topics/ensemble-gumbel-softmax-egs