Learning with Differentiable Perturbed Optimizers

Introduction

Many applications of machine learning benefit from the possibility to train by gradient descent compositional models using end-to-end differentiability. Yet, there remain fields where discrete decisions are required at intermediate steps of a data processing pipeline (e.g., in robotics, graphics or biology). This is the result of many factors: discrete decisions provide a much sought-for interpretability, and discrete solvers are built upon decades of advances in combinatorial algorithms (Schrijver,, 2003) for quick decisions (e.g., sorting, picking closest neighbors, exploring options with beam-search, or with shortest paths problems). These discrete decisions can easily be computed in a forward pass. Their derivatives with respect to inputs are however degenerate: small changes in the inputs either yield no change or discontinuous changes in the outputs. Discrete solvers thus break the back-propagation of computational graphs, and cannot be incorporated in end-to-end learning.

In order to expand the set of operations that can be incorporated in differentiable models, we propose and investigate a new, systematic method to transform discrete optimizers into differentiable operations. Our approach builds upon the method of stochastic perturbations, the theory of which was developed and applied to several tasks of machine learning recently; see Hazan et al., (2016). In a nutshell, we perturb the inputs of a discrete solver with random noise, and consider the perturbed solutions of the problem. The method is both easy to analyze theoretically and simple to implement. We show that the formal expectations of these perturbed solutions are never locally constant and everywhere differentiable, with successive derivatives being expectations of simple expressions.

Related work. Our work is part of growing efforts to modify operations to make them differentiable. Several works have studied the introduction of regularization in the optimization problem to make the argmax differentiable. These works are usually problem-specific, since a new optimization problem needs to be solved. Examples include assignments (Adams and Zemel,, 2011), optimal transport (Bonneel et al.,, 2016; Cuturi,, 2013), differentiable dynamic programming (Mensch and Blondel,, 2018), differentiable submodular optimization (Djolonga and Krause,, 2017). A generic approach is SparseMAP (Niculae et al.,, 2018), based on Frank-Wolfe or active-set algorithms for solving, and on implicit differentiation for Jacobian computation. Like our proposal, SparseMAP only requires access to a linear maximization oracle. However, it is sequential in nature, while our approach is trivial to parallelize. In (Agrawal et al.,, 2019), implicit differentiation on solutions of convex optimization is analyzed. They express the derivatives of the argmax exactly, leading to zero Jacobian almost everywhere when optimizing over polytopes. Vlastelica et al., (2019) proposed to interpolate in a piecewise-linear manner between locally constant regions. The aim is to keep the same value for the Jacobian of the argmax for a large region of inputs, allowing for zero Jacobians as well.

An example of expectation of a perturbed argmax, commonly known as the “Gumbel trick”, dates back to Gumbel, (1954), and random choice models (Luce,, 1959; McFadden,, 1973; Guadagni and Little,, 1983). It is exploited in online learning and bandits to promote exploration, and induce robustness to adversaries (see, e.g., (Abernethy et al.,, 2016) for a survey). It is used for action spaces that are combinatorial in nature (Neu and Bartók,, 2016), and used together with a softmax to obtain differentiable sampling (Jang et al.,, 2016; Maddison et al.,, 2016), and with distributions from extreme value theory (Balog et al.,, 2017).

The use of perturbation techniques as an alternative to MCMC techniques for sampling was pioneered by Papandreou and Yuille, (2011). They are used to compute expected statistics arising in gradients of conditional random fields. They show exactness for the fully perturbed (but intractable case) and propose “low-rank” perturbations as an approximation. These results are extended in (Hazan and Jaakkola,, 2012), proving that the expected maximum with low-rank perturbations provides an upper-bound on the log partition, and replacing the log partition in conditional random fields loss by that expectation. Their results, however, are limited to discrete product spaces. New lower bounds on the partition function are derived in (Hazan et al.,, 2013), as well as a new unbiased sequential sampler for the Gibbs distribution based on low-rank perturbations. These results were further refined in (Gane et al.,, 2014) and (Orabona et al.,, 2014), and these bounds further studied in (Shpakova and Bach,, 2016), who proposed a doubly stochastic scheme. Apart from (Lorberbom et al.,, 2019), who use a finite difference method, we are not aware of any prior work using perturbation techniques to differentiate through an argmax. As reviewed above, all papers focus on (approximately) sampling from the Gibbs distribution, upper-bounding the log partition function, or differentiating through the max.

Contributions. We make the following contributions:

- We propose a new general method transforming discrete optimizers, inspired by the stochastic perturbation literature. This versatile method applies to any blackbox solver without ad-hoc modifications.

- Our stochastic smoothing allows argmax differentiation, through the formal perturbed maximizer. Its Jacobian is well-defined and non-zero everywhere, thereby avoiding vanishing gradients.

- The successive derivatives of the perturbed maximum and argmax are expressed as simple expectations, which are easy to approximate with Monte-Carlo methods.

- Our method yields natural connections to the recently-proposed Fenchel-Young losses by Blondel et al., (2019). We show that the equivalence via duality with regularized optimization makes these losses natural.

- We propose a doubly stochastic scheme for their minimization in learning tasks, and we demonstrate our method on structured prediction tasks, in particular ranking (permutation prediction), for which conditional random fields and the Gibbs distribution are intractable.

Perturbed maximizers

This creates a general and natural model on the variable $Y$, when observations are solutions of optimization problems, with uncertain costs. It enables the modeling of phenomena where agents chose an optimal $y\in\mathcal{C}$ based on uncertain knowledge of $\theta$. We view this as a generalization, or alternative to the Gibbs distribution, rather than an approximation thereof.

Taking expectations with respect to the random perturbation leads to smoothed versions of $F$ and $y^{*}$:

Models of random optimizers for linear problems with perturbed inputs are the subject of a wide litterature in machine learning, under the name of “perturb-and-MAP” Papandreou and Yuille, (2011); Hazan and Jaakkola, (2012), and perturbed leader method in online learning (Hannan,, 1957; Kalai and Vempala,, 2003; Abernethy et al.,, 2014). We refer to it here as the perturbed model.

A generalization of Gumbel-max.

An example of this setting is well-known: when $\mathcal{Y}$ is the set of one-hot-encoding of $d$ classes, $\mathcal{C}$ is the unit simplex, and $Z$ has the Gumbel distribution (Gumbel,, 1954). In that case it is well-known that $p_{\theta}$ is the Gibbs distribution, proportional to $\exp(\langle y,\theta\rangle/\varepsilon)$, $F_{\varepsilon}(\theta)$ is the log-sum-exp function of $\theta$, and $y^{*}_{\varepsilon}(\theta)$ is the vector of softmax (or exponential weights) of the components of $\theta$. Our model is therefore a generalization of the Gumbel-max setting. As $F_{\varepsilon}$ generalizes the log-sum-exp function for Gumbel noise on the simplex, its dual $\Omega$ is a generalization of the negative Shannon entropy (which is the Fenchel dual of the log-sum-exp function). We show this connection, and that the perturbed maximizer can also be defined as the solution of a convex problem, by Fenchel-Rockafellar duality in Proposition 2.1 below. The following table summarizes those parallels. Our framework generalizes these ideas, and proposes to exploit the ease of simulation of $p_{\theta}$ (rather than the explicit forms of Gibbs distributions) for applications in machine learning tasks.

Let $\Omega$ be the Fenchel dual of $F_{1}$, with domain $\mathcal{C}$. We have that

Differentiation and associated loss function.

While these connections have been studied before (Hazan et al.,, 2016; Abernethy et al.,, 2014, 2016), we provide two key new insights. First, the perturbed model allows to take derivatives with respect to the input $\theta$ of $F_{\varepsilon}$ and of $y^{*}_{\varepsilon}$ (Proposition 2.2). These derivatives are also easily expressed as expectations involving $F$ and $y^{*}$ with noisy inputs, as discussed in Section 3. In turn, this yields fast computational methods for these functions and their derivatives. Second, by the duality point of view describing $y^{*}_{\varepsilon}$ as a regularized maximizer, there exists a natural convex loss for this model that can be efficiently optimized in $\theta$, for data $y_{i}\in\mathcal{Y}$. We describe this formalism in Section 4, and apply it in experiments in Section 5.

Properties of the model.

This model modifies the maximum and maximizer by perturbation. Because of the simple action of the stochastic noise , we can analyze their properties precisely.

Assume $\mathcal{C}$ is a convex polytope with non-empty interior, and $\mu$ has positive differentiable density. The perturbed model $p_{\theta}$ and the associated functions $F_{\varepsilon}$, $\Omega=(F_{\varepsilon})^{\star}$, and $y^{*}_{\varepsilon}$ have the following properties, for $R_{\mathcal{C}}=\max_{y\in\mathcal{C}}\|y\|$ and $M_{\mu}=\mathbf{E}[\|\nabla_{z}\nu(Z)\|^{2}]^{1/2}$:

$F_{\varepsilon}$ is strictly convex, twice differentiable, $R_{\mathcal{C}}$-Lipschitz-continuous and its gradient is $R_{\mathcal{C}}M_{\mu}/\varepsilon$-Lipschitz-continuous. Its dual $\Omega$ is $1/(R_{\mathcal{C}}M_{\mu})$-strongly convex, differentiable, and Legendre-type.

Impact of $\varepsilon>0$: we have F_{\varepsilon}(\theta)=\varepsilon F_{1}\big{(}\frac{\theta}{\varepsilon}\big{)},\,F_{\varepsilon}^{*}(y)=\varepsilon\Omega(y),\,y^{*}_{\varepsilon}(\theta)=y^{*}_{1}\big{(}\frac{\theta}{\varepsilon}\big{)}.

For these properties to hold, it is crucial that $\mathcal{C}$ has non-empty interior, i.e., that $\mathcal{Y}$ does not lie in an affine subspace of lower dimension. To adapt to cases where $\mathcal{C}$ lies in a subspace, we consider the set of inputs $\theta$ up to vectors orthogonal to $\mathcal{C}$, or represent $\mathcal{Y}$ in a lower-dimensional subspace. As an example, over the unit simplex and Gumbel noise, the log-sum-exp is not strictly convex, and in fact linear along the all-ones vector $\mathbf{1}$. In such cases, the model is only well-specified in $\theta$ up to the space orthogonal to $\mathcal{C}$, which does not affect prediction tasks.

For any positive temperature $\varepsilon$, these properties imply that there is an informative, well-defined, and nonzero gradient in $\theta$. They also imply the limiting behavior at extreme temperatures.

With the conditions of Proposition 2.2, for $\theta$ such that $y^{*}(\theta)$ is a unique maximum:

For $\varepsilon\to 0$, $F_{\varepsilon}(\theta)\to F(\theta)$ and $y^{*}_{\varepsilon}(\theta)\to y^{*}(\theta)$. For $\varepsilon\to\infty$, $y^{*}_{\varepsilon}(\theta)\to y^{*}_{1}(0)=\operatorname*{arg\,min}_{y\in\mathcal{C}}\Omega(y)$.

For every $\varepsilon>0$, we have $F(\theta)-F_{\varepsilon}(\theta)\leq C\varepsilon$ and $\langle y^{*}(\theta),\theta\rangle-\langle y^{*}_{\varepsilon}(\theta),\theta\rangle\leq C^{\prime}\varepsilon$, for $C,C^{\prime}>0$.

The properties of the distributions $p_{\theta}$ in this model are well studied in the perturbations literature (see, e.g., (Hazan et al.,, 2016) for a survey). They notably do not have a simple closed-form expression, but can be very easy to sample from. By the argmax definition, simulating $Y\sim p_{\theta}$, only requires to sample $\mu$ (e.g., Gaussian, or vector of i.i.d. Gumbel), and to solve the original optimization problem. It is the case in the applications we consider (e.g., max, ranking, shortest paths). This is in stark contrast to the Gibbs distribution, which has the opposite properties.

Differentiation of soft maximizers

As noted above, for the right noise distributions, the perturbed maximizer $y^{*}_{\varepsilon}$ is differentiable in its inputs, with non-zero Jacobian. It is based on integration by parts, not on finite differences as in (Lorberbom et al.,, 2019).

The derivatives are simple expectations. We discuss in the following subsection efficient techniques to evaluate in practice $y^{*}_{\varepsilon}(\theta)$ and its Jacobian, or to generate stochastic gradients, based on these expressions.

For any $\theta$, the perturbed maximizer $y^{*}_{\varepsilon}(\theta)$ is a solution of a convex optimization problem in Eq. (2), allowing computation if $\Omega$ has a simple form. More generally, by their expressions as expectations, the perturbed maximizer and its Jacobian can be approximated with Monte-Carlo methods. This only requires to efficiently sample from $\mu$, and to solve LPs over $\mathcal{C}$.

A Monte-Carlo estimate $\bar{y}_{\varepsilon,M}(\theta)$ of $y_{\varepsilon}^{*}(\theta)$ is given by

Since $\mathbf{E}[y^{(m)}]=y_{\varepsilon}^{*}(\theta)$ for every $m\in\{1,\dots,M\}$, by definition of $p_{\theta}$, it is an unbiased estimate of $y_{\varepsilon}^{*}(\theta)$. Note that the formulae in Proposition 3.1 give several manners to stochastically approximate $F_{\varepsilon}$, $y^{*}_{\varepsilon}$, and their derivatives by using $F(\theta+\varepsilon Z^{(m)})$, $y^{*}(\theta+\varepsilon Z^{(m)})$ and $\nabla_{z}\nu(Z^{(m)})$ and averages. This yields unbiased estimates for $F_{\varepsilon}$, $y^{*}_{\varepsilon}$, and its Jacobian. The plurality of these formulae gives the user several options for practical implementation. For both $y^{*}_{\varepsilon}$ and its Jacobian, we use the first one presented in Proposition 3.1 for our applications.

A great strength of this method is the absence of conceptual or computational overhead. Further, even though our analysis relies on the specific structure of the problem as an LP, these algorithms do not. The Monte-Carlo estimates can be obtained by using a function $y^{*}$ as a blackbox, without requiring knowledge of the problem or of the algorithm that solves it. For instance, for ranking, solving the LP only involves a sort.

If $y^{*}_{\varepsilon}$ or its derivatives are used in stochastic gradient descent for training in supervised learning, a full approximation of the gradients is not always necessary. Taking only $M=1$ (or a small number) of observations is acceptable here, as the gradients are stochastic in the first place.

With parallelization and warm starts, we can alleviate the dependency in $M$ of the running time: We can independently sample the $Z^{(m)}$ and compute the $y^{(m)}=y^{*}(\theta+\varepsilon Z^{(m)})$ in parallel. On the other hand, starting from a solution or near-solution (such as $y^{*}(\theta)$) as initialization can improve running times dramatically, especially at lower temperatures.

Perturbed model learning with Fenchel-Young losses

There is a large literature on learning parameters of a Gibbs distribution based on data $(y_{i})_{i=1,\ldots,n}$, through maximization of the likelihood:

The expression of the gradient justifies the name of moment-matching procedures. The expectation of the Gibbs is however hard to evaluate in some cases. For instance, for permutation problems, it is known to be #P-hard to compute (Valiant,, 1979; Taskar,, 2004). This motivates its replacement by $p_{\theta}$ (perturb-and-MAP in this literature), and to use this method as a proxy for log-likelihood to learn the parameters (Papandreou and Yuille,, 2011).

We show here that this approach can be formally analyzed by the use of Fenchel-Young losses (Blondel,, 2019) in this context. It is equivalent to maximizing a term akin to Eq. (5), substituting the log-partition $Z(\theta)$ with $F_{\varepsilon}(\theta)$. The use of these losses also drastically improves the algorithmic aspects of the learning tasks, by the specific expression of the gradients of the loss.

It is nonnegative, convex in $\theta$, and minimized with value 0 if and only if $\theta$ is such that $y^{*}_{\varepsilon}(\theta)=y$. It is equal to the Bregman divergence associated to $\varepsilon\Omega$, i.e., $L_{\varepsilon}(\theta\,;y)=D_{\varepsilon\Omega}(y,\hat{y}^{*}_{\varepsilon}(\theta))$. As $\theta$ and $y$ interact in this loss only through a scalar product, for random $Y$ we have $\mathbf{E}[L_{\varepsilon}(\theta;Y)]=L_{\varepsilon}(\theta;\mathbf{E}[Y])+C$, where $C$ does not depend on $\theta$. This is particularly convenient in analyzing the performance of Fenchel-Young losses in generative models. The gradient of the loss is

The Fenchel-Young loss can therefore be interpreted as a loss in $\theta$ that is a function of $y^{*}_{\varepsilon}(\theta)$. Moreover, it can be optimized in $\theta$ with first-order methods simply by computing the soft maximizer, without having to compute its Jacobian. It is therefore a particular case of the situation described in Eq.(3) and (4), allowing to even bypass virtually the perturbed maximizer block in the output, and to directly optimize a loss between observation $y$ and model outputs $\theta=g_{w}(x)$.

As described in Remark 2, given observations $(x_{i},y_{i})_{1\leq i\leq n}\in\mathcal{X}^{n}\times\mathcal{Y}^{n}$, we can fit a model $g_{w}$ such that $y_{\varepsilon}^{*}(g_{w}(x_{i}))\approx y_{i}$. The Fenchel-Young loss between $g_{w}(x_{i})$ and $y_{i}$ is a natural way to do so

This is motivated by a generative model where, for some $w_{0}$

Indeed, under this model the population loss $\mathbf{E}[L_{\varepsilon,{\sf emp}}(w)]$ is the average of terms $L_{\varepsilon}(g_{w}(x_{i})\,;y^{*}_{\varepsilon}(g_{w_{0}}(x_{i})))$, up to an additive constant. The population loss is therefore minimized at $w_{0}$. The gradient of the empirical loss is given by

Each term in the sum, gradient of the loss for a single observation, is therefore a stochastic gradient for $L_{\varepsilon,\sf emp}$ (w.r.t. $i$ uniform in $[n]$) or for $L_{\varepsilon,\sf pop}$ (w.r.t. to a random $y_{i}$ from $p_{g_{w_{0}}(x_{i})}$).

The methods we described to stochastically approximate the gradient are particularly adapted here. Indeed, following (Shpakova and Bach,, 2016), given an observation $y_{i}$ and a current value $\theta_{i}=g_{w}(x_{i})$, a doubly stochastic version of the gradient $\nabla_{w}L_{\varepsilon}(g_{w}(x_{i})\,;y_{i})$ is obtained by

This can also be used with a procedure where batches of data points are used to compute approximate gradients, where the number of artificial samples $M$ and the batch size can be chosen separately.

This can be extended to an unsupervised setting, where observations $(y_{i})_{1\leq i\leq n}\in\mathcal{Y}^{n}$ are fitted with a model $p_{\theta}$, motivated by a generative model where $y_{i}=\operatorname*{arg\,max}_{y\in\mathcal{C}}\ \langle\theta_{0}+\varepsilon Z_{i},y\rangle$, that is $y_{i}\sim p_{\theta_{0}}(y)$, for some unknown $\theta_{0}$. We have a natural empirical $\bar{L}_{n}$ and population loss $L_{\theta_{0}}$:

The empirical loss is minimized for $\hat{\theta}_{n}$ such that $y^{*}_{\varepsilon}(\hat{\theta}_{n})=\bar{Y}_{n}$ and the population loss when $y^{*}_{\varepsilon}(\theta)=y^{*}_{\varepsilon}(\theta_{0})$. As a consequence, the whole battery of statistical results, from asymptotic to non-asymptotic, can be leveraged, and we present the simplest one (asymptotic normality).

When $n$ goes to $\infty$, with the assumptions of Proposition 2.2 on the model, we have

in distribution, where $\Sigma_{Y}$ is the covariance of $Y\sim p_{\theta}$.

Experiments

We demonstrate the usefulness of perturbed maximizers in a supervised learning setting, as described in Section 4. We focus on a classification task and on two structured prediction tasks, label ranking and learning to predict shortest paths. Since we focus on the prediction task, the issues raised in Remark 1 do not apply. When learning with the Fenchel-Young losses, we simulate doubly stochastic gradients $\nabla_{w}L_{\varepsilon}(g_{w}(x_{i})\,;y_{i})$ of the empirical loss with $M$ artificial perturbations (see Equation 6).

We will open-source a Python package allowing to turn any black-box solver into a differentiable function, in just a few lines of code. Full details of the experiments are included in Appendix C.

We use the perturbed argmax with Gaussian noise in an image classification task on the CIFAR-10 dataset. This serves two purposes: showing that we perform as well as the cross entropy loss, in a case where a soft max can be easily computed, and exhibiting the impact of the algorithmic parameters. We train a vanilla-CNN with 10 network outputs that are the entries of $\theta$, we minimize the Fenchel-Young loss between $\theta_{i}=g_{w}(x_{i})$ and $y_{i}$, with different temperatures $\varepsilon$ and number of perturbations $M$. We observe competitive performance compared to standard losses as baselines (Fig. 2, left and center).

We analyze the impact of the algorithmic parameters on optimization and generalization abilities. We exhibit the final loss and accuracy for different number of perturbations in the doubly stochastic gradient ($M=1,1000$). We highlight the importance of the temperature parameter $\varepsilon$ on the algorithm (see Figure 2, right). Very high or low temperatures degrade the ability to fit to training and to generalize to test data, by lack of smoothing or loss of information about $\theta$. We also observe that our framework is very robust to the choice of $\varepsilon$, demonstrating its adaptivity.

2 Perturbed label ranking

We consider label ranking tasks, where each $y_{i}$ is a label permutation for features $x_{i}$. We minimize the weights of an affine model $g_{w}$ (i.e., $\theta_{i}=g_{w}(x_{i})$) using our perturbed Fenchel-Young loss, a simple squared loss and the recently-proposed blackbox loss of Vlastelica et al., (2019). Note that our loss is convex in $\theta$ and enjoys unbiased gradients, while (Vlastelica et al.,, 2019) uses a non-convex loss with gradient proxies. We use the same 21 datasets as in (Hüllermeier et al.,, 2008; Cheng et al.,, 2009). We report Spearman’s correlation (higher is better) in Figure 3. Results are averaged over 10-fold CV and parameters tuned by 5-fold CV. We find that our loss performs better or similarly (within a $5$% range) on $76$ % and $90$ % of the datasets, respectively. Detailed experimental setup and results are given in Appendix C.2.

To better understand the complexity of this task, we also created a range of artificial datasets where 100 labels are generated by $y_{i}=\operatorname*{arg\,max}_{y}\langle x_{i}^{\top}w_{0}+\sigma Z_{i},y\rangle$, in dimension 50, for different values of $\sigma$. We minimize the same losses as before in $w$. For almost correct labels ($\sigma\approx 0$), our method accurately generalizes to the test data (see Figure 3, and Figure 7 in Appendix C for other metrics). We observe that the Fenchel-Young loss performs as well or better than the other losses, particularly

in terms of robustness to the noise. All details are included in Appendix C.2.

3 Perturbed shortest path

We replicate the experiment of Vlastelica et al., (2019), aiming to learn the travel costs in graphs based on features, given examples of shortest path solutions (see Figure 5). We use a dataset of 10,000 RGB images of size $96\times 96$ illustrating Warcraft terrains of $12\times 12$ 2D grid networks. The responses $y_{i}$ are a shortest path between the top-left and bottom-right corners, for costs hidden to the network, corresponding to the terrain type. They are $12\times 12$ binary matrices representing the vertices along the shortest path.

Following Vlastelica et al., (2019), we train a network whose first five layers are those of ResNet18 for the Fenchel-Young loss between the predicted costs $\theta_{i}=g_{w}(x_{i})$ and the shortest path $y_{i}$. We optimize over $50$ epochs with batches of size $70$, temperature $\varepsilon=1$ and $M=1$ (single perturbation). We are able, only after a few epochs, to generalize very well, and to accurately predict the shortest path on the test data. We compare our method to two baselines, from (Vlastelica et al.,, 2019): training the same network with their proposed blackbox loss and with a squared loss. We show two metrics: perfect accuracy percentage and cost ratio to optimal path (see Figure 6); full implementation details are in Appendix C.3.

Conclusion

Despite a large body of work on perturbations techniques for machine learning, most existing works focused on approximating sampling, log-partitions and expectations under the Gibbs distribution. Together with novel theoretical insights, we propose to use a general perturbation framework to differentiate through, not only a max, but also an argmax, without ad-hoc modification of the underlying solver. In addition, by defining an equivalent regularizer $\Omega$, we show how to construct Fenchel-Young losses and propose a doubly stochastic scheme, enabling learning in various tasks, and validate on experiments its ease of application.

FB’s work was funded in part by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute). FB also acknowledges support from the European Research Council (grant SEQUOIA 724063).

References

Appendix A Proofs of technical results

The function $\varepsilon\Omega$ is the Fenchel dual of $F_{\varepsilon}$ (see Proposition 2.2, impact of the temperature), and is defined on $\mathcal{C}$. As such, as in Abernethy et al., (2014), we have that

It is maximized at $\nabla_{\theta}F_{\varepsilon}(\theta)=y_{\varepsilon}^{*}(\theta)$, by Fenchel-Rockaffelar duality (see, e.g. Wainwright and Jordan,, 2008, Appendix A). ∎

- $F_{\varepsilon}$ is twice differentiable, as a direct consequence of Proposition 3.1.

- $F_{\varepsilon}$ is $R_{\mathcal{C}}$-Lipschitz

$F$ is the maximum of finitely many functions that are $R_{\mathcal{C}}$-Lipschitz. It therefore also satisfies this property. $F_{\varepsilon}$ is an expectation of such functions, therefore it satisfies the same property.

- $F_{\varepsilon}$ is $R_{\mathcal{C}}M_{\mu}/\varepsilon$-gradient Lipschitz.

As a consequence, by the Cauchy–Schwarz inequality, and Lipschitz property of $F$, it holds that

The function $\varepsilon\Omega$ is the Fenchel dual of $F_{\varepsilon}$, which is strictly convex and $R_{\mathcal{C}}M_{\mu}/\varepsilon$ smooth. As a consequence, $\Omega$ is differentiable on the image of $y^{*}_{\varepsilon}$ – the interior of $\mathcal{C}$ – and it is $1/R_{\mathcal{C}}M_{\mu}$-strongly convex.

The regularization function $\Omega$ is differentiable on the interior. If there is a point $y$ of its boundary such that $\nabla_{y}\Omega$ does not diverge when approaching $y$, then taking $\theta$ such that $\theta-\varepsilon\nabla_{y}\Omega(y)\in\mathcal{N}_{\mathcal{C}}(y)$ (where $\mathcal{N}_{\mathcal{C}}(y)$ is the normal cone to $\mathcal{C}$ at $y$), then $y_{\varepsilon}^{*}(\theta)=y$. However, $y^{*}_{\varepsilon}$ takes image in the interior of $\mathcal{C}$ (see immediately below), leading to a contradiction.

- The perturbed maximizer is in the interior of $\mathcal{C}$

Since the distribution of $Z$ has positive density, the probability that $\theta+\varepsilon Z\in\mathcal{N}_{\mathcal{C}}(y)$ (i.e. $p_{\theta}(y)$) is positive for all $y\in\mathcal{E}$. As a consequence, since

with all positive weights $p_{\theta}(y)$, $y^{*}_{\varepsilon}$ is in the interior of the convex hull $\mathcal{C}$ of $\mathcal{E}$.

- The function $y^{*}_{\varepsilon}$ is differentiable, by twice differentiability of $F_{\varepsilon}$, by Proposition 3.1.

Influence of temperature parameter $\varepsilon>0$

Since $y^{*}_{\varepsilon}(\theta)=\nabla_{\theta}F_{\varepsilon}(\theta)$, and since $F_{\varepsilon}(\theta)=\varepsilon F_{1}(\theta/\varepsilon)$, we have $y^{*}_{\varepsilon}(\theta)=y_{1}^{*}(\theta/\varepsilon)$. ∎

We recall that we assume that $\theta$ yields a unique maximum to the linear program on $\mathcal{C}$. This is true almost everywhere, and assumed here for simplicity of the results. We discuss briefly at the end of this proof how this can be painlessly extended to the more general case.

Limit at low temperatures $(\varepsilon\to 0)$

Since $F$ is convex (see proof of Proposition 2.2), so by Jensen’s inequality

Taking expectations on both sides yields that

As a consequence, when $\varepsilon\to 0$, combining these two inequalities yields that $F_{\varepsilon}(\theta)\to F(\theta)$.

Regarding the behavior of the perturbed maximizer $y^{*}_{\varepsilon}(\theta)$, we follow the arguments of (Peyré and Cuturi,, 2019, Proposition 4.1). By Proposition 2.1 and the definition of $y^{*}(\theta)$, we have

Since $\Omega$ is continuous, it is bounded on $\mathcal{C}$, and the right hand term above is bounded by $C\varepsilon$, for some $\varepsilon>0$. As a consequence, when $\varepsilon\to 0$, $\langle y^{*}_{\varepsilon}(\theta),\theta\rangle\to\langle y^{*}(\theta),\theta\rangle$. For any sequence $\varepsilon_{n}\to 0$, the sequence $y_{n}=y^{*}_{\varepsilon_{n}}(\theta)$ is in a compact $\mathcal{C}$. Therefore, it has a subsequence $y_{\varphi(n)}$ that converges to some limit $y_{\infty}\in\mathcal{C}$. However, since $\langle y^{*}_{\varphi(n)},\theta\rangle\to\langle y^{*}(\theta),\theta\rangle$, we have $\langle y_{\infty},\theta\rangle=\langle y^{*}(\theta),\theta\rangle$, by continuity. Since $y^{*}(\theta)$ is a unique maximizer, $y_{\infty}=y^{*}(\theta)$. As a consequence, all convergent subsequences of $y_{n}$ converge to the same limit $y^{*}(\theta)$: it is the unique accumulation point of this sequence. It follows directly that $y_{n}$ converges to $y^{*}(\theta)$, as it lives in a compact set, which yields the desired result.

Limit at high temperatures By Proposition 2.2, $y^{*}_{\varepsilon}(\theta)=y_{1}(\theta/\varepsilon)$, so the desired result follows by continuity of the perturbed maximizer.

Nonasymptotic inequalities. These inequalities follow directly from those proved to establish limits at low temperatures.

If $\theta$ is such that the maximizer is not unique (which occurs only on a set of measure 0), the only result affected is the convergence of $y^{*}_{\varepsilon}(\theta)$ when $\theta\to 0$. Following the same proof of (Peyré and Cuturi,, 2019, Proposition 4.1), it can be shown to converge to the minimizer of $\Omega$ over the set of maximizer. This point is always unique, as the minimizer of a strongly convex function over a convex set. ∎

By the law of large numbers, $\bar{Y}_{n}$ converges to its expectation $y^{*}_{\varepsilon}(\theta_{0})$ a.s. Since $\varepsilon\nabla_{y}\Omega$, the inverse of $\nabla_{\theta}F_{\varepsilon}$, is also continuous (by the fact that $\Omega$ is convex smooth), we have that $\hat{\theta}_{n}$ converges to $\theta_{0}$ a.s.

We write the first order conditions for $\bar{L}_{\varepsilon,n}$ at $\hat{\theta}_{n}$ and the Taylor expansion with Lagrange remainder for all coordinates, one by one

where $A$ is such that, for all coordinates $i\in[d]$

for some $\bar{\theta}^{(i)}\in[\hat{\theta}_{n},\theta_{0}]$. We note here that since the estimator is not necessarily in dimension 1, $A_{n}$ cannot be written directly as $\nabla^{2}_{\theta}\bar{L}_{\varepsilon,n}(\bar{\theta})$ for some $\bar{\theta}\in[\hat{\theta}_{n},\theta_{0}]$, since the Taylor expansion with Lagrange remainder is not true in its multivariate form. However, doing it coordinate-by-coordinate as here allows to circumvent this issue.

We have that $\nabla^{2}\bar{L}_{\varepsilon,n}=\nabla^{2}F_{\varepsilon}$. Since $\hat{\theta}_{n}\to\theta_{0}$ a.s. we have that $\bar{\theta}^{(i)}\to\theta_{0}$ for all $i\in[d]$, so $A_{n}\to\nabla^{2}F_{\varepsilon}(\theta_{0})$ a.s. Rearranging terms in Eq. (7), we have

By the central limit theorem, \sqrt{n}\big{(}\bar{Y}_{n}-y^{*}_{\varepsilon}(\theta_{0})\big{)}\to\mathcal{N}(0,\Sigma_{Y}) in distribution. As a consequence, by convergence of $A_{n}$ and Slutsky’s lemma, we have the convergence in distribution

Appendix B Examples of discrete decision problems as linear programs

On this set, using Gumbel noise yields the log-sum-exp for $F_{\varepsilon}$, the Gibbs distribution for $p_{\theta}$, and the softmax for $y^{*}_{\varepsilon}$. Using other noise distributions for $Z$ will change the model.

Using different reference vectors $v$ yield different perturbed operations, and $v=(1,2,\ldots,d)$ is commonly used.

Assignment. The linear assignment problem, and more generally the optimal transport problem, can also be written as a linear program. In the case of the assignment problem, it is the Birkhoff polytope of doubly-stochastic matrices, whose extreme points are the permutation matrices

There is a large literature on regularization of this problem, with entropic penalty Cuturi, (2013). This is one of the rare cases where the regularized version of the problem is actually computationally lighter, in stark contrast with the general case in our setting.

Combinatorial problems. Many other problems, such in combinatorial optimization can be formulated exactly (e.g. minimum spanning tree, maximum flow), or approximately via convex relaxations (e.g. traveling salesman problem, knapsack), via relaxations in linear programs. Differentiable versions of these exact or approximate solutions can therefore be obtained via perturbation methods.

Relaxations with atomic norms A wide variety of high-dimensional statistical learning problems can be tackled by regularization via atomic, or otherwise sparsity-inducing norms Chandrasekaran et al., (2012); Bach et al., (2012). Our framework also allows us to consider versions of these estimators that are differentiable in their inputs.

Appendix C Experimental details

In the experiment on perturbed maximum for classification on CIFAR-10, we train a vanilla-CNN made of 4 convolutional and 2 fully connected layers for 600 epochs with batches of size 32.

We train by minimizing two losses in the weights $w$ of the network function $g_{w}$, fitting the outputs $\theta_{i}=g_{w}(x_{i})$ to labels $y_{i}$

Perturbed Fenchel-Young (proposed): our proposed Fenchel-Young loss (see Definition 4.1),

Cross entropy loss, for a soft max layer $s_{\varepsilon}$ and an entrywise $\log$

C.2 Perturbed label ranking

In this experiment, we consider label ranking tasks, where each $y_{i}$ is a ground-truth label permutation for features $x_{i}$. We minimize the weights of an affine model $g_{w}$ (i.e., $\theta_{i}=g_{w}(x_{i})$) using the following losses:

Perturbed Fenchel-Young (proposed): our proposed Fenchel-Young loss (see Definition 4.1),

Perturbed $+$ Squared loss (proposed): $\frac{1}{2}\|y_{i}-y_{\varepsilon}^{*}(g_{w}(x_{i}))\|^{2}$, where gradients can be computed using Proposition 3.1 and the chain rule,

Squared loss: $\frac{1}{2}\|y_{i}-g_{w}(x_{i})\|^{2}$,

Blackbox loss: $\frac{1}{2}\|y_{i}-y^{*}(g_{w}(x_{i}))\|^{2}$, where we use the gradient proxy of Vlastelica et al., (2019), re-implemented for the experiments on ranking.

We use the same 21 datasets as in (Hüllermeier et al.,, 2008; Cheng et al.,, 2009). Detailed results are given in Table 1 and Table 2.

For the experiment on artificial datasets, we use the same setup as above, with a linear model instead of affine. The ground-truth vector $w_{0}$ is obtained by uniform sampling in $\{-1,1\}^{d}$, and the $x_{i}$ are standard isotropic normal. In the experiments presented here, we optimize over 2000 epochs, with a batch size of 32: very good results are obtained even with a smaller number of epochs, but we increased it artificially to better evaluate numerically the final predictive performance of all methods (see Figure 8). In the main text, we present in Figure 4 the metric of perfect rank accuracy over one run of simulations. We present in Figure 7 the same metric, as well as the metric of partial rank accuracy (i.e. the proportion of correctly ordered labels), for completeness, averaged over three runs of the dataset. To further illustrate these results, we include in Figure 8, for two fixed values of the noise level, how these metrics evolve through training.

C.3 Perturbed shortest path

In this experiment, we have followed the setup of Vlastelica et al., (2019), to obtain comparable results. We have replicated the network that they use based on Resnet18, and followed their optimization procedure, using Adam with the same learning rate schedule, changing at epochs 30 and 40 out of 50. We also included the baseline that they used, based on training the same network without an optimizer layer. These results are obtained by using the implementation code that they provide.

We minimize the weights of this model for our proposed Fenchel-Young loss (see Definition 4.1).

The perfect accuracy metric measures the percentage of test instances for which an exactly optimal path is recovered, and the cost ratio to optimal metric measures the ratio between the total cost of the path proposed by taking the shortest path for proposed costs $\theta_{i}=g_{w}(x_{i})$ (after training) to the total cost of the path with true costs (see Figure 6).