On The Representation Properties Of The Perturb-Softmax And The Perturb-Argmax Probability Distributions

## Abstract

Abstract The Gumbel-Softmax probability distribution allows learning discrete tokens in generative learning, while the Gumbel-Argmax probability distribution is useful in learning discrete structures in discriminative learning. Despite the efforts invested in optimizing these probability models, their statistical properties are under-explored. In this work, we investigate their representation properties and determine for which families of parameters these probability distributions are complete, i.e., can represent any probability distribution, and minimal, i.e., can represent a probability distribution uniquely. We rely on convexity and differentiability to determine these statistical conditions and extend this framework to general probability models, such as Gaussian-Softmax and Gaussian-Argmax. We experimentally validate the qualities of these extensions, which enjoy a faster convergence rate. We conclude the analysis by identifying two sets of parameters that satisfy these assumptions and thus admit a complete and minimal representation. Our contribution is theoretical with supporting practical evaluation.

## 1 Introduction

Learning over discrete probabilistic models is an active research field with numerous applications. Examples include learning probabilistic latent representations of semantic classes or beliefs (Kingma et al., ; Mnih & Gregor, ; Salakhutdinov & Hinton, ). The Gumbel-Argmax and Gumbel-Softmax probability distributions are widely applied in machine learning to model and analyze discrete probability distributions.

The Gumbel-Argmax is an equivalent representation of the softmax operation and plays a key role in “follow the perturb-leader” family of algorithms in online learning (Hannan, ; Kalai & Vempala, , ; Rakhlin et al., ). Its extension to Gaussian-Argmax allows better bounds on their gradients and consequently provides better regret bounds in linear and high-dimensional settings Abernethy et al. (, ); Cohen & Hazan (). The argmax operation allows for efficient sampling, making the Perturb-Argmax probability models pivotal in discriminative learning algorithms of high-dimensional discrete structures (Berthet et al., ; Pogančić et al., ; Song et al., ; Cohen Indelman & Hazan, ; Niculae et al., ). The Gumbel-Softmax (or the Concrete distribution) probability distribution, which replaces the argmax operation with a softmax operation is easier to optimize and therefore plays a key role in generative learning models (Jang et al., ; Maddison et al., ). The discrete nature of these probability models provides a natural representation of concepts, e.g., in
zero-shot text-to-image generation (DALL-E) (Ramesh et al., ). While the Gumbel-Argmax and Gumbel-Softmax probability distributions are widely applied in machine learning, their statistical representation is still under-explored.

In this work, we investigate the representation properties of the Gumbel-Argmax and Gumbel-Softmax probability distributions. We aim to determine for which families of parameters these distributions are complete and minimal. A distribution is considered complete if it can represent any probability distribution, and minimal if it can uniquely represent a probability distribution. Our statistical investigation realizes the Gumbel-Argmax and the Gumbel-Softmax probability distributions as gradients of respective convex functions. In Theorems 4.1 and 4.3, we prove the conditions under which $\Theta\subset\mathbb{R}^{d}$ is a complete and minimal representation of the Gumbel-Softmax probability distribution and generalize these results to other random perturbations, e.g., Gaussian-Softmax or more generally, Perturb-Softmax probability distributions. We also extend these methods to Perturb-Argmax probability distributions, for which Gumbel-Argmax is a special case, and state the conditions for which it is complete and minimal in Theorems 5.1 and 5.3. Our findings are illustrated in Figure 1.

We begin by introducing the notation relating parameters and the relevant probability distributions in Section 3. Subsequently, we investigate the Perturb-Softmax probability models as gradients of the expected log-sum-exp convex function and prove their completeness by connecting their gradients to the relative interior of the probability simplex. In Section 4, we determine the minimality of Perturb-Softmax by the strict convexity of the expected log-sum-exp when restricted to the respective parameter space. We then investigate the Perturb-Argmax probability models as sub-gradients of the expected-max convex function and establish the conditions for which their parameter space is complete and minimal, see Section 5. Finally, we empirically demonstrate the qualities of Perturb-Softmax extension in generative and discriminative learning setting, showing improved convergence of Gaussian-Softmax over Gumbel-Softmax beyond the linear high-dimensional setting that was investigated in online learning Abernethy et al. (); Cohen & Hazan ().

Figure: (a) An illustration of the representation properties of the Perturb-Softmax. Its representation is complete for any perturbation distribution under the conditions in Theorem 4.1. The mapping is one-to-one whenever the representation parameters are not linearly constrained (full mapping), and single-valued otherwise (green dashed mapping).
Refer to caption: x1.png

## 2 Related work

The exponential family realized by the softmax operation over its parameters is extensively used in machine learning. However, sampling high-dimensional models is challenging due to its normalizing factor (Geman & Geman, ; Goldberg & Jerrum, ). The Gumbel-Argmax probability distribution measures the stability of the argmax operation over Gumbel random variables. It serves as an equivalent representation of softmax operation, thereby enabling efficient sampling from the exponential family (Gumbel, ; Luce, ). In the context of machine learning, the Gumbel-Argmax probability models underlie the ”follow the perturbed-leader” online learning algorithm (Hannan, ; Kalai & Vempala, , ). The Gaussian-Argmax probability models extend the ”follow the perturbed-leader” family of online algorithms and improve their regret bounds (Rakhlin et al., ; Abernethy et al., , ; Cohen & Hazan, ). Our work complements these studies by exploring the statistical properties of Gumbel-Argmax and Gaussian-Argmax. We prove the conditions under which Perturb-Argmax probability models are complete, making them suitable for use in machine learning, and when they are minimal, meaning they can uniquely identify a probability distribution from Perturb-Argmax probability models.

The Gumbel-Softmax probability models were introduced as an alternative to the exponential family and its Gumbel-Argmax equivalent in the context of generative learning (Maddison et al., ; Jang et al., ). This alternative allows for efficient sampling, making it highly effective for learning with stochastic gradient methods. The discrete nature of the Gumbel-Softmax sampling has been utilized to tokenize the visual vocabulary in the celebrated zero-shot text-to-image generation, DALL-E (Ramesh et al., ). The Gaussian-Softmax probability models were introduced in variational auto-encoders as their closed-form KL-divergence makes it easier to realize as regularization (Potapczynski et al., ). Our work extends these works and sets the statistical properties of Gumbel-Softmax and Gaussian-Softmax. We also show that Gaussian-Softmax enjoys faster convergence as the Gaussian distribution decays faster than the Gumbel distribution when approaching infinity.

Berthet et al. () introduced a framework for optimizing discrete problems based on Perturb-Argmax probability models. This framework applies to discriminative learning using the Fenchel-Young losses and relies on convexity to propagate gradients over discrete choices. Similar to our approach, this work adopts a general view and utilizes convexity to explore the gradient properties of Perturb-Argmax models. Our work differs in that we use convexity and differentiability to investigate their statistical representation properties, specifically when these models are complete and minimal. Other methods include blackbox differentiation based on gradients of a surrogate linearized loss (Pogančić et al., ), the direct loss minimization technique (Hazan et al., ; Keshet et al., ; Song et al., ; Cohen Indelman & Hazan, ) based on gradients of the expected discrete
loss, and entropy regularization techniques (Niculae et al., ; Martins & Astudillo, ). Unlike these methods, we focus on studying the statistical properties of randomized discrete probability models rather than on optimization frameworks.

## 3 Background

We denote by $\Delta$ the probability simplex, i.e., the set of all probabilities over $d$ discrete events, namely $\Delta\triangleq\{p\in\mathbb{R}^{d}:p(i)\geq 0,\sum_{i=1}^{d}p(i)=1\}$. A parameterized discrete probability distribution $p_{\theta}(i)\in\Delta$ is determined by its parameters $\theta\in\Theta^{d}$ that reside in the Euclidean space $\Theta^{d}\subseteq\mathbb{R}^{d}$.

### 3.1 Completeness and minimality of the Softmax operation

The softmax operation $\mbox{softmax}:\mathbb{R}^{d}\rightarrow\Delta$ is the standard mapping from the set of parameters $\Theta$ to the probability simplex $\Delta$. Formally, when we define $p^{sm}_{\theta}$ by softmax relation

$$ $p^{sm}_{\theta}\triangleq\mbox{softmax}(\theta)\triangleq\left(\frac{e^{\theta _{1}}}{\sum_{j=1}^{d}e^{\theta_{j}}},...,\frac{e^{\theta_{d}}}{\sum_{j=1}^{d}e ^{\theta_{j}}}\right).$ (1) $$

A parameterized family of distributions $\mathbb{P}_{\Theta}\triangleq\{p_{\theta}:\theta\in\Theta\}$ is called complete if for every $p\in\Delta$ there exists $\theta\in\Theta$ such that $p_{\theta}=p$. Alternatively, the mapping from the parameters to their probabilities is onto the probability simplex (surjective). Similarly, a parameterized family of distributions is called minimal, if there is one-to-one mapping between its parameters and their corresponding probability distributions (injective). Formally, $p_{\theta}\neq p_{\tau}$ if and only if $\theta\neq\tau$. A complete and minimal mapping is also identifiable, i.e., for every probability $p\in\Delta$ one can identify the unique parameters $\theta\in\Theta$ for which $p=p_{\theta}$.

The identifiability of the softmax mapping was explored in the context of the exponential family of distributions (Wainwright & Jordan, ). One can verify that the set of parameters $\Theta=\mathbb{R}^{d}$ is complete: for every $p\in\Delta$, one can set $\theta=\log p$, for which $p=softmax(\theta)$. However, the set $\Theta=\mathbb{R}^{d}$ is not minimal, as the two parameter vectors $\theta$ and $\theta+c1$ both realize the same probability, i.e., $\mbox{softmax}(\theta)=\mbox{softmax}(\theta+c1)$. Conversely, a set of parameters $\Theta$ is minimal for the softmax mapping if there are no $\theta,\tau\in\Theta$ for which $\theta_{i}=\tau_{i}=c$ for every $i=1,...,d$.(^1^11If $\mbox{softmax}(\theta)=\mbox{softmax}(\tau)$, and $p_{\theta}=\mbox{softmax}(\theta)$, $p_{\tau}=\mbox{softmax}(\tau)$, then $p_{\theta}(i)/p_{\tau}(i)=1$ for every $i$, and $\log p_{\theta}(i)-\log p_{\tau}(i)=0$. From the softmax mapping, this translates to $\theta_{i}-\tau_{i}=c$ for every $i$ while $c=\log(\sum_{j}e^{\theta_{j}})-\log(\sum_{j}e^{\tau_{j}})$.) Consequently, identifiable sets of parameters for the softmax operation can be $\Theta=\{\theta\in\mathbb{R}^{d}:\sum_{j}\theta_{j}=0\}$ or $\Theta=\{\theta\in\mathbb{R}^{d}:\theta_{1}=0\}$. In Appendix 7.2.1 we prove that these sets are both complete and minimal.

### 3.2 Gumbel-Softmax and Gumbel-Argmax probability distributions

The Gumbel-Softmax probability distribution emerged as a smooth approximation of the Gumbel-Argmax representation of $p^{sm}_{\theta}$. We turn to describe the Gumbel-Argmax and Gumbel-Softmax discrete probability distributions.

The Gumbel distribution is a continuous distribution whose probability density function is $\hat{g}(t)=e^{-e^{-(t+c)}}$, where $c\approx 0.5772$ is the Euler-Mascheroni constant. We denote by $\gamma=(\gamma_{1},...,\gamma_{d})$ the vector of $d$ independent random variables that follow the Gumbel distribution law and by $g(\gamma)=\prod_{i=1}^{d}\hat{g}(\gamma_{i})$ the probability density function (pdf) of the independent Gumbel distribution.

We denote by $p^{gam}_{\theta}$ the Gumbel-Argmax probability distribution, which relies on a one-hot representation of the maximal argument. The indicator function $1[condition]$ equals one when the condition holds and zero otherwise. Then, the Gumbel-Argmax probability distribution takes the form:

$$ $\displaystyle p^{gam}_{\theta}$ $\displaystyle\triangleq$ $\displaystyle\mathbb{E}_{\gamma\sim g}[\arg\max(\theta+\gamma)]$ (2) $\displaystyle p^{gam}_{\theta}(i)$ $\displaystyle\triangleq$ $\displaystyle\mathbb{E}_{\gamma\sim g}[i=\arg\max(\theta+\gamma)]\triangleq \int_{\mathbb{R}^{d}}g(\gamma)1[i=\arg\max(\theta+\gamma)]d\gamma$ (3) $$

The fundamental theorem of extreme value statistics asserts the equivalence between the softmax distribution in Equation 4 and the Gumbel-Argmax distribution in Equation 2, namely $p^{sm}_{\theta}=p^{gam}_{\theta}$, cf. Gumbel (). Therefore, the statistical representation properties of completeness and minimality of the softmax operation are identical to the statistical properties of the Gumbel-Argmax probability distribution. Unfortunately, the argmax operation is non-smooth and requires special treatment when used in learning its parameters using gradient methods.

The Gumbel-Softmax probability distribution $p^{gsm}_{\theta}(i)$ was developed as a smooth approximation of its Gumbel-Argmax counterpart:

$$ $\displaystyle p^{gsm}_{\theta}$ $\displaystyle\triangleq$ $\displaystyle\int_{\mathbb{R}^{d}}g(\gamma)\mbox{softmax}(\theta+\gamma)d \gamma\triangleq\mathbb{E}_{\gamma\sim g}[\mbox{softmax}(\theta+\gamma)]$ (4) $$

We note that we define $p^{gsm}_{\theta}$ as a $d$-th dimensional vector, and the integral with respect to $\gamma$ (or expectation) is taken with respect to each coordinate of the softmax operation.

One can verify that the Gumbel-Softmax is indeed a probability distribution: (i) it is non-negative since it is the average of the non-negative softmax operation, and (ii) it sums up to unity since it is an average of the softmax operations $\mbox{softmax}(\theta+\gamma)$ that sum up to unity for every $\gamma$.

### 3.3 Differentibility properties of convex functions

We investigate the representation properties of the Gumbel-Softmax and Gumbel-Argmax probability models when they result from gradients of multivariate functions over their set of parameters $\Theta$.

The softmax function is the gradient of the log-sum-exp function:

$$ $\mbox{softmax}(\theta)=\nabla\log\left(\sum_{i=1}^{d}e^{\theta_{i}}\right).$ (5) $$

As shown in Section 3.1 this gradient mapping is complete and minimal over a convex subsetset $\Theta\subset\mathbb{R}^{d}$. The extension of this argument to Gumbel-Softmax requires notions of convexity, covered in Appendix 7.1.1.
The (sub)differential of different convex functions $f(\theta)$ is used in our study as a (multi-valued) mapping between the convex set of the primal domain $\Theta$ and its dual domain ${\cal P}$, which is the Gumbel-Softmax or the Gumbel-Argmax probability model.
We define the conditions for which:

- I
$\partial f$ is a single-valued mapping, i.e., $\partial f=\nabla f$. In this case, for every $\theta\in\Theta$ matches a single $\nabla f(\theta)\in{\cal P}$. If this property does not hold then the parameters $\theta$ that generate a probability $p$ are not identifiable under this mapping.
- II
The gradient mapping $\nabla f(\theta)$ is onto the probability simplex $\Delta$. In this case, the set of parameters $\Theta$ is complete, i.e., it can represent (and learn) any probability $p\in\Delta$ using its gradients $\nabla f(\theta)$.
- III
The gradient mapping $\nabla f(\theta)$ is one-to-one. In this case, the set of parameters $\Theta$ is minimal, i.e., there are no two parameters $\theta,\tau$ that represent the same probability distribution $p\in{\cal P}$.

Our framework allows establishing these relations for any random perturbations, which we refer to as Perturb-Softmax and Perturb-Argmax.

## 4 Perturb-Softmax probability distributions

In this section, we explore the statistical representation properties of the Perturb-Softmax operation as a generalization of the Gumbel-Softmax operation.
Our exploration emerges from the connection between the softmax operation and the log-sum-exp convex function, as described in Equation 5. We establish a similar relation between perturb-log-sum-exp and perturb-softmax:

$$ $\displaystyle f(\theta)$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\gamma}\left[\log\left(\sum_{i=1}^{d}e^{\theta_{i}+ \gamma_{i}}\right)\right]$ (6) $\displaystyle\nabla f(\theta)$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\gamma}[\mbox{softmax}(\theta+\gamma)]$ (7) $$

The function $f(\theta)$ is defined for any random perturbation $\gamma$, whether $\gamma$ values are from a discrete, a bounded, or an unbounded set (^2^22Formally, for unbounded random perturbations $\gamma$ we restrict ourselves to probability density functions for which $f(\theta)<\infty$.). The function $f(\theta)$ is differentiable since it is the expectation of the differentiable log-sum-exp function and $\nabla f(\theta)$ is attained by the Leibniz rule for differentiation under the integral sign(^3^33Formally, $\nabla f(\theta)$ is finite whenever the dominant convergence theorem holds. For unbounded $\gamma$ this holds for any probability density function $p(\gamma)$ for which $\lim_{\gamma\rightarrow\infty}p(\gamma)\log\left(\sum_{i=1}^{d}e^{\theta_{i}+ \gamma_{i}}\right)=0$. This happens for Gumbel, Gaussian, and other standard probability density functions.). Also, the function $f(\theta)$ is convex, as it is an expectation of convex log-sum-exp functions. We exploit the convexity of $f(\theta)$ to define the conditions on the parameter space $\Theta$ for which the Perturb-Softmax probability distributions span the probability simplex.

The gradient $\nabla f(\theta)$ maps parameters $\theta$ to a probability, as the softmax vector $\mbox{softmax}(\theta+\gamma)$ sums up to unity for any $\gamma$ and therefore also in expectation over $\gamma$ (Corollary 7.2). In the next theorem, we determine the conditions for which the gradient mapping spans the relative interior of the probability simplex, i.e., the set of all possible positive probabilities.

###### Theorem 4.1 (Completeness of Perturb-Softmax) .

Let $\Theta\subseteq\mathbb{R}^{d}$ be a convex set and let $\gamma=(\gamma_{1},...,\gamma_{d})$ be a vector of random variables whose cumulative distribution decays to zero as $\gamma$ approaches $\pm\infty$. Let $h_{i}(\theta)=\theta_{i}-max_{j\neq i}\theta_{j}$ be a continuous function over $\Theta$. If $h_{i}(\theta)$ are unbounded then $\Theta$ is a complete representation of the Perturb-Softmax probability models:

$$ $ri(\Delta)\subseteq\mathbb{E}_{\gamma}[\mbox{softmax}(\theta+\gamma)]\subseteq\Delta$ (8) $$

###### Proof.

The proof relies on fundamental notions of the conjugate dual function $f^{*}(p)$ and its convex domain ${\cal P}$, cf. Equations (17, 18) in the Appendix.

We begin by considering $f(\theta)=\mathbb{E}_{\gamma}[\log(\sum_{i=1}^{d}e^{\theta_{i}+\gamma_{i}})]$ and its gradient, which is the Perturb-Softmax model $\nabla f(\theta)=\mathbb{E}_{\gamma}[\mbox{softmax}(\theta+\gamma)]$. Equation 21 implies that its gradients, i.e., the Perturb-Softmax probability distributions, reside in their convex domain ${\cal P}$ (cf. Equation 20), and contain its relative interior, thus:

$$ $ri({\cal P})\subseteq\left\{\mathbb{E}_{\gamma}[\mbox{softmax}(\theta+\gamma)] :\theta\in\Theta\right\}\subseteq{\cal P}$ (9) $$

To conclude the proof, we prove in Appendix 7.3 that the zero-one probability vectors reside in the closure of $\cal P$. Since the closure of $\cal P$ is a convex set, we conclude that it is the probability simplex.
∎

The above theorem implies that the set of all Perturb-Softmax probability distribution is an almost convex set that resides within the convex set of all probabilities $\Delta$ and contains its relative interior, i.e., the convex set of all positive probabilities $ri(\Delta)$.

Next, we describe the conditions for which the parameter space $\Theta$ is minimal. In this case, two different parameters $\theta\neq\tau$ result in two different Perturb-Softmax models $\mathbb{E}_{\gamma}[\mbox{softmax}(\theta+\gamma)]\neq\mathbb{E}_{\gamma}[
\mbox{softmax}(\tau+\gamma)]$. Interestingly, minimality is tightly tied to strict convexity. We begin by proving that $f(\theta)=\mathbb{E}_{\gamma}[\log(\sum_{i=1}^{d}e^{\theta_{i}+\gamma_{i}})]$ is strictly convex when restricted to $\Theta$.

###### Lemma 4.2 (Strict convexity) .

Let $\Theta\subseteq\mathbb{R}^{d}$ be a convex set and let $\gamma=(\gamma_{1},...,\gamma_{d})$ be a vector of random variables and let $f(\theta)=\mathbb{E}_{\gamma}[\log(\sum_{i=1}^{d}e^{\theta_{i}+\gamma_{i}})]$. If
$\Theta$ has no two vectors $\theta\neq\tau\in\Theta$ that are affine translations of each other, for which $\theta_{i}=\tau_{i}+c$ for every $i=1,...,d$ and some constant $c$ then $f(\theta)$ is strictly convex over $\Theta$, i.e., for any $\theta\neq\tau\in\Theta$ and any $0<\lambda<1$ it holds that

$$ $f(\lambda\theta+(1-\lambda)\tau)<\lambda f(\theta)+(1-\lambda)f(\tau).$ (10) $$

The proof, based on Hölder’s inequality, is provided in Appendix 7.4.
The condition that $\theta$ and $\tau$ are not a translation of each other guarantees strict convexity. If $\tau_{i}=\theta_{i}+c$ for every $i$, then $f(\tau)=f(\theta)+c$. This linear relation implies that the convexity condition holds with equality.

The minimality theorem is a direct consequence of Lemma 4.2, as strict convexity of differentiable function implies the gradient mapping is one-to-one (cf. Rockafellar (), Theorem 26.1).

###### Theorem 4.3 (Minimality of Perturb-Softmax) .

Let $\Theta\subseteq\mathbb{R}^{d}$ be a convex set and let $\gamma=(\gamma_{1},...,\gamma_{d})$ be a vector of random variables. $\Theta$ is a minimal representation of the Perturb-Softmax probability models if there are no two parameter vectors $\theta\neq\tau\in\Theta$ that are affine translations of each other, for which $\theta_{i}=\tau_{i}+c$ for every $i=1,...,d$ and some constant $c$.

###### Proof.

Lemma 4.2 implies that $f(\theta)$ is strictly convex and Equation (7) implies it is differentiable. Recall the conjugate dual function and its domain (Equations (17, 18) in the Appendix) and its gradient mapping $\nabla f:\Theta\rightarrow{\cal P}$. Since $f(\theta)$ is strictly convex then the function $g(\theta)=\langle p,\theta\rangle-f(\theta)$ is strictly concave hence its maximal argument $\theta^{*}$ is unique. The gradient vanishes at the maximal argument, $\nabla g(\theta^{*})=0$ or equivalently, $p=\nabla f(\theta^{*})$. Since $\theta^{*}$ is unique then $p$ is unique as well. Therefore $\nabla f(\theta)$ is a one-to-one mapping.
∎

To conclude, one can use any convex set $\Theta\subset\mathbb{R}^{d}$ that satisfies the conditions of Theorem 4.1 and Theorem 4.3. Similarly to the softmax probability model, $\Theta$ that is both complete is minimal can be $\Theta=\{\theta\in\mathbb{R}^{d}:\sum_{j}\theta_{j}=0\}$ or $\Theta=\{\theta\in\mathbb{R}^{d}:\theta_{1}=0\}$.

###### Theorem 4.1 (Completeness of Perturb-Softmax) .

###### Proof.

###### Lemma 4.2 (Strict convexity) .

###### Theorem 4.3 (Minimality of Perturb-Softmax) .

###### Proof.

## 5 Perturb-Argmax probability distributions

In this section, we explore the statistical representation properties of the Perturb-Argmax operation as a generalization of the Gumbel-Argmax operation. Throughout our investigation, we treat the Perturb-Argmax probability model as the sub-gradient of the expected Perturb-Max function, which we prove in Corollary 8.1 in the Appendix for completeness:

$$ $\displaystyle f(\theta)$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\gamma}\left[\max_{i}\{\theta_{i}+\gamma_{i}\}\right]$ (11) $\displaystyle\partial f(\theta)$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\gamma}[\arg\max(\theta+\gamma)]$ (12) $$

Different than the softmax operation, the argmax operation is not continuous everywhere. This difference arises from the fact that unlike the differentiable log-sum-exp function, the max function is not everywhere differentiable. However, since it is a convex function, its sub-gradient always exists.
In the following, we prove that the sub-gradients span the set of all positive probability distributions, i.e., the relative interior of the probability simplex.

###### Theorem 5.1 (Completeness of Perturb-Argmax) .

Let $\Theta\subseteq\mathbb{R}^{d}$ be a convex set and let $\gamma=(\gamma_{1},...,\gamma_{d})$ be a vector of random variables. Then, $\Theta$ is a complete representation of the Perturb-Argmax probability models:

$$ $ri(\Delta)\subseteq\mathbb{E}_{\gamma}[\arg\max(\theta+\gamma)]\subseteq\Delta$ (13) $$

###### Proof.

The proof technique follows the argument of Theorem 4.1. Given the conditions on $h_{i}(\theta)$, we can construct a series $\{\theta^{(n)}\}_{n=1}^{\infty}$ for which $h_{i}(\theta^{(n)})=n$ for every $n\in\mathbb{N}$. To conclude the proof, we prove in Appendix 8.1 that $\mathbb{E}_{\gamma}[\arg\max(\theta^{(n)}+\gamma)]$ approaches the zero-one probability vector as $n\rightarrow\infty$.
This proves that the zero-one distributions are limit points of probabilities in ${\cal P}$, i.e., $cl({\cal P})=\Delta$.
∎

The above theorem holds for any type of random perturbation $\gamma$. Next, we show that the statistical properties of the Perturb-Argmax probability models depend on their perturbation type.
The minimality of the representation of Perturb-Argmax probability models holds for non-discrete random perturbation $\gamma$. It relies on the differentiability properties of its probability density function $p(\gamma)$.

###### Lemma 5.2 (Differentiability of Perturb-Max) .

Let $\gamma=(\gamma_{1},...,\gamma_{d})$ be a vector of random variables with differentiable probability density function $p(\gamma)=\prod_{i=1}^{d}p_{i}(\gamma_{i})$ and let $f(\theta)=\mathbb{E}_{\gamma}[\max\{\theta+\gamma\}]$. Then, $f(\theta)$ is differentiable and its gradient is

$$ $\nabla f(\theta)=\mathbb{E}_{\gamma}[\arg\max(\theta+\gamma)]$ (14) $$

The proof is provided in Appendix 8.2.
Lemma 5.2 shows a single-valued mapping from the parameter space to the probability space. In the following, we show that this mapping brings forth a minimal representation of the Perturb-Argmax probability models under certain conditions.

###### Theorem 5.3 (Minimality of Perturb-Argmax) .

Let $\Theta\subseteq\mathbb{R}^{d}$ be a convex set and let $\gamma=(\gamma_{1},...,\gamma_{d})$ be a vector of random variables whose probability density functions $p_{i}(\gamma_{i})$ are differentiable and positive. $\Theta$ is a minimal representation of the Perturb-Argmax probability models if there are no two parameter vectors $\theta\neq\tau\in\Theta$ that are affine translations of each other, for which $\theta_{i}=\tau_{i}+c$ for every $i=1,...,d$ and some constant $c$.

The proof is provided in Appendix 8.3. It is based on showing that under these conditions, the function $f(\theta)=\mathbb{E}_{\gamma}[\max\{\theta+\gamma\}]$ is strictly convex. We rely on its one-dimensional function $g(\lambda)\triangleq f(\theta+\lambda v)$ and show that $g^{\prime\prime}(\lambda)>0$. Since the function $g(\lambda)$ is convex then $g^{\prime\prime}(\lambda)\geq 0$, and it is enough to show that $g^{\prime}(\lambda)$ depends on $\lambda$, for which it follows that $g^{\prime\prime}(\lambda)\neq 0$ and consequently $g^{\prime\prime}(\lambda)>0$.

Our theorem conditions require the probability density function to be positive, to ensure that the second derivative is positive as it always accounts for a change in the perturbation space.

###### Theorem 5.1 (Completeness of Perturb-Argmax) .

###### Proof.

###### Lemma 5.2 (Differentiability of Perturb-Max) .

###### Theorem 5.3 (Minimality of Perturb-Argmax) .

### 5.1 Non-minimal representation for bounded perturbations

In the following, we analyze an example of a Perturb-Argmax distribution when the probability density function of the perturbation is differentiable almost everywhere but bounded. In this case, one can construct a non-minimal representation.

###### Proposition 5.4 .

Let $\theta\in\mathbb{R}^{2}$, and consider i.i.d. random variables with a smooth bounded probability density function $\gamma\sim U(1,-1)$. A single-valued mapping exists between $f(\theta)$ and the Perturb-Argmax probability distribution. However, a one-to-one mapping does not exist.

###### Proof.

The perturb-max function $f(\theta)$ can be expressed as

$$ $\displaystyle f(\theta)=\theta_{2}+\mathbb{E}_{\gamma_{1},\gamma_{2}}\left[ \max\{\theta_{1}-\theta_{2}+\gamma_{1}-\gamma_{2},0\}\right],$ (15) $$

when the distribution of $\gamma$ is omitted for brevity.
We can express $f(\theta)$ by the pdf of the random variable $\gamma_{1}-\gamma_{2}$ (Equation 102 in the appendix). Since $f(\theta)$ is a smooth function, a single-valued mapping exists (Theorem 26.1 Rockafellar ()). However, $f(\theta)$ is not strictly convex, hence a one-to-one mapping does not exist and it can be concluded that $\Theta$ is not a minimal representation of the Perturb-Argmax probability.
The derivatives of $f(\theta)$, corresponding to the probabilities of the $\arg\max$, are illustrated in Figure 3.

Figure: Figure 2: An illustration of $\partial f(\theta)$ for perturbations with a smooth bounded probability density function $\gamma\sim U(-1,1)$. $\partial f(\theta)$ is a single-valued mapping between the parameters and the Perturb-Argmax probability.
Refer to caption: x3.png

We defer all details to Section 8.5 in the Appendix.
∎

###### Proposition 5.4 .

###### Proof.

### 5.2 Discrete perturbations and identifiablity

We next analyze the identifiability of the Perturb-Argmax distribution representation. Since the max function is not always differentiable, the Perturb-Max function $f(\theta)$ (Equation 11) is not always differentiable. However, since the max function is convex, its sub-differential $\partial f(\theta)$ exists. Unfortunately, the function $\partial f(\theta)$ is a multi-valued function, i.e., for some parameters $\theta\in\Theta$ there exist $p\neq q\in{\cal P}$ that both are its sub-gradient. Thus, the probability cannot be identified from the parameters when the perturbation is discrete. This property is demonstrated in the following proposition:

###### Proposition 5.5 .

Let $\Theta=\mathbb{R}^{2}$ and $\gamma=(\gamma_{1},\gamma_{2})$ be a vector of discrete random variables $\gamma_{i}\in\{1,-1\}$ that are uniformly distributed: $\mathbb{P}[\gamma_{i}=1]=\mathbb{P}[\gamma_{i}=-1]=\frac{1}{2}$. Then, the Perturb-Argmax probability distribution is not identifiable.

The proof is provided in Appendix 8.4, and it is based on computing the function $f(\theta)=\mathbb{E}_{\gamma}[\max\{\theta+\gamma\}]$ analytically by taking the expectation w.r.t. $\gamma$.
The function $f(\theta)$, illustrated in Figure 6 in the appendix, is continuous and differentiable almost everywhere.
However, in its overlapping segments, the function is not differentiable, i.e., it has a sub-differential $\partial f(\theta)$ which is a set of sub-gradients. To prove that the Perturb-Argmax probability model is unidentifiable, we show that $\partial f(\theta)$ is a multi-valued mapping when $\theta_{1}=\theta_{2}$. The sub-differential mapping is illustrated in Figure 3.

###### Proposition 5.5 .

## 6 Experiments

In this section, we demonstrate the advantage of the Gaussian-Softmax over the commonly used Gumbel-Softmax. Experiments in density estimation and variational inference exhibit that, compared to the Gumbel-Softmax, the Gaussian-Softmax enjoys a faster convergence rate and better approximate discrete distributions.

### 6.1 Approximating discrete distributions

We compare the Gumbel-Softmax and the Normal-Softmax in approximating discrete distributions. The $L_{1}$ objective function is minimized between the Perturb-Softmax function applied to the fitted parameters of the probability density function and the target discrete distribution, denoted by $p_{0}$. We consider two target discrete distributions with finite support: a binomial distribution with parameters $n=12,p=0.3$, and a discrete distribution with $p=(\frac{10}{68},\frac{3}{68},\frac{4}{68},\frac{5}{68},\frac{10}{68},\frac{10
}{68},\frac{3}{68},\frac{4}{68},\frac{5}{68},\frac{10}{68})$. Figure 4 shows that the Normal-Softmax better approximates both distributions and exhibits faster convergence than the Gaussian-Softmax.

Following the experiment in Potapczynski et al. (), we consider discrete distributions with countably
infinite support: a Poisson distribution with $\lambda=50$, and a negative binomial distribution with $r=50,p=0.6$.
Differently from our method, the identifiability of the parameters is lost with the invertible Gaussian parameterization method. Results show that the Normal-Softmax has similar benefits over the Gaussian-Softmax as exhibited for discrete distributions with finite support (Figure 9 in the appendix).
Table 1 in the Appendix shows that the approximation based on the Normal-Softmax probability model consistently achieves lower mean and standard deviation $L_{1}$.
See more details in Appendix 9.1.

Figure: (a) Approximation with dimension d=12 of a binomial distribution.
Refer to caption: x5.png

### 6.2 Variational inference

We compared the training ELBO-based loss of categorical Variational-Autoencoders for $N=10$ variables, each is a $K$-dimensional categorical variable, $K\in[10,30,50]$ on the binarized MNIST LeCun & Cortes (), the Fashion-MNIST (Xiao et al., ), and the Omniglot (Lake et al., ) datasets for different smooth perturbation distributions. The architecture consists of an encoder of $X\rightarrow FC(300)\rightarrow\mathrm{ReLU}\rightarrow N*K$, and matching decoder $N*K\rightarrow FC(300)\rightarrow\mathrm{ReLU}\rightarrow X$. We compare the training convergence when propagating gradients with the Normal-Softmax or the Gumbel-Softmax, both performed with a temperature equal to $1$.
Results show that the former achieves better and faster learning convergence in all experiments (Figures 5 and 10).

Figure: (a) K=10
Refer to caption: x11.png

## 7 APPENDIX

### 7.1 Related work

#### 7.1.1 Convexity

We consider convex functions $f:\Theta\rightarrow\mathbb{R}$ over its convex domain $\Theta=dom(f)$ and follow the notation in Rockafellar (). A function over a domain $\Theta$ is convex if for every $\theta,\tau\in\Theta$ and $0\leq\lambda\leq 1$ it holds that $f(\lambda\theta+(1-\lambda)\tau)\leq\lambda f(\theta)+(1-\lambda)f(\tau)$. A function is strictly convex if for every $\theta\neq\tau\in\Theta$ and $0<\lambda<1$, it holds that $f(\lambda\theta+(1-\lambda)\tau)<\lambda f(\theta)+(1-\lambda)f(\tau)$. Whenever $f(\theta)$ is twice differentiable, a function is strictly convex if its Hessian is positive definite.

Convexity is a one-dimensional property. A function $f:\Theta\rightarrow\mathbb{R}$ is convex if and only if its one-dimensional reduction $g(\lambda)\triangleq f(\theta+\lambda v)$ is convex in every admissible direction $v$, i.e., whenever $\theta,\theta+\lambda v\in\Theta$. A twice differentiable function $f(\theta)$ is strictly convex if the second derivative of $g(\lambda)$ is positive in every admissible direction $v$. In this case, we denote $g^{\prime}(\lambda)$ by $\nabla_{v}f(\theta)$ and call it a directional derivative:

$$ $\nabla_{v}f(\theta)\triangleq\lim_{\epsilon\rightarrow 0}\frac{f(\theta+ \epsilon v)-f(\theta)}{\epsilon}$ (16) $$

A multivariate function is differentiable if its directional derivative is the same in every direction $v\in\mathbb{R}^{d}$, namely $\nabla f(\theta)=\nabla_{v}f(\theta)$ for every $v\in\mathbb{R}^{d}$.

Convexity admits duality correspondence. Any primal convex function $f(\theta)$ has a dual conjugate function $f^{*}(p)$

$$ $\displaystyle f^{*}(p)$ $\displaystyle=$ $\displaystyle\max_{\theta\in\Theta}\left\{\langle p,\theta\rangle-f(\theta)\right\}$ (17) $\displaystyle{\cal P}$ $\displaystyle\triangleq$ $\displaystyle dom(f^{*})\triangleq\{p:f^{*}(p)<\infty\}$ (18) $$

Since $f^{*}(p)$ is a convex function, its domain ${\cal P}$ is a convex set.

A sub-gradient $p\in\partial f(\theta)$ satisfies $f(\tau)\geq f(\theta)+\langle p,\tau-\theta\rangle$ for every $\tau\in\Theta$. The sub-gradient is intimately connected to directional derivatives. Theorem 23.2 in Rockafellar () states that

$$ $p\in\partial f(\theta)\hskip 14.22636pt\mbox{ iff }\hskip 14.22636pt\nabla_{v} f(\theta)\geq\langle p,v\rangle,\;\forall\mbox{ admissible }v.$ (19) $$

The vector $v$ is admissible if $\theta+\epsilon v\in\Theta$ for small enough $\epsilon$.

The set of all sub-gradients is called sub-differential at $\theta$ and is denoted by $\partial f(\theta)$. A convex function is differentiable at $\theta$ when $\partial f(\theta)$ consists of a single vector, and it is denoted by $\nabla f(\theta)$.
The sub-differential is a multi-valued mapping between the primal parameters and dual parameters, i.e.,

$$ $\partial f:\Theta\rightarrow{\cal P}$ (20) $$

One can establish with this property the definition of sub-gradient at the optimal point $\theta^{*}=\arg\max_{\theta}\{\langle p,\theta\rangle-f(\theta)\}$. In this case, $0\in\partial\left(\langle p,\theta^{*}\rangle-f(\theta^{*})\right)$, where the sub-gradient is taken with respect to $\theta$ at the maximal argument $\theta^{*}$. From the linearity of the sub-gradient, there holds: $p\in\partial f(\theta^{*})$, or equivalently, $\{\partial f(\theta):\theta\in\Theta\}\subseteq{\cal P}$. Using the connection between sub-gradients and directional derivatives, one can show that whenever directional derivatives exist, one can infer a sub-gradient, i.e., the set of all sub-gradients contains the relative interior of $\mathbb{P}$, cf. Theorem 23.4 Rockafellar ():

$$ $ri({\cal P})\subseteq\{\partial f(\theta):\theta\in\Theta\}\subseteq{\cal P}$ (21) $$

### 7.2 Perturb-Softmax probability distributions

#### 7.2.1 Completeness and minimality of the Softmax operation

###### Theorem 7.1 .

The representation of the softmax distribution defined over $\theta\in\mathbb{R}^{d}$ is complete. It is minimal when the corresponding log-sum-exp

$$ $f(\theta)=\log(\sum_{i=1}^{d}e^{\theta_{i}})$ (22) $$

is a strictly convex function (A paraphrase of Wainwright & Jordan (), Proposition 3.1).

###### Proof.

First, we note that the derivatives of $f(\theta)$ (Eq. 22),

$$ $\frac{\partial f(\theta)}{\partial\theta_{j}}=\mbox{softmax}(\theta),$ (23) $$

correspond to the softmax probabilities $p^{sm}_{\theta}$ (Eq. 1). $p^{sm}_{\theta}$ is a Gibbs model, hence the representation is complete.

Let $\theta,\tau\in\mathbb{R}^{d}$ and denote $e^{\theta_{i}}=u_{i}$, $e^{\tau_{i}}=v_{i}$. Then, for $\lambda\in(0,1)$

$$ $\displaystyle f(\lambda\theta+(1-\lambda)\tau)$ $\displaystyle=$ $\displaystyle\log(\sum_{i=1}^{d}e^{\lambda\theta_{i}+(1-\lambda)\tau_{i}})$ (24) $\displaystyle=$ $\displaystyle\log(\sum_{i=1}^{d}u_{i}^{\lambda}v_{i}^{1-\lambda}).$ (25) $$

Applying Hölder’s inequality to Equation 25 :

$$ $\displaystyle\log\left(\sum_{i=1}^{d}u_{i}^{\lambda}v_{i}^{1-\lambda}\right)$ $\displaystyle\leq\log(\sum_{i=1}^{d}u_{i}^{\lambda\frac{1}{\lambda}})^{\lambda }(\sum_{i=1}^{d}v_{i}^{(1-\lambda)\frac{1}{1-\lambda}})^{1-\lambda})$ (26) $\displaystyle=\lambda\log(\sum_{i=1}^{d}u_{i})+(1-\lambda)\log(\sum_{i=1}^{d}v _{i})$ (27) $\displaystyle=\lambda f(\theta)+(1-\lambda)f(\tau).$ (28) $$

Therefore, it holds that

$$ $f(\lambda\theta+(1-\lambda)\tau)\leq\lambda f(\theta)+(1-\lambda)f(\tau),$ (29) $$

proving that $f(\theta)$ is convex.

Hölder’s inequality holds with equality if and only if there exists a constant $c\in\mathbb{R}$ such that

$$ $\displaystyle|v_{i}^{1-\lambda}|$ $\displaystyle=$ $\displaystyle c|u_{i}^{\lambda}|^{\frac{1}{\lambda}-1}\quad\forall i$ (30) $\displaystyle e^{(1-\lambda)\tau_{i}}$ $\displaystyle=$ $\displaystyle ce^{\lambda\theta_{i}(\frac{1}{\lambda}-1)}\quad\forall i$ (31) $\displaystyle\tau_{i}$ $\displaystyle=$ $\displaystyle\frac{\log(c)}{1-\lambda}+\theta_{i}\quad\forall i,$ (32) $$

in which case $\theta$ and $\tau$ are linearly constrained and there exists some $\alpha\in\mathbb{R}^{2},\alpha\neq 0$ such that $\alpha_{1}\theta+\alpha_{2}\tau=const$. Therefore, when the representation of $f(\theta)$ is minimal $f(\theta)$ is strictly convex.

Then, consider any $\theta,\tau\in\Theta^{d}$, such that $\Theta=\{\theta\in\mathbb{R}^{d}:\theta_{1}=0\}$. Hölder’s inequality holds strictly as there can not exist a constant $c$ such that Equation 32 holds for all $i$ if $d>1$, proving that the representation of $\Theta=\{\theta\in\mathbb{R}^{d}:\theta_{1}=0\}$ is minimal.

Consider $\Theta=\{\theta\in\mathbb{R}^{d}:\sum_{j}\theta_{j}=0\}$. Then, let $\theta,\tau\in\mathbb{R}^{d}:\sum_{i}\theta_{i}=0,\sum_{i}\tau_{i}=0$, and denote $p_{i}\propto e^{\theta_{i}}$ and $q_{i}\propto e^{\tau_{i}}$. The proof requires showing that if there exists $i:p_{i}\neq q_{i}$ such that $\theta_{i}\neq\tau_{i}+c$.
Equivalently, it requires proving that if it holds that $\theta_{i}-\tau_{i}=0$ for any $i$, then $p_{i}=q_{i}$ for all $i$.
Explicitly,

$$ $\displaystyle p_{i}$ $\displaystyle=$ $\displaystyle\frac{e^{\theta}_{i}}{\sum_{j}e^{\theta}_{j}}$ (33) $\displaystyle q_{i}$ $\displaystyle=$ $\displaystyle\frac{e^{\tau}_{i}}{\sum_{j}e^{\tau}_{j}}.$ (34) $$

Then, by marginalization it holds that

$$ $\displaystyle\log(p_{i})=\log(qi)$ $\displaystyle\longleftrightarrow$ $\displaystyle\theta_{i}-\log(\sum_{j}e^{\theta_{j}})=\tau_{i}-\log(\sum_{j}e^{ \tau_{j}})$ (35) $\displaystyle\longleftrightarrow$ $\displaystyle\sum_{i}\theta_{i}-\sum_{i}\log(\sum_{j}e^{\theta_{j}})=\sum_{i} \tau_{i}-\sum_{i}\log(\sum_{j}e^{\tau_{j}})$ (36) $\displaystyle\longleftrightarrow$ $\displaystyle d\log(\sum_{j}e^{\theta_{j}})=d\log(\sum_{j}e^{\tau_{j}}),$ (37) $$

which concludes the proof.

$\Theta=\{\theta\in\mathbb{R}^{d}:\sum_{j}\theta_{j}=0\}$ is complete by the conditions of our completeness theorems by setting $n$ at the $i^{th}$ positions and $-a/d$ everywhere else.

$\Theta=\{\theta\in\mathbb{R}^{d}:\theta_{1}=0\}$ is complete by the conditions of our completeness theorems by setting $n$ at the $1^{st}$ positions and 0 0 everywhere else.
∎

###### Corollary 7.2 .

The derivative of the expected log-sum-exp $f(\theta)$ (Equation 6) is a probability function.

###### Proof.

Denote $p_{\gamma,j}=\frac{\partial f(\theta)}{\partial\theta_{j}}$, then

$$ $\displaystyle\sum_{j=1}^{d}p_{\gamma,j}$ $\displaystyle=$ $\displaystyle\sum_{j=1}^{d}\mathbb{E}_{\gamma}\left[\frac{e^{\theta_{j}+\gamma _{j}}}{\sum_{i=1}^{d}e^{\theta_{i}+\gamma_{i}}}\right]$ (38) $\displaystyle=$ $\displaystyle\mathbb{E}_{\gamma}\left[\sum_{j=1}^{d}\frac{e^{\theta_{j}+\gamma _{j}}}{\sum_{i=1}^{d}e^{\theta_{i}+\gamma_{i}}}\right]=1.$ (39) $$

Also,

$$ $\mathbb{E}_{\gamma}\left[\frac{e^{\theta_{j}+\gamma_{j}}}{\sum_{i=1}^{d}e^{ \theta_{i}+\gamma_{i}}}\right]\geq 0\quad\forall j.$ (40) $$

∎

###### Theorem 7.1 .

###### Proof.

###### Corollary 7.2 .

###### Proof.

### 7.3 Supporting proof for Theorem 4.1

First, we prove that $f(\theta)$ (Equation 6) is a closed proper convex function and is also essentially smooth.
$f(\theta)$ is a convex function as a maximum of convex (linear) functions. Then,
$f(\theta)$ is proper as its effective domain is nonempty and it never attains the value $-\infty$, since $\theta\in\mathbb{R}^{d}$.
$f(\theta)$ is infinitely differentiable throughout the domain, therefore it is a smooth function throughout its domain.
$f(\theta)$ is a smooth convex function on $\mathbb{R}^{d}$, therefore it is in particular essentially smooth.
The smoothness of $f(\theta)$ guarantees its continuity, and since $\mathbb{R}^{d}$ can be considered a closed set, then $f(\theta)$ is a closed function.

Given the conditions of the theorem on $h_{i}(\theta)$, we can construct a series $\{\theta^{(n)}\}_{n=1}^{\infty}$ for which $h_{i}(\theta^{(n)})=n$ for every $n\in\mathbb{N}$. We prove that $\mathbb{E}_{\gamma}[\arg\max(\theta^{(n)}+\gamma)]$ approaches the zero-one probability vector as $n\rightarrow\infty$.

$$ $\displaystyle\mathbb{E}_{\gamma}\left[\frac{e^{\theta^{(n)}_{i}+\gamma_{i}}}{ \sum_{j=1}^{d}e^{\theta^{(n)}_{j}+\gamma_{j}}}\right]$ $\displaystyle=\mathbb{E}_{\gamma}\left[\frac{1}{\sum_{j=1}^{d}e^{\theta^{(n)}_ {j}+\gamma_{j}-\theta^{(n)}_{i}-\gamma_{i}}}\right]$ (43) $\displaystyle=\mathbb{E}_{\gamma}\left[\frac{1}{1+\sum_{j\neq i}e^{\theta^{(n) }_{j}-\theta^{(n)}_{i}+\gamma_{j}-\gamma_{i}}}\right]$ $\displaystyle\geq\mathbb{E}_{\gamma}\left[\frac{1}{1+\sum_{j\neq i}e^{-n+ \gamma_{j}-\gamma_{i}}}\right]\stackrel{{\scriptstyle n\rightarrow\infty}}{{ \rightarrow 1}}$ $$

The limit argument holds since the probability of $\gamma_{1},...,\gamma_{d}$ decay as they tend to infinity. This proves that the zero-one distributions are limit points of probabilities in ${\cal P}$, i.e., $cl({\cal P})=\Delta$.

### 7.4 Proof of Lemma 4.2

In the following, we prove the strict convexity of the expected log-sum-exp of Lemma 4.2.

###### Proof.

Let $\theta,\tau\in\mathbb{R}^{d}$ and $0<\lambda<1$. Then

$$ $\displaystyle f(\lambda\theta+(1-\lambda)\tau)=\mathbb{E}_{\gamma}\left[\log \left(\sum_{i=1}^{d}e^{\lambda(\theta_{i}+\gamma_{i})+(1-\lambda)(\tau_{i}+ \gamma_{i})}\right)\right]=\mathbb{E}_{\gamma}\left[\log\left(\sum_{i=1}^{d}u_ {i}v_{i}\right)\right],$ (44) $$

where $u_{i}\triangleq e^{\lambda(\theta_{i}+\gamma_{i})}$ and $v_{i}\triangleq e^{(1-\lambda)(\tau_{i}+\gamma_{i})}$. Applying Hölder’s inequality $\langle u,v\rangle\leq\|v\|_{1/\lambda}\cdot\|u\|_{1/(1-\lambda)}$ we obtain the convexity condition of the log-sum-exp function:

$$ $f(\lambda\theta+(1-\lambda)\tau)\leq\lambda f(\theta)+(1-\lambda)f(\tau).$ (45) $$

To prove strict convexity we note that Hölder’s inequality for non-negative vectors $u_{i},v_{i}$ holds with equality if and only if there exists a constant $\alpha\in\mathbb{R}$ such that $v_{i}=\alpha u_{i}^{\frac{1-\lambda}{\lambda}}$ for every $i=1,...,d$, or equivalently:

$$ $e^{(1-\lambda)(\tau_{i}+\gamma_{i})}=\left(e^{\lambda(\theta_{i}+\gamma_{i})} \right)^{\frac{1-\lambda}{\lambda}}\longleftrightarrow\tau_{i}=\theta_{i}+c$ (46) $$

Where $c=\frac{\log\alpha}{1-\lambda}$. Therefore, if $\tau\neq\theta+c$, then $\langle u,v\rangle<\|v\|_{1/\lambda}\cdot\|u\|_{1/(1-\lambda)}$ for every $\gamma$ and consequently it also holds when applying the logarithm function and taking an expectation with respect to $\gamma$.
∎

###### Proof.

## 8 Perturb-Argmax probability distributions

###### Corollary 8.1 .

We prove that the derivative of the expected maximizer is the probability of the $\arg\max$. Namely, that $\frac{\partial}{\partial\theta}\mathbb{E}_{\gamma\sim g}\left[\max_{i}\{\theta
_{i}+\gamma_{i}\}\right]=P_{\gamma}\left(\arg\max_{i}\{\theta_{i}+\gamma_{i}\}
=i\right)$.

###### Proof.

First, by differentiating under the integral:

$$ $\partial\mathbb{E}_{\gamma}[\max\{\theta+\gamma\}]=\mathbb{E}_{\gamma}[ \partial\max\{\theta+\gamma\}]$ (47) $$

Writing a subgradient of the max-function using an indicator function (an application of Danskin’s Theorem):

$$ $\partial\max_{i}\{\theta_{i}+\gamma_{i}\}=\mathds{1}[\arg\max_{i}(\theta_{i}+ \gamma_{i})=i]$ (48) $$

The proof then follows by applying the expectation to both
sides of Equation 48.
∎

###### Corollary 8.1 .

###### Proof.

### 8.1 Supporting proof for Theorem 5.1

Given the conditions on $h_{i}(\theta)$, we can construct a series $\{\theta^{(n)}\}_{n=1}^{\infty}$ for which $h_{i}(\theta^{(n)})=n$ for every $n\in\mathbb{N}$. We show that $\mathbb{E}_{\gamma}[\arg\max(\theta^{(n)}+\gamma)]$ approaches the zero-one probability vector as $n\rightarrow\infty$.

$$ $\displaystyle\mathbb{P}[i=\arg\max(\theta^{(n)}+\gamma)$ $\displaystyle=\mathbb{P}\left[\theta^{(n)}_{i}+\gamma_{i}\geq\max_{j\neq i}\{ \theta^{(n)}_{j}+\gamma_{j}\}\right]$ $\displaystyle\geq\mathbb{P}\left[\theta^{(n)}_{i}+\gamma_{i}\geq\max_{j\neq i} \{\theta^{(n)}_{j}\}+\max_{j\neq i}\{\gamma_{j}\}\right]$ $\displaystyle\geq\mathbb{P}\left[\gamma_{i}\geq-n+\max_{j\neq i}\{\gamma_{j}\} \right]\stackrel{{\scriptstyle n\rightarrow\infty}}{{\rightarrow 1}}$ $$

The limit argument holds since the probability of $\gamma_{1},...,\gamma_{d}$ decay as they tend to infinity.

###### Corollary 8.2 .

The convex conjugate of $f(\theta)$ takes the following values:

$$ $f^{*}\left(\lambda\right)=\begin{cases}-\mathbb{E}_{\gamma\sim g}\left[\gamma_ {\hat{i}}\right]&\quad\text{if }\lambda=\lambda^{*}\\ \infty&\quad\text{otherwise},\end{cases}$ (51) $$

where $\gamma_{\hat{i}}$ denotes $\gamma_{i}$ for which $\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=\hat{i}$.

###### Proof.

Then, the convex conjugate of $f(\theta)$, when $\lambda^{*}_{i}$ is denoted by $p_{i}$ is

$$ $\displaystyle f^{*}\left(\lambda\right)$ $\displaystyle=$ $\displaystyle\sum_{i}\theta_{i}p_{i}-\mathbb{E}_{\gamma\sim g}\left[\max_{i}\{ \theta_{i}+\gamma_{i}\}\right]$ (52) $\displaystyle=$ $\displaystyle\sum_{i}\theta_{i}p_{i}-\mathbb{E}_{\gamma}\left[\left(\sum_{\hat {i}}\mathds{1}_{\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=\hat{i}}\right)\max_{i} \{\theta_{i}+\gamma_{i}\}\right]$ (53) $\displaystyle=$ $\displaystyle\sum_{i}\theta_{i}p_{i}-\sum_{\hat{i}}\mathbb{E}_{\gamma}\left[ \mathds{1}_{\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=\hat{i}}\max_{i}\{\theta_{i} +\gamma_{i}\}\right]$ (54) $\displaystyle=$ $\displaystyle\sum_{i}\theta_{i}p_{i}-\sum_{\hat{i}}\mathbb{E}_{\gamma}\left[ \mathds{1}_{\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=\hat{i}}(\theta_{\hat{i}}+ \gamma_{\hat{i}})\right]$ (55) $\displaystyle=$ $\displaystyle\sum_{i}\theta_{i}p_{i}-\sum_{\hat{i}}\mathbb{E}_{\gamma}\left[ \mathds{1}_{\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=\hat{i}}\theta_{\hat{i}} \right]-\sum_{\hat{i}}\mathbb{E}_{\gamma}\left[\mathds{1}_{\arg\max_{i}\{ \theta_{i}+\gamma_{i}\}=\hat{i}}\gamma_{\hat{i}}\right]$ (56) $\displaystyle=$ $\displaystyle\sum_{i}\theta_{i}p_{i}-\sum_{\hat{i}}p_{\hat{i}}\theta_{\hat{i}} -\sum_{\hat{i}}\mathbb{E}_{\gamma}\left[\mathds{1}_{\arg\max_{i}\{\theta_{i}+ \gamma_{i}\}=\hat{i}}\gamma_{\hat{i}}\right]$ (57) $\displaystyle=$ $\displaystyle-\sum_{\hat{i}}\mathbb{E}_{\gamma}\left[\mathds{1}_{\arg\max_{i} \{\theta_{i}+\gamma_{i}\}=\hat{i}}\gamma_{\hat{i}}\right]$ (58) $\displaystyle=$ $\displaystyle-\sum_{\hat{i}}p_{\hat{i}}\mathbb{E}_{\gamma}\left[\gamma_{\hat{i }}\right]$ (59) $\displaystyle=$ $\displaystyle-\mathbb{E}_{\gamma}\left[\gamma_{\hat{i}}\right]$ (60) $$

where $\gamma_{\hat{i}}$ denotes $\gamma_{i}$ for which $\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=\hat{i}$.

∎

###### Corollary 8.2 .

###### Proof.

### 8.2 Proof of Lemma 5.2

###### Proof.

By reparameterization

$$ $\displaystyle f(\theta)$ $\displaystyle=$ $\displaystyle\int_{\mathbb{R}^{d}}p(\gamma)\max\{\theta+\gamma\}d\gamma=\int_{ \mathbb{R}^{d}}p(\theta-\gamma)\max\{\gamma\}d\gamma$ (61) $$

The proof concludes by differentiating under the integral sign while noting that $p(\theta-\gamma)$ is differentiable.
∎

###### Proof.

### 8.3 Proof of Theorem 5.3

In what follows we prove the minimality of the Perturb-Argmax of Theorem 5.3.

###### Proof.

Similarly to the Perturb-Softmax setting, we prove that under these conditions the function $f(\theta)=\mathbb{E}_{\gamma}[\max\{\theta+\gamma\}]$ is strictly convex. For this we rely on its one dimensional function $g(\lambda)\triangleq f(\theta+\lambda v)$ and show that $g^{\prime\prime}(\lambda)>0$. Since the function $g(\lambda)$ is convex then $g^{\prime\prime}(\lambda)\geq 0$, and it is enough to show that $g^{\prime}(\lambda)$ depends on $\lambda$, for which it follows that $g^{\prime\prime}(\lambda)\neq 0$ and consequently $g^{\prime\prime}(\lambda)>0$.

$g^{\prime}(\lambda)$ is the directional derivative $\nabla_{v}f(\theta)$ in every admissible direction $v=\tau-\theta$, for $\tau,\theta\in\Theta$. The theorem conditions assert that $v$ is not the constant vector, i.e., $v\neq c1$, where $c$ is some constant and $1=(1,...,1)$ is the all-one vector.

We assume, without loss of generality, that $\max\{\theta+\gamma\}\triangleq\max_{i=1,...,d}\{\theta_{i}+\gamma_{i}\}$ is chosen between two indexes, namely $\max\{\theta+\gamma\}=\max\{\theta_{1}+\gamma_{1},\theta_{j}+\gamma_{j}\}$. This is possible as we treat $j$ to be $j=\arg\max_{i\neq 1}\{\theta_{j}+\gamma_{j}\}$. We denote by $p_{1}(\gamma_{1})$ the differentiable probability density function of $\gamma_{1}$. We denote remaining random variables as
$\gamma_{-1}\triangleq(\gamma_{2},...,\gamma_{d})$, their probability density function by their measure $d\mu(\gamma_{-1})$.

We analyze $g^{\prime}(\lambda)=\lim_{\epsilon\rightarrow 0}\frac{1}{\epsilon}(g(\lambda+
\epsilon)-g(\lambda))$

$$ $\displaystyle g(\lambda)$ $\displaystyle=$ $\displaystyle E_{\gamma}[\max\{\theta_{1}+\lambda v_{1}+\gamma_{1},\theta_{j}+ \lambda v_{j}+\gamma_{j}\}]$ $\displaystyle=$ $\displaystyle\int d\mu(\gamma_{-1})\int_{-\infty}^{\alpha(\gamma_{j})}p_{1}( \gamma_{1})(\theta_{j}+\lambda v_{j}+\gamma_{j})d\gamma_{1}+\int d\mu(\gamma_{ -1})\int_{\alpha(\gamma_{j})}^{\infty}p_{1}(\gamma_{1})(\theta_{1}+\lambda v_{ 1}+\gamma_{1})d\gamma_{1}$ $$

$\alpha(\gamma_{j})$ is the threshold for which $\gamma_{1}$ shifts the maximal value to $\theta_{1}+\lambda v_{1}+\gamma_{1}$, namely $\alpha(\gamma_{j})=\theta_{j}+\lambda v_{j}+\gamma_{j}-\theta_{1}-\lambda v_{1}$.

With this notation, $g(\lambda+\epsilon)=$

$$ $\int d\mu(\gamma_{-1})\int_{-\infty}^{\alpha(\gamma_{j})+\epsilon(v_{j}-v_{1}) }p_{1}(\gamma_{1})(\theta_{j}+(\lambda+\epsilon)v_{j}+\gamma_{j})d\gamma_{1}+ \int d\mu(\gamma_{-1})\int_{\alpha(\gamma_{j})+\epsilon(v_{j}-v_{1})}^{\infty} p_{1}(\gamma_{1})(\theta_{1}+(\lambda+\epsilon)v_{1}+\gamma_{1})d\gamma_{1}$ (63) $$

Their difference is composed of four terms:

$$ $\displaystyle g(\lambda+\epsilon)-g(\lambda)=$ $\displaystyle\int d\mu(\gamma_{-1})\int_{-\infty}^{\alpha(\gamma_{j})+\epsilon (v_{j}-v_{1})}p_{1}(\gamma_{1})\epsilon v_{j}d\gamma_{1}+\int d\mu(\gamma_{-1} )\int_{\alpha(\gamma_{j})+\epsilon(v_{j}-v_{1})}^{\infty}p_{1}(\gamma_{1}) \epsilon v_{1}d\gamma_{1}$ $\displaystyle+\int d\mu(\gamma_{-1})\int_{\alpha(\gamma_{j})}^{\alpha(\gamma_{ j})+\epsilon(v_{j}-v_{1})}p_{1}(\gamma_{1})(\theta_{j}+(\lambda+\epsilon)v_{j} +\gamma_{j})d\gamma_{1}$ $\displaystyle-\int d\mu(\gamma_{-1})\int_{\alpha(\gamma_{j})}^{\alpha(\gamma_{ j})+\epsilon(v_{j}-v_{1})}p_{1}(\gamma_{1})(\theta_{1}+(\lambda+\epsilon)v_{1} +\gamma_{1})d\gamma_{1}$ $$

Taking the limit to zero $\lim_{\epsilon\rightarrow 0}\frac{1}{\epsilon}(g(\lambda+\epsilon)-g(\lambda)$, the last two terms cancel out, since when taking the limit then $\gamma_{1}=\alpha(\gamma_{j})$ and $\alpha(\gamma_{j})=\theta_{j}+\lambda v_{j}+\gamma_{j}-\theta_{1}-\lambda v_{1}$ by definition, or equivalently, $\alpha(\gamma_{j})+\theta_{1}+\lambda v_{1}=\theta_{j}+\lambda v_{j}+\gamma_{j}$. Therefore

$$ $\displaystyle g^{\prime}(\lambda)=\int d\mu(\gamma_{-1})\int_{-\infty}^{\alpha (\gamma_{j})}p_{1}(\gamma_{1})v_{j}d\gamma_{1}+\int d\mu(\gamma_{-1})\int_{ \alpha(\gamma_{j})}^{\infty}p_{1}(\gamma_{1})v_{1}d\gamma_{1}$ (64) $$

We conclude that by the conditions of the theorem $\alpha(\gamma_{j})=\theta_{j}+\lambda v_{j}+\gamma_{j}-\theta_{1}-\lambda v_{1}$ is a function of $\lambda$ since there exists $j$ for which $v_{j}-v_{1}\neq 0$ and the probability density function $p_{1}(\gamma_{1})>0$ therefore it assigns mass on the intervals $[-\infty,\alpha(\gamma_{j})]$ and $[\alpha(\gamma_{j}),\infty]$. Therefore $g^{\prime}(\lambda)$ is non-constant function of $\lambda$ and $g^{\prime\prime}(\lambda)\neq 0$.
∎

###### Proof.

### 8.4 Proof of Proposition 5.5

Let $\theta\in\mathbb{R}^{2}$, and consider i.i.d. random variables $\gamma\in\{1,-1\}$, such that $P(\gamma=1)=P(\gamma=-1)=\frac{1}{2}$. Let $f(\theta)=\mathbb{E}_{\gamma}\left[\max_{i}\{\theta_{i}+\gamma_{i}\}\right]$ denote the expected perturbed maximum over the domain $\mathbb{R}^{d}$. Let $f(\theta)$ take values over the extended real domain $\mathbb{R}\cup\{\pm\infty\}$. Clearly, $P(1=\arg\max_{i}\{\theta_{i}+\gamma_{i}\})=P(\theta_{1}+\gamma_{1}\geq\theta_{
2}+\gamma_{2})$.

Then, $f(\theta)$ can be explicitly expressed as:

$$ $\displaystyle f(\theta)$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\gamma}\left[\max_{i}\{\theta_{i}+\gamma_{i}\}\right]$ $\displaystyle=$ $\displaystyle P(\gamma_{1}=1,\gamma_{2}=1)\left(1+\max\{\theta_{1},\theta_{2} \}\right)+P(\gamma_{1}=-1,\gamma_{2}=-1)\left(\max\{\theta_{1},\theta_{2}\}-1\right)$ $\displaystyle+$ $\displaystyle P(\gamma_{1}=1,\gamma_{2}=-1)\left(\max\{\theta_{1}+1,\theta_{2} -1\}\right)+P(\gamma_{1}=-1,\gamma_{2}=1)\left(\max\{\theta_{1}-1,\theta_{2}+1 \}\right)$ $\displaystyle=$ $\displaystyle\frac{1}{2}\left(\max\{\theta_{1},\theta_{2}\}\right)+P(\gamma_{1 }=1,\gamma_{2}=-1)\left(\max\{\theta_{1}+1,\theta_{2}-1\}\right)+P(\gamma_{1}= -1,\gamma_{2}=1)\left(\max\{\theta_{1}-1,\theta_{2}+1\}\right)$ $\displaystyle=$ $\displaystyle\frac{1}{2}\left(\max\{\theta_{1},\theta_{2}\}\right)+\frac{1}{2} \left(\frac{\max\{\theta_{1}+1,\theta_{2}-1\}}{2}+\frac{\max\{\theta_{1}-1, \theta_{2}+1\}}{2}\right)$ $$

Equation 8.4 suggests that $f(\theta)$ takes the following form:

$$ $f(\theta)=\begin{cases}\theta_{1}&\quad\text{if }\theta_{1}\geq\theta_{2}+2\\ \frac{3}{4}\theta_{1}+\frac{1}{4}\theta_{2}+\frac{1}{2}&\quad\text{if }\theta_ {2}+2\geq\theta_{1}\geq\theta_{2}\\ \frac{3}{4}\theta_{2}+\frac{1}{4}\theta_{1}+\frac{1}{2}&\quad\text{if }\theta_ {2}-2\leq\theta_{1}\leq\theta_{2}\\ \theta_{2}&\quad\text{if }\theta_{1}\leq\theta_{2}-2\end{cases}$ (66) $$

Figure: Figure 6: An illustration of the function $f(\theta)$ (Equation 66) and the direction vector $v:v_{1}\geq v_{2}$ at $\theta=(0,0)$.
Refer to caption: x17.png

Recall, that we aim to prove that the Perturb-Argmax probability model is unidentifiable. The function $f(\theta)$, illustrated in Figure 6 in the appendix, is continuous and differentiable almost everywhere.
However, in its overlapping segments, i.e., when $\theta_{1}=\theta_{2}+2$, $\theta_{1}=\theta_{2}$ and $\theta_{1}+2=\theta_{2}$, the function is not differentiable, i.e., it has a sub-differential $\partial f(\theta)$ which is a set of sub-gradients. To prove that the Perturb-Argmax probability model is unidentifiable, we show that $\partial f(\theta)$ is a multi-valued mapping when $\theta_{1}=\theta_{2}$. In particular, we show that every probability distribution $p=(p_{1},p_{2})$ with $p_{1}\in[\frac{1}{4},\frac{3}{4}]$ satisfies $p\in\partial f(\theta)$.

For this task, we recall the connection between sub-gradients and directional derivatives: $p\in\partial f(\theta)$ if $\nabla_{v}f(\theta)\geq\langle p,v\rangle$ for every $v\in\mathbb{R}^{2}$.
When $\theta_{1}=\theta_{2}=c$, then $f(\theta)=c+\frac{1}{2}$, thus for the direction $v=(v_{1},v_{2})$ for which $v_{1}\geq v_{2}$ holds $\nabla_{v}f(\theta)=\frac{3}{4}v_{1}+\frac{1}{4}v_{2}$.
Recall that $p\in\partial f(\theta)$ if $\frac{3}{4}v_{1}+\frac{1}{4}v_{2}\geq\langle p,v\rangle$ for every $v_{1}\geq v_{2}$. Thus we conclude that $p$ must satisfy $p_{1}\leq\frac{3}{4}$. Since the same holds to $v_{2}\geq v_{1}$ then $p_{1}\geq\frac{1}{4}$. Taking both these conditions, $p\in\partial f(\theta)$ when $\frac{1}{4}\leq p_{1}\leq\frac{3}{4}$.
Therefore, $\partial f(\theta)$ is multi-valued mapping, or equivalently, the parameters $\theta=(\theta_{1},\theta_{2})$ are not identifying probability distributions. The sub-differential mapping is

$$ $(\partial f(\theta))_{1}=\begin{cases}1&\text{if }\theta_{1}>\theta_{2}+2\\ [\frac{3}{4},1]&\text{if }\theta_{1}=\theta_{2}+2\\ \frac{3}{4}&\text{if }\theta_{2}+2>\theta_{1}>\theta_{2}\\ [\frac{1}{4},\frac{3}{4}]&\text{if }\theta_{1}=\theta_{2}\\ \frac{1}{4}&\text{if }\theta_{2}-2<\theta_{1}<\theta_{2}\\ [0,\frac{1}{4}]&\text{if }\theta_{1}=\theta_{2}-2\\ 0&\text{if }\theta_{1}<\theta_{2}-2.\end{cases},$ (67) $$

as illustrated in Figure 3.

### 8.5 Proof of Proposition 5.4

The perturb-max function $f(\theta)$ can be expressed as

$$ $\displaystyle f(\theta)$ $\displaystyle=\mathbb{E}_{\gamma_{1},\gamma_{2}\sim[-1,1]}\left[\max_{i}\{ \theta_{i}+\gamma_{i}\}\right]$ (68) $\displaystyle=\mathbb{E}_{\gamma_{1},\gamma_{2}}\left[\max\{\theta_{1}+\gamma_ {1},\theta_{2}+\gamma_{2}\}\right]-E_{\gamma_{2}}\left[\gamma_{2}\right]$ (69) $\displaystyle=\mathbb{E}_{\gamma_{1},\gamma_{2}}\left[\max\{\theta_{1}+\gamma_ {1},\theta_{2}+\gamma_{2}\}-\gamma_{2}\right]$ (70) $\displaystyle=\mathbb{E}_{\gamma_{1},\gamma_{2}}\left[\max\{\theta_{1}+\gamma_ {1}-\gamma_{2},\theta_{2}\}+\theta_{2}-\theta_{2}\right]$ (71) $\displaystyle=\theta_{2}+\mathbb{E}_{\gamma_{1},\gamma_{2}}\left[\max\{\theta_ {1}-\theta_{2}+\gamma_{1}-\gamma_{2},0\}\right],$ (72) $$

when the second equation holds since $E_{\gamma_{2}\sim[-1,1]}\left[\gamma_{2}\right]=0$, and the distribution of $\gamma$ is omitted for brevity.

Define $\theta=\theta_{1}-\theta_{2}$ and $Z=\gamma_{1}-\gamma_{2}$. Then, the random variable $Z$ has a triangular distribution. The random variable $Z$ has the following cdf:

$$ $\displaystyle F_{Z}(z)$ $\displaystyle=P(\gamma_{1}-\gamma_{2}\leq z)$ (73) $\displaystyle=\begin{cases}0&\quad\text{if }z<-2\\ \frac{1}{4}\frac{(2+z)^{2}}{2}&\quad\text{if }0>z\geq-2\\ \frac{1}{4}(4-\frac{(2-z)^{2}}{2})&\quad\text{if }2\geq z\geq 0\\ 0&\quad\text{if }z>2\end{cases}$ (74) $$

Figure: Figure 7: Simulation of the density of the difference between 1M iid uniform(-1,1) random variables.
Refer to caption: x18.png

The random variable $Z$ has the following pdf, also verified in simulation of the density of the difference between 1M iid $U(-1,1)$ random variables (Figure 7):

$$ $\displaystyle f_{Z}(z)$ $\displaystyle=$ $\displaystyle\begin{cases}\frac{1}{4}(2+z)\quad\text{if }0>z\geq-2\\ \frac{1}{4}(2-z)\quad\text{if }2\geq z\geq 0\\ 0\quad\text{otherwise }\end{cases}$ (75) $$

With the pdf of the random variable $Z$ (Equation 75) consider $f(\theta)$ the appropriate range of $\theta$.

- 1.
Case: $-\theta<-2$
In this case $\theta_{1}-\theta_{2}>2$, therefore
$\displaystyle f(\theta)$
$\displaystyle=\mathbb{E}_{\gamma_{1},\gamma_{2}\sim[-1,1]}\left[\max_{i}\{
\theta_{i}+\gamma_{i}\}\right]$
(76)
$\displaystyle=\mathbb{E}_{\gamma_{1}\sim[-1,1]}\left[\theta_{1}+\gamma_{1}\ \right]$
(77)
$\displaystyle=\theta_{1}+\mathbb{E}_{\gamma_{1}\sim[-1,1]}\left[\gamma_{1}\ \right]$
(78)
$\displaystyle=\theta_{1}$
(79)
- 2.
Case: $0\leq-\theta\leq 2$
$\displaystyle f(\theta)$
$\displaystyle=\theta_{2}+\mathbb{E}_{z}\left[\max\{\theta+z,0\}\right]$
(80)
$\displaystyle=\theta_{2}+\int_{-\theta}^{2}f_{Z}(z)(\theta+z)dz$
(81)
$\displaystyle=\theta_{2}+\int_{-\theta}^{2}\frac{1}{4}(2-z)(\theta+z)dz$
(82)
$\displaystyle=\theta_{2}+\frac{1}{2}\int_{-\theta}^{2}\theta dz+\frac{1}{2}
\int_{-\theta}^{2}zdz-\frac{1}{4}\int_{-\theta}^{2}z\theta dz-\frac{1}{4}\int_
{-\theta}^{2}z^{2}dz$
(83)
$\displaystyle=\frac{1}{4}\left(\frac{4}{3}+2(\theta_{1}+\theta_{2})+(\theta_{1
}-\theta_{2})^{2}+\frac{1}{6}(\theta_{1}-\theta_{2})^{3}\right)$
(84)
- 3.
Case: $-2\leq-\theta\leq 0$
$\displaystyle f(\theta)$
$\displaystyle=\theta_{2}+\mathbb{E}_{z}\left[\max\{\theta+z,0\}\right]$
(85)
$\displaystyle=\theta_{2}+\int_{-\theta}^{0}f_{Z}(z)(\theta+z)dz+\int_{0}^{2}f_
{Z}(z)(\theta+z)dz$
(86)
$\displaystyle=\theta_{2}+\int_{-\theta}^{0}\frac{1}{4}(2+z)(\theta+z)dz+\int_{
0}^{2}\frac{1}{4}(2-z)(\theta+z)dz$
(87)
$\displaystyle=\theta_{2}+\frac{1}{4}(\theta^{2}-\frac{1}{6}\theta^{3})+\frac{1
}{4}(2\theta+\frac{4}{3})$
(88)
$\displaystyle=\frac{1}{4}\left(\frac{4}{3}+2(\theta_{1}+\theta_{2})+(\theta_{1
}-\theta_{2})^{2}-\frac{1}{6}(\theta_{1}-\theta_{2})^{3}\right)$
(89)
$\displaystyle\int_{-\theta}^{0}\frac{1}{4}(2+z)(\theta+z)dz$
$\displaystyle=\frac{1}{2}\int_{-\theta}^{0}\theta dz+\frac{1}{2}\int_{-\theta}
^{0}zdz+\frac{1}{4}\int_{-\theta}^{0}z\theta dz+\frac{1}{4}\int_{-\theta}^{0}z
^{2}dz$
(90)
$\displaystyle=\frac{1}{2}\left[\theta z\right]_{-\theta}^{0}+\frac{1}{2}\left[
\frac{z^{2}}{2}\right]_{-\theta}^{0}+\frac{1}{4}\left[\frac{z^{2}\theta}{2}
\right]_{-\theta}^{0}+\frac{1}{4}\left[\frac{z^{3}}{3}\right]_{-\theta}^{0}$
(91)
$\displaystyle=\frac{1}{4}(\theta^{2}-\frac{1}{6}\theta^{3})$
(92)
$\displaystyle\int_{0}^{2}\frac{1}{4}(2-z)(\theta+z)dz$
$\displaystyle=\int_{0}^{2}\frac{1}{4}(2\theta+2z-z\theta-z^{2})dz$
(93)
$\displaystyle=\frac{1}{2}\int_{0}^{2}\theta dz+\frac{1}{2}\int_{0}^{2}zdz-
\frac{1}{4}\int_{0}^{2}z\theta dz-\frac{1}{4}\int_{0}^{2}z^{2}dz$
(94)
$\displaystyle=\frac{1}{2}\left[\theta z\right]_{0}^{2}+\frac{1}{2}\left[\frac{
z^{2}}{2}\right]_{0}^{2}-\frac{1}{4}\left[\frac{z^{2}\theta}{2}\right]_{0}^{2}
-\frac{1}{4}\left[\frac{z^{3}}{3}\right]_{0}^{2}$
(95)
$\displaystyle=\frac{1}{4}(2\theta+\frac{4}{3})$
(96)
- 4.
Case: $-\theta>2$
In this case $\theta_{1}-\theta_{2}<-2$, therefore
$\displaystyle f(\theta)$
$\displaystyle=\mathbb{E}_{\gamma_{1},\gamma_{2}\sim[-1,1]}\left[\max_{i}\{
\theta_{i}+\gamma_{i}\}\right]$
(97)
$\displaystyle=\mathbb{E}_{\gamma_{2}\sim[-1,1]}\left[\theta_{2}+\gamma_{2}\ \right]$
(98)
$\displaystyle=\theta_{2}+\mathbb{E}_{\gamma_{2}\sim[-1,1]}\left[\gamma_{2}\ \right]$
(99)
$\displaystyle=\theta_{2}$
(100)

To conclude,

$$ $\displaystyle f(\theta)$ $\displaystyle=\begin{cases}\theta_{1}\quad\text{if }-\theta<-2\\ \frac{1}{4}\left(\frac{4}{3}+2(\theta_{1}+\theta_{2})+(\theta_{1}-\theta_{2})^ {2}+\frac{1}{6}(\theta_{1}-\theta_{2})^{3}\right)\quad\text{if }0\leq-\theta \leq 2\\ \frac{1}{4}\left(\frac{4}{3}+2(\theta_{1}+\theta_{2})+(\theta_{1}-\theta_{2})^ {2}-\frac{1}{6}(\theta_{1}-\theta_{2})^{3}\right)\quad\text{if }-2\leq-\theta \leq 0\\ \theta_{2}\quad\text{if }-\theta>2\end{cases}$ (101) $$

Alternatively, one writes

$$ $\displaystyle f(\theta)$ $\displaystyle=\begin{cases}\theta_{1}\quad\text{if }\theta>2\\ \frac{1}{4}\left(\frac{4}{3}+2(\theta_{1}+\theta_{2})+(\theta_{1}-\theta_{2})^ {2}-\frac{1}{6}(\theta_{1}-\theta_{2})^{3}\right)\quad\text{if }2\geq\theta \geq 0\\ \frac{1}{4}\left(\frac{4}{3}+2(\theta_{1}+\theta_{2})+(\theta_{1}-\theta_{2})^ {2}+\frac{1}{6}(\theta_{1}-\theta_{2})^{3}\right)\quad\text{if }0\geq\theta \geq-2\\ \theta_{2}\quad\text{if }\theta<-2\end{cases}$ (102) $$

The derivative of $f(\theta)$ (Equation 101), $\frac{\partial}{\partial\theta}f(\theta)=(\frac{\partial}{\partial\theta_{1}}f
(\theta),\frac{\partial}{\partial\theta_{2}}f(\theta))$ corresponds to the probabilities of the arg max, $(P_{\gamma}\left(\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=1\right),P_{\gamma}
\left(\arg\max_{i}\{\theta_{i}+\gamma_{i}\}=2\right))$ :

$$ $\displaystyle\frac{\partial}{\partial\theta}f(\theta)$ $\displaystyle=\begin{cases}(1,0)\quad\text{if }-\theta<-2\\ (\frac{1}{2}+\frac{1}{2}\theta+\frac{1}{8}\theta^{2},\frac{1}{2}-\frac{1}{2} \theta-\frac{1}{8}\theta^{2})\quad\text{if }0\leq-\theta\leq 2\\ (\frac{1}{2}+\frac{1}{2}\theta-\frac{1}{8}\theta^{2},\frac{1}{2}-\frac{1}{2} \theta+\frac{1}{8}\theta^{2})\quad\text{if }-2\leq-\theta\leq 0\\ (0,1)\quad\text{if }-\theta>2\end{cases}$ (103) $$

$$ $\displaystyle\frac{\partial}{\partial\theta}f(\theta)$ $\displaystyle=\begin{cases}(1,0)\quad\text{if }\theta>2\\ (\frac{1}{2}+\frac{1}{2}\theta-\frac{1}{8}\theta^{2},\frac{1}{2}-\frac{1}{2} \theta+\frac{1}{8}\theta^{2})\quad\text{if }2\geq\theta\geq 0\\ (\frac{1}{2}+\frac{1}{2}\theta+\frac{1}{8}\theta^{2},\frac{1}{2}-\frac{1}{2} \theta-\frac{1}{8}\theta^{2})\quad\text{if }0\geq\theta\geq-2\\ (0,1)\quad\text{if }\theta<-2\end{cases}$ (104) $$

Then, the partial derivatives $\frac{\partial}{\partial\theta}f(\theta)\in[0,1]$ and sum to $1$, as expected.

$$ $\frac{\partial}{\partial\theta_{1}}f(\theta)+\frac{\partial}{\partial\theta_{2 }}f(\theta)=1$ (105) $$

- 1.
Case: $0\leq-\theta\leq 2$
$\displaystyle\frac{\partial}{\partial\theta}f(\theta)$
$\displaystyle=(\frac{1}{4}(2+2\theta_{1}-2\theta_{2}+\frac{3\theta_{1}^{2}-6
\theta_{1}\theta_{2}+3\theta_{2}^{2}}{6}),\frac{1}{4}(2+2\theta_{2}-2\theta_{1
}+\frac{-3\theta_{1}^{2}+6\theta_{1}\theta_{2}-3\theta_{2}^{2}}{6}))$
(106)
$\displaystyle=(\frac{1}{2}+\frac{1}{2}\theta+\frac{1}{8}\theta^{2},\frac{1}{2}
-\frac{1}{2}\theta-\frac{1}{8}\theta^{2})$
(107)
The global minimum of the derivative $f(\theta)$ w.r.t. $\theta_{1}$, $\min\frac{\partial}{\partial\theta_{1}}f(\theta)=0$ for $\theta=-2$, since
$\frac{\partial\frac{1}{4}(2+2\theta_{1}-2\theta_{2}+\frac{3\theta_{1}^{2}-6
\theta_{1}\theta_{2}+3\theta_{2}^{2}}{6})}{\partial\theta_{1}}=\frac{1}{4}(2+
\theta_{1}-\theta_{2}),$
(108)
in which case $\frac{\partial}{\partial\theta_{2}}f(\theta)=1$.
The global maximum of the derivative $f(\theta)$ w.r.t. $\theta_{1}$, $\max\frac{\partial}{\partial\theta_{1}}f(\theta)=\frac{1}{2}$ for $-\theta=0$, in which case $\frac{\partial}{\partial\theta_{2}}f(\theta)=\frac{1}{2}$.
The global maximum of the derivative $f(\theta)$ w.r.t. $\theta_{2}$, $\max\frac{\partial}{\partial\theta_{2}}f(\theta))=1$ for $\theta_{2}=\theta_{1}+2$, since
$\frac{\partial\frac{1}{4}(2+2\theta_{2}-2\theta_{1}+\frac{-3\theta_{1}^{2}+6
\theta_{1}\theta_{2}-3\theta_{2}^{2}}{6}))}{\partial\theta_{2}}=\frac{1}{4}(2+
\theta_{1}-\theta_{2}),$
(109)
in which case $\frac{\partial}{\partial\theta_{1}}f(\theta)=0$.
The global minimum of the derivative $f(\theta)$ w.r.t. $\theta_{2}$, $\max\frac{\partial}{\partial\theta_{2}}f(\theta)=\frac{1}{2}$ for $-\theta=0$, in which case $\frac{\partial}{\partial\theta_{1}}f(\theta)=\frac{1}{2}$.

## 9 Experiments

We use a $16$GB 6-Core Intel Core i7 CPU in both experiments.

### 9.1 Approximating discrete distributions

Our experiments are based on the publicly available code of Potapczynski et al. ().
In all experiments, a thousand samples from a target distribution $p_{0}$ are sampled to approximate its probability density function parameters, based on the $L_{1}$ objective function. Optimization is based on the Adam optimizer (Kingma & Ba, ) with a learning rate $1.e-2$. The fitted parameters are initialized uniformly, as visualized in Figure 8.
Figure 9 shows results for fitting discrete distributions with countably
infinite support: a Poisson distribution with $\lambda=50$, and a negative binomial distribution with $r=50,p=0.6$ respectively. The Normal-Softmax distribution better approximates both distributions and exhibits faster
convergence than the Gaussian-Softmax.

Table 1 shows the average and standard deviation $L_{1}$ objective between the approximated probability density function and the four target discrete distributions of both models after $300$ iterations, computed over the dimension $d$ of the fitted models. The approximation based on the Normal-Softmax probability model achieves lower mean and standard deviation $L_{1}$ in all cases.

**Table 1: $L_{1}$ average and standard deviation between the target and approximated probability density function of various target discrete distributions of the Normal-Softmax and Gumbel-Softmax probability models.**
| Target Distribution | Normal-Softmax | Gumbel-Softmax |
| --- | --- | --- |
| Discrete | 0.026 $\pm$ 0.002 | 0.027$\pm$0.002 |
| Binomial | 0.036$\pm$0.003 | 0.177 $\pm$0.014 |
| Poisson | 0.090$\pm$0.001 | 0.618 $\pm$0.007 |
| Negative Binomial | 0.083 $\pm$ 0.001 | 0.417 $\pm$ 0.004 |

Figure: (a) a binomial distribution
Refer to caption: x19.png

Figure: (a) Approximation of a Poisson distribution.
Refer to caption: x27.png

### 9.2 Variational inference

This experiment is based on the publicly available implementation of the Gumbel-Softmax-based implementation of the discrete VAE in [Direct-VAE](https://github.com/GuyLor/Direct-VAE). Optimization is based on the Adam optimizer (Kingma & Ba, ) with a learning rate of $1.e-3$. Batch size is set to $100$.
We use the regular train/ test splits and follow previous research splits (e.g., as in Lorberbom et al. ()). The MNIST and the Fashion-MNIST datasets’ training set comprises $60,000$ images, and the test set comprises $10,000$ images. For the Omniglot, the training set comprises $24,345$ images, and the test set comprises $8,070$ images.

Figure 10 shows the results of the variational inference experiment for the Fashion-MNIST dataset (Xiao et al., ) with $N=10$ discrete variables, each is a $K$-dimensional categorical variable, $K\in[10,30,50]$.

Figure: Figure 10: Categorical VAE with Perturb-Softmax training loss on the the Fashion-MNIST dataset with a $K$-dimensional categorical variable, $K\in[10,30,50]$.
Refer to caption: x33.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: - How can Theorem 5.1 be corrected? Does the mistake also apply to Theorem 4.1?

- Do you have any insights as to why Gaussian-Softmax seems to be more learnable in these settings than Gumbel-Softmax? Can those insights be generalized to help guide to other noise distributions that might be even better?
soundness: 2
strengths: The idea of alternatives to the Gumbel trick is interesting. It indeed relates to FTPL (and similarly, I think, to various non-exponential mechanisms used in differential privacy), but in this context to my knowledge the Gumbel version is the only one that has been thoroughly considered.

The experimental results are promising that learning with Gaussian noise may be preferable to Gumbel noise.
summary: The paper studies representation properties of "Perturb-Softmax" and "Perturb-Argmax" distributions, generalizations of the Gumbel-Softmax/Gumbel-Max distributions used in a significant line of prior machine learning work on discrete random variables. They show theorems establishing general representation results, and experiments arguing that perturbing with a Gaussian rather than a Gumbel yields improved learning performance.
weaknesses: While the setting of this paper is interesting, the presentation is a little bit disjointed:
- Your conclusion says that "Our main contribution is a theoretical study of the representation properties of Gumbel-Softmax and Gumbel-Argmax probability models." I think this is not really true; the basic derivation of Gumbel-Argmax shows that it is obviously equivalent to general softmax parameterizations of discrete distributions, in which case the completeness and minimality properties are well-known. While I don't think either of the two Gumbel-Softmax papers made explicit arguments about completeness, both included a temperature parameter, and the equivalence of Gumbel-Softmax to Gumbel-Argmax in the limit of low temperatures – explicitly stated as Proposition 1 part (c) of Maddison et al. – immediately and obviously implies completeness there. Thus most of these results were implicit in the literature already and the new ones are not _especially_ interesting in the machine learning context.
- What is quite interesting, though, is the claim that perturbing with other forms of noise can yield the same kind of completeness (and, less relevantly for learning, minimality) properties. The experiments with Gaussian-Softmax also indicate that other kinds of noise are worth exploring, although it would be nice to have a better understanding of _why_ (see discussion about the "convergence rates" argument below) or of how much improvement can be seen in a wider variety of settings. So, I think the paper should be somewhat reframed as being about the generalized Perturb-*max models, and how other choices than Gumbel can be better.
    - To this end: while the experiments here are definitely not _nothing_, they also contain no detail other than "here is the final approximation quality," and they're also not of an especially exciting scale in 2024. Ideally, I'd like to see some better poking at why optimization worked better in a simple case, e.g. one of the fixed discrete target distributions (is the parameter vector $\\theta$ to get a good approximation simpler? if not, is something about the optimization landscape simpler?). It would also be nice to have experiments for a more modern model using this as a component; just take one of the many recent models citing these papers and swap out the Gumbel.

But, before doing that, a vital issue:

- Theorem 5.1 as stated here is **plainly false**. In the theorem statement, there is no constraint on the form of random variable $\gamma$. First, they are presumably intended to be iid, but maybe this doesn't matter. So, consider using the random variable $\gamma_i$ which is a point mass at 0. Then the Perturb-Argmax distribution is $\\mathbb{E}_\\gamma[ \\mathrm{arg\\,max}\\; \\theta + \\gamma ] = \\mathrm{arg\\,max}\\; \\theta$; depending on exactly how we define argmax's behaviour in the case of ties, this can _only_ produce one-hot distributions. Since the structure of your proof is that one-hots are in the closure of $\\mathcal P$ and that this closure is convex, presumably you didn't notice this when adapting the proof of Theorem 4.1, since indeed that first property is certainly true. Presumably, then, the claim that the closure of $\\mathcal P$ is convex is false in this case. I don't know where the flaw is, but something must be wrong. In particular, since Theorem 4.1 follows a similar structure, it is particularly concerning that there is some unknown flaw that might also apply to that theorem's proof. While it's not immediately obvious to me what the problem is, I'm not so confident in how to carefully apply Rockafellar's results that you use (which are, I think, stated for $\\mathbb R^d$) to the probability simplex, and so am not confident that it's a mistake in Theorem 5.1 specifically. I cannot see arguing for anything other than rejecting the paper unless this mistake is identified and corrected, and a full correct proof provided during the rebuttal process.

Perhaps relatedly, the mathematical discussion and definition in several aspects of this paper is quite sloppy. Here are a few reasonably important ones:
- Line 195 claims that "The fundamental theorem of extreme value statistics asserts the equivalence between the softmax distribution in Equation 4 and the Gumbel-Argmax distribution in Equation 2." I don't see how this is in any way true. You only cited a 60-page lecture series which does not use this phrase to identify a theorem; presumably, you mean a [version of this theorem](https://en.wikipedia.org/wiki/Fisher%E2%80%93Tippett%E2%80%93Gnedenko_theorem), which states that the maximum of $n$ iid values can have one three asymptotic behaviours as $n \\to \\infty$, one of which is the Gumbel distribution. I don't see at all what this has to do with the distribution of a _soft_ max of a finite number of values. Luckily, though, you don't seem to use this seemingly-incorrect claim again.
- Section 3.4 "Convergence rates" compares the upper bounds on concentration of Lipschitz functions of a Gaussian variable and a Gumbel variable, and since the former is faster than the latter, vaguely imply that this means that Gaussian-Softmax is better than Gumbel-Softmax. But the purpose of the remainder of your paper is that you would use the two to identify _the exact same discrete distribution_, and thus estimating $\mathbb E\_{\\gamma \\sim g}[ f(\\theta, x, \\gamma) ]$ by Monte Carlo samples of $\\gamma$ will be _identical_ between equivalent choices of $\\theta$ with different distributions for $\\gamma$. These differences in upper bounds (which are worst-case over functions $\\gamma \\mapsto f(\\theta, x, \\gamma)$) are not relevant here.

And here are a few examples that are easy to fix, but indicative that perhaps you were not as careful in writing as you should have been:
- On line 172 you say that the probability density function of the Gumbel is $\\hat g(t) = e^{-e^{-(t+c)}}$ where $c \\approx 0.5772$. This is not true; this is the _cumulative_ density function, specifically of a Gumbel with mean $c$ and scale $1$. While scale $1$ is standard for the Gumbel-max trick, there is absolutely no reason to use mean $c$; it doesn't break anything, but by far the more typical choice would be to use mean $0$.
- Theorem 4.1 says "let $\gamma = (\\gamma\_1, \dots, \\gamma\_d)$ be a vector of random variables whose cumulative distribution decays to zero as $\\gamma$ approaches $\\pm \infty$." First, presumably you mean that the $\\gamma_i$ are iid, and it is the distribution of each of those that has the appropriate decay property. Secondly: it is trivially true of any real-valued random variable that its cumulative distribution function decays to zero as you approach $-\\infty$, and to _one_ as you approach $\\infty$. Perhaps you meant that the measure of sets $(-\\infty, -t)$ and $(t, \\infty)$ should approach zero as $t \\to \\infty$? But this is again automatically true of any real-valued variable. I actually have no idea what you mean here.
- Line 719: "A multivariate function is differentiable if its directional derivative is the same in every direction $v \\in \\mathbb R^d$, namely $\\nabla f(\\theta) = \\nabla_v f(\\theta)$ for every $v \\in \\mathbb R^d$." This is not at all true, as should be clear from the fact that your equation is equating a vector to a scalar. Rather, the directional derivatives should be _consistent_, i.e. $\\nabla f(\\theta) \cdot v = \\nabla_v f(\\theta)$ for all $v$.

Minor points of terminology:
- Your definition of identifiability is different than the one I've always heard before, and is also used e.g. [by Wikipedia](https://en.wikipedia.org/wiki/Identifiability), which is exactly what you call "minimal." I don't see much reason in giving a special name to the bijection case here, and particularly in giving one that is so widely used to mean something slightly different; I would strongly encourage you to change the terminology to "complete" (or perhaps "universal" would be more common in learning contexts) when the map from parameters to probability distributions is (almost) a surjection, and "identifiable" when it is an injection.
- Typically one would put `\appendix` at the start of the appendix, which would name the appendix A rather than 8.
- Line 713: "Convexity is a one-dimensional property." This is a strange statement; while you can define it in one dimension and extend it to other cases as you do here, there are also (many) direct definitions of convexity that work for any input vector space.

Trivial typos, etc that I happened to notice:
- line 160: you wrote $softmax$ instead of $\\mathrm{softmax}$
- "subsetset" on line 209
- Footnote 3: "dominant convergence theorem" should be "dominated convergence theorem"

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: In practice, the Gumbel-softmax is more widely utilized than the Gumbel-max, and the temperature parameter is naturally added to the distribution setting. However, it seems that it is dropped in the theoretical analysis. Why is that?
soundness: 2
strengths: The paper is clearly written and well-organized.

The work is theoretically grounded.
summary: The paper studies the Gumbel-max distributions in the mathematical view. The authors especially investigate the completeness, identifiability, and minimality of the distributions, i.e., the Purturb-Maxes. Finally, the advantage of Gaussian-softmax, which is an example of the Purturb-softmax, is demonstrated in the experiment section.
weaknesses: Even though this is a theoretical paper, the work requires empirical evidence more than the ones suggested in the manuscript. 

Also, it is not clear how the current experiments are related to the theories provided in the main body.

The proofs can be moved to the appendix.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - Should line 160 only consider $p \in \operatorname{ri}(\Delta)$?
- Should Theorem 3.1 include the sub-Gaussianity parameter?
- Should line 240 say "Gumbel" instead of "exponential family"?
- In Theorem 4.1, what does "whose cumulative distribution decays to zero" mean? Can you make this precise, as CDFs do not decay to zero?
soundness: 3
strengths: - The paper is concise and overall well-written (see Weaknesses for feedback). In particular, the proofs in the main text and footnotes are helpful and executed well; usually, this type of presentation is hard to get right.
- The mathematical analysis is elegant and general enough to account for any perturbation distribution.
- The experimental testbed is generally helpful, although I do not believe the provide much insight regarding the theory (see Weaknesses).
- The appendices are helpful for making the paper self-contained overall, but the proof appendices should not be theorem-proof lists. Try to make them self-contained as well; concretely, please reintroduce the theorem, add proof outlines, discuss any techniques that are novel to this paper versus those that are adapted from others with citations, etc)/
summary: The paper characterizes the the ability of particular parametrized families of probability mass functions (PMFs) to fit an arbitrary PMF on a finite probability space. The main result is a set of conditions under which the Perturbed-Softmax and Perturbed-Argmax families enjoy completeness (whether the parameter-to-distribution map is onto) and/or minimality (whether the parameter-to-distribution map is one-to-one).
weaknesses: - The earlier part of the presentation could be improved, in that components from the Related Work and Preliminaries could be incorporated into the Introduction. For example, I don't believe the paper mentions one of the major motivations of the Gumbel-Softmax as an approximation to the discrete sampling operation in VAEs that can be backpropagated through. Similarly, while online learning methods are stated as a motivation, the reader should be able to know where the perturbed-argmax/softmax operations appear in their methods, even without having a background in online learning. Equations for the perturbed operations could appear much earlier in the text, so that the Gumbel-softmax/argmax, Gaussian-softmax/argmax, and their perturbed variants are not so abstract. 
- Unless I may be misunderstanding, their seem to be some technical errors in the main text proofs. For example, the proof of Theorem 4.1 does not use the $h_i$ function, so it must be incomplete. This assumption is in fact used in Theorem 5.1, but not included in the theorem statement.
- I do not believe the current experiments align well with rest of the paper. In particular, what is shown is that Gaussian perturbations result in a predictor that learns faster with respect to validation loss. However, the rest of the paper is about the representation capabilities of parametrized models, and if I understand correctly, the Gumbel and Gaussian perturbed models have the same representation properties. Moreover, this difference between Gaussian and Gumbel performance seems to already be understood in theory based on these concentration properties shown earlier in the paper. I felt that an experiment which included a model that did not satisfy completeness conditions (within the same setup as 6.2) and shows the bias of this model at the learned minimizer would be much more illuminating.

Meta-review: This work investigates the representation properties of the Gumbel-Softmax and Gumbel-Argmax probability distributions, which are used for learning discrete tokens and structures in generative and discriminative models, respectively. The study identifies the conditions under which these distributions are complete (able to represent any probability distribution) and minimal (able to represent a distribution uniquely), using convexity and differentiability. The analysis extends to general probability models like Perturb-Softmax and Perturb-Argmax, concluding that certain parameter sets allow for complete and minimal representations, with experimental results validating the faster convergence of Gaussian-Softmax compared to Gumbel-Softmax.

The problem addressed by the paper is clearly important and practical.  Although the paper is meant to be theoretical, the experimental results are really minimal and lack some important comparisons as the reviewers have pointed out.  There are also a few concerns raised by Reviewer e4WY that remain not fully satisfactorily addressed in the rebuttal.  So the paper needs one more round of revision.

Decision: Reject