Stability Regularization for Discrete Representation Learning
Authors: Adeel Pervez
Efstratios Gavves
Venue: ICLR 2022 Poster (2022)

Abstract: We present a method for training neural network models with discrete stochastic variables.
The core of the method is \emph{stability regularization}, which is a regularization procedure based on the idea of noise stability developed in Gaussian isoperimetric theory in the analysis of Gaussian functions.
Stability regularization is method to make the output of continuous functions of Gaussian random variables close to discrete, that is binary or categorical, without the need for significant manual tuning.
The method allows control over the extent to which a Gaussian function's output is close to discrete, thus allowing for continued flow of gradient.
The method can be used standalone or in combination with existing continuous relaxation methods.
We validate the method in a broad range of experiments using discrete variables including neural relational inference, generative modeling, clustering and conditional computing.

Under review as a conference paper at ICLR 2022
STABILITY REGULARIZATION FOR DISCRETE REPRESENTATION LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
We present a method for training neural network models with discrete stochastic variables, The core of the method is stability regularization, which is regu- larization procedure based on the idea of noise stability developed in Gaussian isoperimetric theory in the analysis of Gaussian functions. Stability regularization is a method to make the output of continuous functions of Gaussian random vari- ables close to díscrete, that is binary or categorical, without the need for significant manual tuning. The method allows control over the extent to which a Gaussian function's 's output is close to discrete, thus allowing for continued flow of gradient. The method can be used standalone or in combination with existing continuous relaxation methods. We validate the method in broad range of settings, showing competitive performance against the state-of-the-art.
INTRODUCTION
Neural networks are universal approximators of continuous functions. Often, however, discrete computations are desirable, whether for the intermediate neurons and their representations (Oord al., (017), the parameters (Courbariaux et al.), or the outputs. Current methods for training neural networks require differentiability which means that it is not straightfor- ward to train neural networks with discrete variables. This has led to the development of several approximate methods (Williams) 1992; Jang et al. 2017 Bengio et al. [2013] Tucker jet al. 2017 Pervez al. 2020) with various trade-off's of bias, variance, and complexity. In this work we focus on neural networks with discrete intermediate representations. Building upon techniques from the analysis of functions in Gaussian spaces (Janson et al. 1997), and specifically the notion of sta- bility of Gaussian functions, propose novel regularization strategy on representations that yields precise and hassle-free discrete representations.
Log Probability
Stability Reg.
Gum Softmax
Steps
Figure 1: Training curves in MNIST
structured prediction with ResNet model with a 3x7 dimensional bi- nary latent space. Gumbee--ofttax diverges later in training (T = 0.5).
Several approaches have been introduced in the literature for learning discrete representations with backpropagation. The simplest approach is the Shraight-Througe estimator (Bengio et al., (2013). which essentially ignores the intermediate discrete function allowing the gradients to flow. Another popular choice is the Gumbel-Softmax (Maddison al. 2017 Jang et al.) (017). which replaces the discrete categorical variables with relaxed stochastic continuous ones. In both approaches the discrete variables are replaced by approximations and the model is biased with respect to the original discrete Frraight-Throue
objective. When employed with complex architectures. estimators often underperform due this bias, as in figure П The reason that with continuous relaxation methods there is tension between obtaining better optima and objective function values, and obtaining discrete representations. Importantly, the more complex the optimization (or the model) is the greater well is the pressure towards non-discrete solutions, thus increasing bias further. Adding to the complexity of obtaining discrete representations, with current methods there no direct incentive for the optimization procedure to obtain discrete representations: with Gumbel-Softma the extent how close-to close-to-diseres to-eeerete a representation is obtained is controlled by a temperature variable, which must be manually tuned.
and Gumbel-Softmaxx

---

Under review as a conference paper at ICLR 2022
Unbiased estimators like REINFORCE (Williams 1992) and reduced variance extensions like REBAR (Tucker et al.] 2017) and RELAX (Grthhohhl et al. 2018)) - have also been explored. However, these methods tend to be computationally expensive. which limits their usefulness for complex models. All in all, whether due to bias, high variance, high computational complexity, or the need for manual tuning, there remains need for alternative methods for obtaining hassle-free discrete representations with neural networks, especially when increasing their complexity.
In this work we present regularization procedure for discrete representations, which can be used either as a standalone method, or in combination with existing continuous relaxation or straight- through estimators. In its standalone form, the method replaces discrete variable by a parameterized continuous function whose output (say a sigmoid or softmax function) corresponds to the discrete variable, and which then regularized to produce discrete outputs. In combination with continuous relaxations such as Gumbel-SSoftaa the method can be used to regularize the logits input to the sampling procedure to serve as an implicit temperature control by making the logits noise stable
We achieve this by resting upon the notion of noise stability developed in the analysis of Gaussian functions (Borell, 1985| Mossel & Neeman 2012). Roughly speaking, the noise stability a Gaussian function is a measure of its resilience to noise. Given a Gaussian function / and correlated Gaussian variables e, e' E Rd noise stability of f is defined as Stab = E.,e' [f') Borell's isoperimetric theorem [Borell 1985). states that for bounded functions of some fixed volume with range [0. 1]. noise stability is maximized by functions that are indicator functions of half spaces. Given that half space indicators maximize noise stability in Gaussian space, we suggest that optimizing stability is a very simple and effective method of transforming Gaussian inputs to binary vectors, thus simplifying the process of obtaining discrete representations.
In summary, we demonstrate how the concept of noise stability can be used to regularize stochastic neural networks with Gaussian variables to train hassle-free neural networks with discrete (Bernoulli or categorical) variables.
In the following, we first give a short introduction noise stability in Gaussian analysis. We then motivate our proposal for using noise stability to regularize Gaussian functions for learning discrete representations. We validate by experiments in the Neural Variational Inference framework to learning graph structured latent spaces, learning discrete (deterministic) autoencoders, clustering with Gaussian Mixture VAEs, gating ResNets, and structured prediction.
NOISE STABILITY OF GAUSSIAN FUNCTIONS
2.1
STABILITY AND GAUSSIAN ISOPERIMETRY
Noise Stability of a Gaussian function f R - IR is defined for a noise parameter E 10 as
Stab = ((()
e'= vo -1-
where e' are called p-correlated Gaussian pairs and e, ~ (0, 1). are samples from the standard normal distribution. Stability is defined here in terms of standard normal Gaussian variables, but it is easily extended to any distribution of independent Gaussian variables by reparameterization, given the mean and standard deviation. For the special case where f = LA, the indicator function set A, stability measures the probability that both e and e' remain within A.
Stab,[f] = Pie E AA é' e A],
(3)
By Borell's Gaussian isoperimetric theorem (Borell 1985 Mossel & Neeman 2012), stability is related to the Gaussian isoperimetric inequality. According to Gaussian isoperimetric inequality, geometric objects with minimum boundary (i.e., surface area) in Gaussian space with fixed Gaussian volume (i.e., [[ for f the object's 's indicator function) are half spaces.
Theorem (Borell Isoperimetric Theorem (Borell 1985) ). For fixed p E (0,1) and [R²(
in Gaussian space with range [0, 1] and fixed volume E a, Stab,lf] is maximized by f = 1æ where 1g is an indicator function of a half space with volume a.
As a consequence, given a parameterized bounded Gaussian function, maximizing the stability makes the function / approach the indicator function of some half space illustrated in figure 2
2

---

Under review as a conference paper at ICLR 2022
x²
Figure 2: Increasing stability (left to right) makes a function stable to p-correlated noisy version of as it approaches Gaussian halfspace while keeping fixed Gaussian volume.
For further details on noise stability in the context of Gaussian analysis of functions refer the interested reader to ('Donelll 2014).
3 STABILITY REGULARIZATION
Stability regularization can be used either as standalone method or in combination with Gumbel- Softmax style continuous relaxation.
3.1 REGULARIZATION FOR DISCRETE VARIABLES
We start from a continuous model, say neural network with L layers or modules, f Next, we describe how to employ stability regularization that any arbitrary intermediate function fi(z; 0) with bounded output learns to output discrete variables. Given input = for fi we estimate the stability of and, thereafter, maximize it by adding it to the loss objective
regularizing term.
For a single input vector 2 E R² following the definition of stability in equation ( we sample p-correlated Gaussian variables e, as in equation ( We then evaluate fi twice: once for + and once for + é'. The expectation of their product, [fee( + f))( + +e)] is the stability, [ta),[fi(=)]. Maximizing the stability for single input 2, Stab,[fi(2) for fixed (0,1), the function fi approaches an indicator as described by Borell's 's theorem.
In a batch setting, we compute a Monte Carlo estimate of the expected stability over the input, that is Ez (s)a), by sampling one p-correlated Gaussian pair per batch element. Given 2,6,6' E Bnnkk so that 6, and éf are p-correlated Gaussian, the estimate computed as
1 fiz + +4)f(( TI
Ee[Stab,[fi(t) =
(4)
where n is the batch size and the arithmetic operations are done clement-wise. To maximize stability we sum average the estimate in equation (4) across dimension and add the result an additional regularization term to the loss function with which we train the model.
3.2
STABILITY REGULARIZATION
The regularization makes the function stable relative to correlated Gaussian noise by moving the inputs 2; further apart. In some cases the inputs can become too far separated which can hurt optimization if left uncontrolled, For such problematic cases we introduce mean-centered stability regularization which preserves the expected value of fi given the input z; ensuring that the induced separation remains limited.
The idea behind mean-centered stability regularization is compute the stability of fi E[f{) so that maximizing stability reorients the separating hyperplane without changing the Gaussian volume of the corresponding halfspace. According to Borell's 's theorem maximizing stability of fi for any fixed Eiji causes it to approach Gaussian halfspace indicator (figure 2 Since it can be expensive to compute the expectation of neural network, we maximize the difference two stability computations: Given parameters P1: P2, P2 < 01 we optimize Stab,, [fi] Stab pz [fi]. We show that this objective is equal to Stab p1 -ra ft - [[[]] to first order with an error that is quadratic in Pl, P2.
Proposition 1. Given a Gaussian function f R" + [0, 1] and parameters P1.P2 E (0, 1),
Stab -pa [[ = Stab,, [f] Stab [ (pp2(p2 pi)).
3

---

Under review as a conference paper at ICLR 2022
Algorithm 1 Stability Regularization
Require: Input E <g stability layer fi with range (0, 1)"; noise parameter E (0, 1); stability constraint
(0, 1)
1: Sample e' E R p-correlated Gaussian vectors.
2: Compute y1 = [i(= +), y/2 = [i(z ((2+22 +')
3: Estimate average stability over batch per dimension 5 Σi 11,12225
4: Apply stability constraint per dimension: S = clip(S, ax).
5: Sum S across dimensions and optimize by gradient descent
Algorithm 2 Stability Regularization with Mean Centering
Require: Input z E Ruxk; stability layer fi with range (0, 10"; noise parameter P1, P2 (0,
< P1
Sample (c1, ci) Pr pr-correlated and (c2, <2) P2 po-collllaeed Gaussian vectors from Ba*k fi(z + f2) 6))
2: Compute y1 = fi(z + fi) ()),
3: Estimate average stability over batch per dimension S = E,(,, V2,1).
4: Sum S across dimensions and optimize by gradient descent
See appendix A for proof.
For stability regularization without mean centering we clip the stability at a maximum value to prevent the network output from becoming overly saturated. In figure 2 this would correspond to the points becoming far from the boundary leading to saturation and slowdown of optimization. Without mean centering, constraint stability limits how far the points can be from the boundary improving optimization.
Borell's 's theorem guarantees that the function will converge to halfspace for any p. We did not observe the method to be sensitive to in our experiments.
The precise procedures are described in algorithms 2 and
Stability Regularized Layers. The stability regularized neural network layers be any arbitrary bounded output neural network layer, possibly even a layer of activations without learned parameters. We use sigmoid activations for Bernoulli and softmax for categorical variables.
Probabilistic Models and Gumbel-Softmax. We use stability regularization alongside Gumbel noise in probabilistic models such VAEs where it is important to be able to compute log probabilities of obtained samples. Given a block of layers fi with a Gumbel softmax (or Gumbel sigmoid) GunbelSoffmax(logiss. we compute stability using standard softmax or sigmoid without adding the Gumbel noise as output input to the downstream network.
activation, i.e.,
and use the Gumbel softmax
The optimization procedure with continuous relaxations provides no incentive to encourage discrete representation. The consequence is that such methods work better when there is little pressure from the optimization pressure to be non-discrete, as happens with larger latent space dimension. With smaller bottleneck latent spaces, however, there is greater optimization pressure be continuous and the optimization with continuous relaxations becomes harder because of the need to manually tune the temperature. With stability regularization the regularization procedure form of implicit temperature control and the extent of how discrete a representation becomes is controlled by the extent of regularization.
Computational Complexity. Stability regularization is easy to implement and adds some extra computations due to the extra evaluations for correlated Gaussians. We emphasize that any extra computation is local to the stability layer fi. The rest of the network unaffected. Depending the application, there usually exist only few such layers in a large model, in which case the stability computation is a small fraction of the total cost and we do not observe noticeable increase in computational cost in experiments.
4

---

Under review as a conference paper at ICLR 2022
4 RELATED WORK
A number of methods have appeared in the literature for training neural networks with discrete variables. The methods can broadly divided into two categories: score function based methods, and pathwise gradient methods (Mohamed et al. 2019].
Score function based methods are usually unbiased - the prototypical example being REINFORCE (Williams 1992). The simplest form of these methods has high variance and number of approaches to reduce the variance of the REINFORCE estimator have appeared. NVIL (Mnih & Gregor 2014) subtracts learned MLP baseline from the REINFORCE estimator reduce variance, while MuProp (Gu et al. 2016) subtracts a sample-dependent baseline based on the first order Taylor series expansion. REBAR (Tucker al. (2017) combines the REINFORCE estimator with continuous relaxation baseline, while RELAX (Grathwohl et al.
2018) generalizes the REBAR control variate to one parameterized by neural network. Yin et al. propose ARSM, finite difference estimator with adaptive evaluation.
Pathwise gradient methods make use of the functional form of the operations that are applied to the discrete random samples, and consequentialyy tend to have lower variance than score- core-funntion methods (Mohamed af.] (209). However, since gradients through discrete nodes are not defined, approximations have to be made, which makes such estimates biased in general. Continuous relaxation estimators replace discrete nodes with continuous stochastic variables. Gumbel-Softmax (Maddison) er al. 2017 Jang et al. (017)) is an example of this, which relaxes the Gumbel-Max parameterization of discrete variables. (Potapezynski et al. 2020 extend Gumbel- Softbmlx to Gaussian variables.
Straight-Through estimators (Bengio et al. 2013) are another category of biased estimators, which provide hard samples in the forward pass but skip the discrete node in the backward pass. straight-through estimators can also be combined with continuous relaxation estimators (Jang et al. (017) provide hard samples. In DARN (Gregor et al. 2014) develop another pathwise gradient estimator be unbiased for quadratic func- tions but is shown to be a lower bias Straight- Through estimator in (Pervez et al., 2020). FouST (Pervez et al., (2200) also em- ploys similar correlated samples but in the context of variance reduction Bernoulli-input networks.
0.8
Probability
0.34
Activations
5 EXPERIMENTS
Probability
We validate our stability regularization procedure on number of tasks with Bernoulli and categorical variables to show its ef- fectiveness and wide applicability. First, with a simple example we show that stability regularization can indeed be used to learn Bernoulli random variables with Gaussian functions and that we can control the extent of saturation (or of how close func- tion is to being Boolean) by constraining stability some fixed upper bound. We then demonstrate experiments on 1) latent structure recovery with Neural Relational Inference (Kipf et al., 2018) 2) autoencoders with Bernoulli variables and generative models for images, 3) Gaussian mixture VAEs for unsupervised clustering and 4) conditional computing, where individual lay- ers of a ResNet model can be dynamically switched off to save computation.
Activations
Figure 3: A toy regression exam- ple. The Gaussian function's 's mean output approaches the target proba- bility of 0.8 with and without stabil- ity regularization. The histograms show the outputs clustering around the target without regularization (top) and approaching binary values with regularization (bottom) which average to 0.8.
5.1 A TOY DISCRETE REGRESSION EXAMPLE
As a simple illustration of stability regularization we train a small two-layer neural network with sigmoid output and standard normal input to output binary values. We train the neural network using a regression objective so that the expected output of the network matches a target t, i.c. minimizing (E] 1)2 with and without stability regularization on f. The results for with target of
5

---

Under review as a conference paper at ICLR 2022
Table 1: Categorical VAE experiment with stability regularization showing negative best validation ELBO (best of 5 runs) where is the Gumbel-softtax temperature.
Method
Gumbel-Softmax
Straight-Through
T =
0.1
=.5.
T ==1
T=
1.5
T=
Without Stab Reg.
107 107.4
107.7
114
122.7
132.4
116.7
With Stab Reg.
106.8
112.8
116.8
122.7
111.9
0.8 is shown in figure 3 As shown, without regularization the outputs cluster around the target, as is expected with regression, while with regularization the network outputs binary values with the mean output matching the regression target.
5.2 CATEGORICAL VAE
To judge the effect of stability regularization we perform experiments with categorical VAE on dynamically binarized MNIST. We train the VAE using Gumbel-Softmax and Straight-Thron Through both with and with- out stability regularization. The aim is to show that traditional gradient estimators benefit from stability regularization and that stability regu- larization can be interpreted as performing implicit temperature control when used with continuous relaxation.
Table 2: Regularization
coefficient with categor- ical VAE with Gumbel-
Softmax at 7 = 2
Value
-V. ELBO
80
124.6
We use VAE with the following architecture. The encoder has one layer of 200 units and the ELU non-linearity. The decoder has one layer of 400 units followed by the ELU non- non-linearity. The stochastic layer consists of 100 categorical variables with categories and is implemented as a single layer of 200 units computing the logits for each of the categories. For stability regularization we use exactly the same architecture; in particular do not use separate stability layer and regularization is performed directly on the logits. This is a restriction but has the advantage that the methods remain directly comparable and also that the KL divergences are easily computable because of independence.
100
122.7
120
122.2
150
123.1
No Reg.
132.4
We train all experiments for 500 epochs. The input is binarized in each batch by sampling Bernoulii values independently for each pixel by treating the pixel value as the probability. We train using Adam with a learning rate of le-4. For Gumbel-Softman we use temperatures in {0.1, 0.5, 1.0). The stability regularization coefficients are chosen from {20, 50, 100). The parameter p is set to 0.9. Frraigh-TThroug is implemented by sampling 1-hot vectors according to logits and passing the gradient to the softmax (instead of the logits) the backward pass.
Comparisons with Gumbel-Softmax and Straight-Through with and without stability regularization are shown in Table [ We see that stability regularization improves the Frraight-Througe trained model significantly (~~5 nats) and the Gumbel-Softmax trained models when the temperature is higher (~ 10 nats at r = 2). This shows evidence for our suggestion that stability regularization provides implicit temperature control with continuous relaxation.
To study the effect of the regularization term coefficient on training we show validation ELBO for cat- egorical VAE trained with Gumbel-softmax with fixed temperature set to 2 and stability regularization with various coefficient values in Table 2 Coefficient values are chosen from {80, 100, 120, 150]- The table shows that training performance is relatively stable (within +1 nats) across range of coefficient values. Compared with the range of ELBO values obtained when changing the temperature in Table I we conclude that the action of the regularization coefficient is different from that of the temperature.
5.3 GRAPH STRUCTURED LATENT SPACES
We perform experiments with Neural Relational Inference (Kipf et al., 2018), which a variational autoencoder model to recover interactions and learn dynamics given observational data. The latent space of the NRI VAE model is graph structured, and the encoder and decoder are graph neural networks. To obtain sample interaction graphs in the latent space, NRI models Gumbei-Softmax to sample edges producing a graph sample and to propagate gradients through the latent graph. For the NRI experiments, we stability regularization along with Suftee-Softtax.

---

Under review as a conference paper at ICLR 2022
Physies Simulation Experiments The first set of NRI experiments are physics simulations from the original NRI proposal Kipf et al. (2018). This has three types of systems of T) particles connected by springs, 2) charged particles and 3) phase- coupled oscillators with the Kuramoto model [Kuramoto| 1975). Synthetic data was generated using the authors' code with 50k training and 10k validation and test samples each. For each type of system we have further two types of experiments with either or 10 objects.
Table 3: NRI Physics Simulation Accuracy. indicates result different from number reported.
Model
Springs
Charged
Kuramoto
5 Objects
Supervised
99.9+0.0
95.0+u.3
99.7=0.0
NRI-GS NRI-GS+Stab
99. 99.9
82. 20.6
93,9*
99.9+0.0
88.1 +8.2
95.3+0.1
10 Objects
Supervised
98.8+0.0
94.6+0.2
97.1=0.1
NRI-GS NRI-GS+Stab
98.4+0.0
70.8+0.4
66.5*
98.4+0.0
75.0
68.7=0.6
We compare the NRI physics experiments with and without stability regularization.
For this experiment used the same baseline model with new parameters and the stability regularization is used with a softmax applied to the logit parameters going into the Gumbel-Softmax. For this experiment used stability regularization without mean centering, p value of 0.9, and constraint of 0.9 for the sum of stability across the softmax dimensions. We add the stability loss to the optimization objective with a multiplied by factor of 100.
For the 5 and 10 particles springs experiments, we match the NRI baseline, which already matches the supervised baseline. For the 5 and 10 charged particles experiments we get a significant im- provement with stability regularization, achieving an accuracy improvement of about 6% and 5%. For the Kuramoto oscillator experiments, we achieve accuracy improvements of 1.5% to 2% over and 10 oscillator experiments. Here we note that the orig- inally reported accuracy values for the Kuramoto model are higher (96% and 75.7%) but we were un- able to reproduce the numbers in our experiments. The Kuramoto values we report were obtained by running the provided code against which com- pare our method. We note that the improvements are obtained solely by including stability regularization without model changes.
Table 4: Comparing with Paulus et al. 2022 in latent spanning tree recovery with 3,5 and 10 steps, w.r.t. test ELBO, precision, recall.
Method
ELBO
Pr
Re
T=10
Top|V|-1 Topp V
-2100
41
41
-1766
92
92
Sp
-1080
91
91
SpTree+Stab
-1175
89.3
89.3
T=5
Latent Tree Structure Recovery We perform fur- ther experiments with NRI on latent spanning tree recovery proposed by Paulus et al. (202T). Here a dataset is generated by sampling tree over 10 ver- tices, embedding the vertices and applying a force directed graph drawing algorithm (Fruchterman & Reingold 1991). The dataset consists of particle lo- cations at each step obtained by applying the drawing algorithm for T steps. Paulus al. 12021) develop a NRI method for sampling tree structured latent spaces and show improved performance with structure recovery compared with sampling independent edges as in the baseline NRI model and Top [V| 1 selection where V is the vertex set. Their conclusion is that latent tree sampling works better than the baseline for T=10 iterations.
SpTree Sparee+Stab
-516
82.3
82.3
-501
82.6
82.6
T=3
SpTree SpTree+Stab
-221
65.9
65.9
-196
70.52
70.52
Using the authors' code we performed experiments with shorter T=3,5 and T=10 step trajectories with the spanning tree and Top V - sampling with and without stability regularization. recovering the interaction structure is harder when shorter trajectories observed. We used mean centered stability regularization with a one layer stability network, which we took as the last linear layer of the encoder network. We used A1 0.9, P2 (0.89, 0.8, 0.5). We ran random search over the Gumbel- Softmaxx temperature and learning rate and ran the experiments for multiple random seeds for the best hyperparameters found by the search.
general,
We first focus on the spanning tree methods, which resemble the true latent structure of the data. Including stability regularization matches performance T=5,10 steps, and improve with T=3 steps
7

---

Under review as a conference paper at ICLR 2022
from around 66% 70%. What is more interesting, when focusing the Top [V| method for T= 10, we find that having stability regularization we can match the spanning tree method with over 90% precision and recall. This is remarkable considering that the latent space is that of a spanning tree, and that using a Top [V] - without regularization from Paulus et al. (2021) scores a 41 % accuracy. When reducing the number of steps, however, the Top V does not have the capacity to infer the latent structure, attaining only a low precision and recall of 35% for the best validation score,
For this experiment we find that stability regularization made the method significantly less sensitive hyperparameters and random seed. We conclude that with stability regularization helps in the low data limit for this task or when the latent space has less structure than the ground truth alongside making the method less sensitive to hyperparameters and random seed.
5.4 UNSUPERVISED CLUSTERING WITH GAUSSIAN MIXTURE VAES
To show an application of stability regularization categorical variables we turn to Gaussian mixture VAEs and use them for unsupervised clustering for MNIST and OMNIGLOT. We of the Gaussian mixture components are functions of a sampled categorical variable. In particular, we do not marginalize over the categorical variable and resort to sampling only.
very simple generative model where the mean and variance
Table 5: Unsupervised classification accuracy for MNIST
b ~ p(b) Categorical(1)
~ 6) = (Ne, (b), (b)) ~~ >=p(z)))
Model
Accuracy
AAE (Makhzani et al.)
95.9
where ,e,, To, are MLP networks with categor- ical input. We use the following approximate posterior.
GMVAE (Dilottanakkl et al.
92.7
SB-VAE (allsnnik & Smyth)
92.35
Stab. Reg. (20 Clusters)
93.3
~~ (2/[x) N(uo, (o,)) b p~ (b|2) = Categorical( (z))
The networks in p(x])). are ResNets with two ResNet blocks, while Jos is implemented as an MLP. The categorical variable is implemented as softmax at the output of a stability network in Jøa which we regularize to produce categorical outputs. Computing KL divergence requires probabilities for the categories, which approximate using separate two-layer MLP with softmax output, trained to regress over the categorical outputs. To train the model we maximize the ELBO plus the auxiliary stability and regression objectives. We do not use Gumbel Softmax and use standalone stability regularization for the categorical variables.
We obtain average unsupervised classification accuracy MNIST of 93.3% (see table 5). Label assignment is done using the label of the maximum probability example in a cluster. We observe that while relying simple model, we obtain classification accuracy that is better than close to more complex models from the state-of-the-arr. We provide generation examples from 15 clusters for MNIST and OMNIGLOT in figure in the appendix.
5.5 GATING RESNETS
As an application of Boolean variables to gating, we consider the task of selectively turning off layers of a deep ResNet using gates. We chose the adaptive inference network architecture, which builds on ResNet- 1 10 architecture for CIFAR-10 (Veit & Belongie 2017). The architecture incorporates Boolean variables to serve as gates, which are trained using Gumbel-Softtax The architecture also allows setting target rate at which individual layers can be turned off.
We use the authors' PyTorch implementation and replace the Gumbel-Sofinee dalone stability regularization. We use the average activation over the batch probability estimate in algorithm [] and stop stability gradients at the stability block input. We use the default settings of parameters found in the implementation. With target rates of 0.5 and 0.6 we were able to obtain
training with stan-
8

---

Under review as a conference paper at ICLR 2022
accuracy values of 92.86 and 93.31 for the baseline implementation with Gumbel-Softmax training. For our stability regularization method we obtain accuracy values of 93.2 and 93.46 with the rate set to 0.5 and 0.6. While the improvements crease hard. With additional tweaks, Veit & Belongie (2017) report a higher accuracy of 94.24. In the above experiments, however, we use the published implementation with the default hyperparameters. We conclude that stability regularization can also be deployed for dynamically switchable layers in deep architectures.
modest, the already high accuracies makes any further in-
5.6 DISCRETE
AUTOENCODERS AND GENERATIVE MODELING FOR IMAGES
As an application of stability regularization, we train discrete autoen- coders natural images solely with binary latent variables, where the encoder and decoder are deep ResNets. We use a Gaussian output model for the decoder. After training the discrete autoencoder we train a PixelCNN prior over the binary latent variables to get generative model. This similar to the setup used by Dord et al. (2017). with the difference that we do not quantize, instead we use standalone stability regularization to obtain a discrete latent space.
To generate new samples, we first sample a binary latent variable vector from the PixeICNN prior model, which is then fed to the decoder to generate samples. We train the model on 100 with 50k 32x32 labeled images and STL-10 with 100k 96x96 unlabeled images. For CIFAR-10 10 and STL-10 use binary latent spaces of dimensions 16x16x4 and 24x24x4 respectively. This corresponds to a bit-size 24 times for CIFAR-10 and 96 times for 16x 16x4
reduction of 32x32x3x8
STL-10.
Sampled images from the models trained on CIFAR-10 and STL-10 are shown in figure 4 and B.1 in the appendix. We also train conditional PixcICNN prior model for CIFAR- 10 by using the CIFAR- 10 training labels. Samples from the conditional model are shown in figure 4 The generations show that the model is able to capture global information better than PixelCNN model trained on raw images. For the CIFAR- 10 unconditional model we obtain an FID score of 64.3 which is a slight improvement over the value reported by Ostrovski et al. for PixelCNN model trained on raw images. For the conditional CIFAR- 10 model we obtain an FID score of 54.3. A comparison of FID scores for various models on CIFAR-10 10 is in table the appendix. The FID scores were obtained on a highly compressed representation with a standard PixelCNN. The scores can likely be improved by using lower compression or an improved PixcICNN but we do not explore this in this work.
Figure 4: CIFAR-10 recon- structions with a 16x16x4 binary latent space (top); generated samples with a pixeICNN prior on the la- tent space per class (mid- dle); and conditional sam- ples from a single class (bot- tom)
6 CONCLUSION
We presented a new method for regularizing stochastic Gaussian neural networks, so that to train accurately and hassle-free models with Boolean and categorical stochastic variables. For this, we rely upon the notion of noise stability developed in the analysis of Gaussian functions, which is maximized by functions that are indicator functions of half spaces. We validate successfully stability regularization on a wide array of experiments and settings, where Boolean and categorical random variables are required, including physics simulations with graph latent variable models, Gaussian mixture models for clustering, gating applications for large neural networks, and autoencoders. Importantly, we find that stability regularization requires limited tuning compare to other continuous relaxation methods, making it strong contender for models with discrete variables in practice.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Jyoti Aneja, Alexander Schwing, Jan Kautz, and Arash Vahdat. NCP-VAE: Variational autoencoders with noise contrastive priors. URL /ttp:://:/arx www.wiwww..a....aaa.aoooooooooee wwt. 200lg wwtp////..////////0..ooo. 02917
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pp. 214-223. PMLR, 2017. ISSN: 2640-3498.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: vrXiiv:308...0 /308.3432, 2013.
Christer Borell. Geometric bounds on the ornstein-uhlenbes velocity process. 70(1):1-13, 1985. Publisher: Springer.
Matthieu Courbariaux, Itay Hubara, Daniel Soudry, Ran EI- El-iv,,v, and Yoshua Bengio. Binarized neural networks: Training deep neural networks with weights and activations constrained to +1 or -1. URL BRRLhhtp://axiiii org/ aww//10000 02830
Nat Dilokthanakul, Pedro A. M. Mediano, Marta Gamelo, Matthew C. H. Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture variational autoencoders. Jttttt/:///axxi http://p///ax....oorr/////////// 02648
Thomas MJ Fruchterman and Edward M. Reingold. Graph drawing by force-directed placement. 21 (11):1129-11) 1991. Publisher: Wiley Online Library.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoff Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In ICLR 2018: International Conference Learning Representations 2018, 2018.
Karol Gregor, Ivo Danihelka, Andriy Mnih, Charles Blundell, and Daan Wierstra. Deep AutoRegres- sive networks. In /CML, 2014.
Shixiang Gu, Sergey Levine, Ilya Sutskever, and Andriy Mnih. MuProp: Unbiased backpropagation for stochastic neural networks. ICLR (Poster), 2016.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel- In ICLR 2017 International Conference Learning Representations 2017, 2017.
Svante Janson et al. Gaussian hilbert spaces, volume 129. Cambridge university press, 1997.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. 2018.
04687
Yoshiki Kuramoto. Self-entrainment of a population of coupled non-linear oscillators. In International symposium on mathematical problems in theoretical physics, pp. 420-422. Springer, 1975.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous relaxation of discrete random variables. In ICLR 2017 International Conference Learning Representations 2017, 2017.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In Proceedings of The 31 International Conference on Machine Learning, pp. 1791-1799, 2014.
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo Gradient Estimation in Machine Learning. orXiv:9906 10652 [cs, math, stat], June 2019. URL http:
http://bb/19900 1 10652
arXiv:
Elchanan Mossel and Joe Neeman. Robust optimality of gaussian noise stability. 2012. URL https::/arxiv.orgg// atwp//g/ab////// 4126v3
Eric Nalisnick and Padhraic Smyth. Stick- boeck-breakigg variational autoencoders. URL attt:://arxlt
arg/abs/1660 06197
10

---

Under review as a conference paper at ICLR 2022
Ryan O' O'Donnell. Analysis of boolean functions. Cambridge University Press, 2014.
Aäron den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel Recurrent Neural Networks. In ICML, 2016.
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural Discrete Representation Learning. NIPS, 2017.
Georg Ostrovski, Will Dabney, and Rémi Munos. Autoregressive quantile networks for generative modeling. URL http: org. abs 111.//866.000 05575
Max B. Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris J. Maddison. Gradient estimation with stochastic softmax tricks. 2021. URL
08063
Adeel Pervez, Taco Cohen, and Efstratios Gavves. Low bias low variance gradient estimates for boolean stochastic networks. ICML, 2020.
Andres Potapczynski, Gabriel Loaiza-Ganent and John P Cunningham. Invertible gaussian reparam- eterization: Revisiting the gumbel-softmax. 2020.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupprriseddeeppreentttiio learning with deep con- volutional generative adversarial networks. htll:/tcce//arrrv....r///////// 06434
Tapani Raiko, Mathias Berglund, Guillaume Alain, and Laurent Dinh. Techniques for learning binary stochastic feedforward neural networks. arXiv preprint arXiv:1406.2989, 2014.
George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems, pp. 2627-2636. 2017.
Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. 2020. URL
Andreas Veit and Serge Belongie. Convolutional networks with adaptive inference graphs. 2017.
Ronald J Williams. Simple statistical fralient-following algorithms for connectionist reinforcement learning. Machine learning, (3344)229-2550 1992.
Mingzhang Yin, Yuguang Yue, and Mingyuan Zhou. ARSM: Saaeeem--RENFOOE--Swaaea
estimator for gradient backpropagation through categorical variables. pp. 10.

---

Under review as a conference paper at ICLR 2022
A PROOF OF PROPOSITION 1
Proof. We use the following expression of stability in terms of the Hermite expansion of f.
Stabp[f] = Edla
(5)
where a is multi-index indexing the Hermite polynomial basis has and f(a) is the coefficient of in the Hermite expansion of f corresponding to the basis function ha (O'Donnell 2014).
Given parameters Pl, P2 P2 < Pt, we have
Stab ptabpi-pin Pf B[S]]=
E (p1 (21-1))lo f(a))
(6)
a: a/20
where we use the fact that E[ f(0). Also
Stab [f]
Stab pz [f]
Σ (p) P2 (
(7)
a:/a|20
Stab -pe E[fl] + Opp2(p2 p1))
(8)
B FURTHER EXPERIMENTAL DETAILSS
B.1 DISCRETE AUTORNCODERS
For CIFAR-10 we downsample once using strided convolutions the encoder. After the downsam- pling the encoder has 3 ResNet blocks with 80 feature maps. This is followed by the stability network which has two convolutions layers followed by two ResNet blocks all with 80 features maps followed by convolutional layer with 4 features maps. We upsample using transposed convolutions in the decoder followed by ResNets blocks.
For STL-10 we downsample twice using strided convolutions in the encoder. After the first downsam- pling we use one convolutional layer followed by a ResNet block. After the second downsampling we use 6 ResNet block. This is followed by the stability network which has two convolutions layers followed by two ResNet blocks all with 80 features maps followed by convolutional layer with features maps. We use 3 ResNet blocks with 80 feature maps. The first upsampling operation in the decoder followed by ResNets blocks with 80 features maps. The second upsampling operation in decoder followed by ResNets blocks with 64 features maps.
Table 6: FID scores for CIFAR-10
Model
FID
PixelCNN Ostrovski et al.)
Oord al.
2016
65.93
PixelIQN OOstrovski et al.)
49.46
DCGAN (Rafford al.
Arjovsky et al. 2017
37.11
NVAE
Vahdat & Kautz
2020 Aneja et al.)
51.10
Discrete AE+Stab. Reg.
64.39
Discrete AE+Stab. Reg. (cond.)
54.40
The output model in both cases Gaussian.
We E {0.8. 0.9, 0.95) and stability constraint of 0.6. To evaluate threshold the output of the stability network at 0.5 so that greater values become and the rest O. We train with Adam with a learning rate of 8e-5.
12

---

Under review as a conference paper at ICLR 2022
Figure 5: STL-10 reconstructions with a 24x24x4 binary latent space (left); generated samples with pixelCNN prior on the latent space (right).
B.2 STRUCTURED PREDICTION
As further validation of stability regularization to train models with Boolean variables we experiment with the MNIST structured prediction task where the goal generate half of the image given the other half. This task was used in Raiko et al. (2014); Tucker et al. (2017) to test and compare various gradient estimators for Boolean latent variables. We use the dynamically binarized dataset for this task. We use an MLP architecture we layers of 200 tanh units in the encoder and decoder and similarly stability block with layers of 200 tanh units. The output of the encoder fed the stability block and the output of the stability block goes into the encoder. To evaluate we threshold output of the stability block so that the decoder only sees Boolean values during evaluation. We compare validation log likelihood curves with single sample against Gumbel-Softmax, Rebar and MuProp for which we use code from the Rebar code repository. Since our model has more parameters due to the stability block, to make the comparison fairer. We show validation curves in figure
encoder and decoder architectures of 250 tanh units for these models
Log Probability
-04
Stability Reg.
-66
Gumbel-Softmax
Rebar
Dynamic Rebar
MuProp
0.6
122
Steps
Figure 6: Validation curves for structured prediction on dynamic MNIST We find that, terms of validation likelihood, Rebar performs worse than dynamic MNIST dataset on this task while stability regularization outperforms all others by large margin.
with the
B.3 FURTHER ABLATIONS
B.3.1 NOISE PARAMETER
We study the effect of the noise parameter p on VAE performance. Using the same setup in Section
5.2 and using Humbel-Softman with a fixed temperature of 1, we train models with parameters
0..5 0.6, 0.8, 0.9). We run the experiment for 300 epochs instead of 500 as in Section 5.2 The results shown in Table 7 indicate that the performance minimally affected by varying choices of all runs achieving validation BLBO of about 13.8.
13

---

Under review as a conference paper at ICLR 2022
Table Negative validation ELBO when varying p with categorical VAE with Gumbel-Softmax at Here study the effect of using Gumbel Noise also in the stability computation. We use the same setup in Section 5.2 and use stability regularization with Gumbel noise for various temperatures. The results are shown Table Here we find that for stability computation with Gumbel noise we get an improvement of about 2 nats for higher temperatures we get on-par or marginally worse performance. We conclude that under specific conditions Gumbel noise in regularization can lead to small benefit, but the overall effect is marginal.
-V. ELBO
0.5
113.8
0.6
113.9
0.8
113.8
0.9
113.7
B.3.2 GUMBEL NOISE IN REGULARIZATION
When computing the stability regularization objective alongside Gumbel-Softmax we evaluate the stability layer without Gumbel noise in the experiments, using plain softmax the output of the stability layer. Gumbel noise still used for the input to the downstream decoder network.
1.5, 2. Whereas for lower temperatures
Table 8: Categorical VAE experiment with stability regularization using Gumbel noise for stability. Showing negative validation ELBO.
Method
Gumbel-Softmax
0.1
== 0.5
T
T 1.5
Without Stab Reg.
107
107.7
114
122.7
132.4
Stab Reg. Stab Reg.+Gumbel Noise
107.4
106.8
112.8
116.8
122.7
107.4
107.3
112.9
114.4
120.1
4
9
55
1 X
17 X3
9
9
47
0
0
0
0
0
D
0
0
6
A
14
4
2
2
2
2
2
2
2
2
37
2
2 8 333
1
JH
4
IR
X
8
8
8
8
9
B
3
8
8
R
15
JD
T
SE
3
5
3
3
3
3
3
3
3
3
3
T
3
9
4
67
9
4
9
4
4
9
9
4
3
9
8
7
7
7
9
7
9
8
7 2 1 4
4
6
47
2
2
2
2
2
2
2
2
2
2
2
8
8
8
8
8
2
8
2
8
8
4
//
4
4
4
4
4
4
4
6
-
6
6
6
6.
1
1
1
1
1
1
1 000 9
0
0
0
0
0
0
0
7
9
7
5
5
5
5
5
5
5
3
5
S
5 5
Figure 7: Unsupervised clustering with Gaussian Mixture VAE trained with stability regularization.
The images show 15 clusters (one per row) and corresponding generated images for MNIST (left) and OMNIGLOT (right).
14


Reviews:
[{"review_id": "v-dXsHU4xes", "reviewer": "Reviewer_k6BN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think the proposed new method is interesting and worth further study, but I also feel that the contribution of the submission is somewhat limited. \n\nIn particular, I have several concerns about the submission. \n\n1. I think the study is not thorough enough. Though we see some performance improvement, can we compare the proposed method against Gumbel-softmax at a more detailed level. For example, Gumbel-softmax uses temperature to control the extremeness of function values. I think the corresponding hyper-parameter in the proposed method is the strength of regularization. Can we do a side-by-side comparison between the two hyper-parameters. For example, checking the performance vs temperature, and the performance vs regularization strength. \n\nMy hypothesis is that the temperature value is hard to tune while the regularization strength is easier to tune. The performance difference might be from this reason. \n\nIn general, I think the optimization involving discrete variables is hard, and the backpropagation through these variables does not behave. I don't feel the proposed method is changing the property. \n\n2. The performance improvement does not seem to be consistent. Sometimes we see performance drops. Is the experiment in 5.3 a good comparison? The proposed method uses 20 clusters while the first two baselines use 30 clusters. More clusters should leads to higher performance, right? The authors may want to increase the number of clusters and check the performance. How do you get the performance of SB-VAE? I don't find it in the original paper. \n\n3. Which experiments use mean-centering? Do you see a performance improvement by using mean-centering? I don't see a reason for mean centering in experiments. I feel that the model will learn the function and decide an underlying mean even without mean-centering. \n\n\n\n\n", "summary_of_the_paper": "The paper proposes an interesting regularization method for encouraging function values to be discrete. The main idea is very interesting: with two correlated random vectors, the function is encouraged to maximize the correlation of outputs from the two vectors. The function that maximizes the correlation is the indicator function of a half-space. Combining this objective and original training objectives, a learning model can output random values that are near 0 or 1. The model plays the same role as the Gumbel-softmax trick. \n\nThe submission has done extensive experiments comparing the proposed technique against Gumbel-softmax. The results indicate that the new technique improves the performance in several learning tasks, though the improvement is not consistent. ", "summary_of_the_review": "Overall I think the technique is interesting, but I also feel the overall contribution of the submission is somewhat limited. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "N-eLdLHpR3L", "reviewer": "Reviewer_btCU", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work introduces stability regularization for training neural network with discrete stochastic variables. Unlike most of existing methods for handling discrete stochastic variables, the proposed method promotes the output of continuous functions of Gaussian random variables to discrete, which leverages the nice property of the Gaussian noise stability. The authors have conducted various experiments on different benchmarks to showcase its superiority over existing methods.\n\nIt has been known that halfspaces maximize noise stability subject to a measure constraint since Borell's work in 1985. But the theorem only states about the relationship between noise stability and partitions with two parts. This can be directly related to the binary variables, but it is not obvious how this could be extended to the categorical variables (that is, partitions with more than two parts). The way that the authors handle this is to sum across dimension, which is more like treating each dimension as independent binary variables rather than a single categorical variable. This does not prevent the output from generating all zeros or multiple ones, which is no longer a categorical variable. It has been shown that the standard simplex partition fails to maximize multi-part noise stability unless all of the parts have equal measure [1, 2]. Therefore, the analogue of Borell's theorem in the multi-part partitions case seems still an open problem. The authors might want to discuss more regarding the way that they extend the two-parts result to the categorical case and how the proposed regularization (4) promotes a single one for the categorical variable.\n\nThe authors present two versions of stability regularization: one with mean centering and the other without centering. It seems natural to compare these two versions quantitatively and give recommendations for situations where one is better than the other. However, in all 6 experiments, the authors only choose to provide results for either one of the two versions without any direct comparisons between them and the choice of which one to apply in each experiment appears a bit arbitrary. On the other hand, there are a few hyper-parameters introduced for the regularization: the correlation variable, the maximum stability, and the regularization strength. The authors claim that these either do not impact the performance much or are easier to tune without any quantitative results. The most ideal case is that one set of hyperparameters can be applied to multiple problems and gives reasonably good performance for all of them. But this does not seem to be the case here. The maximum stabilities are different across applications, and the regularization strength is only mentioned for the neural relational inference experiments. Since the \"hassle-free\" is one of the most important features of the proposed method, it is critical to provide sensitivity analysis of the introduced parameters and demonstrate that competitive performance can be obtained without careful tuning.\n\nThe authors also discuss how the proposed regularization can be combined with the Gumbel-Softmax method, where the output is from the Gumbel-Softmax computation but the stability is computed with the vanilla softmax without adding the Gumbel noise. This creates a mismatch here: the regularization is not applied on the \"real\" output. The authors might want to provide the rationale behind this and why the Gumbel noise is necessary to produce the output as input for the downstream network. One benefit of combining the stability regularization and the continuous relaxations, as the authors argue, is to provide \"a form of implicit temperature control\" and mitigate the need to manually tune the temperature. There are quantitative results supporting the combination is better than using Gumbel-Softmax alone. But it would be more interesting to see either the less sensitive to the temperature with the combination, or the universal improvement of the combination over the Gumbel-Softmax alone for different temperatures.\n\nIt is really nice that the authors conduct various experiments across different problem domains, but they are mostly in the same flavor without providing complementary support. I would suggest move some of the experiments to the supplementary material. The saved space can be used for the sensitivity analysis mentioned above, and the detailed comparisons with existing methods handling stochastic neural networks similar to the experiments done in [3] (the same toy experiment) and the Section B.2 in the supplementary material. After all, this is a work of training stochastic neural network. It would be nicer to give a thorough comparison with methods discussed in the Related Work.\n\nSome typos:\n- In Algorithm 2, $v_1$ and $v_2$ are used in Line 2, but $y_1, y_2$ are used in Line 3.\n- The error of the mean-centered stability approximation in Proposition 1 is $O(\\rho_2(\\rho_2 - \\rho_1))$, but $O(\\rho_1 (\\rho_1 - \\rho_2))$ in Section A in the supplementary material.\n\n[1] S. Heilman, E. Mossel, and J. Neeman. Standard simplices and pluralities are not the most noise stable. ITCS 2015.\n\n[2] A. De, E. Mossel, and J. Neeman. Noise stability is computatble and approximately low-dimensional. Theory of Computing, 2019.\n\n[3] A. Pervez, T. Cohen, and E. Gavves. Low bias low variance gradient estimates for boolean stochastic networks. ICML 2020.", "summary_of_the_paper": "The paper proposes a method for regularizing neural networks with boolean and categorical stochastic variables. The proposed stability regularization is based on the Gaussian noise stability, which is maximized by the indicator functions of half spaces. Experiments with various benchmarks show promising results comparing to existing methods.", "summary_of_the_review": "The proposed stability regularization has nice theoretical foundation and shows good performance across different problems involving discrete stochastic variables. But I am lean to reject mainly due to the following reasons: 1) the lack of theoretical foundation for the extension to the categorical variable, 2) the insufficient quantitative results supporting the \"hassle-free\" characteristic of the proposed method.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Bb1VL9QpJBl", "reviewer": "Reviewer_ZtpK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper builds on the concept of noise stability of functions of Gaussian variables.  Noise stability measures the resilience of such functions \nto noise by estimating how their outputs correlate when they receive as input correlated Gaussian variables; the higher the correlation the more \nstable the function. For a family of functions, bounded and fix volume, the stability is maximized for these functions that are indicator functions \nof half spaces which is a natural way to construct binary variables. The paper proposes to maximize stability in order to obtain categorical \nvariables.\n\nThe implementation of the method is rather simple and it comes with a small computational overhead due to the fact that it requires a double evaluation\nof these computational units whose output we would want to push to be discrete. \n\nThe paper explores the performance of the method in a diverse set of experiments learning discrete latent spaces, gating computational units using binary\nvariables, and generating discrete outputs. I appreciate the breadth of scenarios, I only have a small comment here on the fact that the comparisons are \nonly done with the Gumbel-softmax, why this choice and not compare also with some of the other approaches that the authors discuss, e.g. straigh-through\nand/or some variants of the score based estimators?\n\n", "summary_of_the_paper": "The paper presents a new method to train neural networks with stochastic discrete variables called stability regularisation. The method pushes \nthe outputs of functions of Gaussian random variables to be close to discrete, and unlike other methods used with discrete variables, \nsuch as Gumbel Softmax which requires temperature annealing, it is easy to tune. The method is demonstrated on a very rich variety of\ntasks and models and shows state of the art performance. \n", "summary_of_the_review": "Overall this is a nice contribution to what is an important problem, learning discrete (stochastic) variables. The idea is simple and the authors demonstrate it in a variety of datasets. The only thing that I would have liked to see is a comparison with other methods that are meant to also work in such discrete settings, e.g. straight-through, score-based estimators. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4JuZn73KquE", "reviewer": "Reviewer_jxs7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\nThe method implemented in the paper is innovative and exploits a nice and not-so-popular property of the Gaussian distributions. Connecting real and discrete outputs through the expectation of the indicator function is a promising idea and may have a good impact on various discrete optimization applications. The paper has an extensive set of experiments.\n\nweaknesses:\nI have two main concerns about this work. One is the computational cost of the proposed method. Probabilistic approaches to discrete optimization often suffer from high computational complexity, e.g. the likelihood is intractable. The proposed procedure is based on a different idea but it looks like the resulting optimization requires Monte Carlo sampling or ELBO approximations. The second concern is about a possible comparison with other methods for discrete optimization, e.g. STE or more recent versions of it. If the authors have some available results, these could have been put in more evidence in the experiments section.\n\nquestions:\n- Would it be possible to use the method for learning a binary NN where only the weights and not the layer representations are discrete?\n- Is this the first time Gaussian stability is used for learning approximations of discrete-valued functions?\n- Could be the proposed method compared with other approaches for discrete optimization, e.g. STE, on both the quality of the output and the computational cost?\n- How many artificial Gaussian variables are needed for a d-dimensional optimization? Is the procedure expected to be scalable for networks with millions of parameters?", "summary_of_the_paper": "The paper describes a computational approach to discrete optimisation, based on a probabilistic regularization procedure that enforces the output of a continuous function to be quasi-discrete. The idea is to use a property of Gaussian noise described by Borell's Isoperimetric Theorem.\n", "summary_of_the_review": "A very nice idea, probably a bit expansive on the computational side.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a method to train neural networks based on so-called stability regularisation. The method encourages the outputs of functions of Gaussian random variables to be close to discrete and does not require temperature annealing like the Gumbel Softmax. All reviewers agreed that the proposed method was novel and of interest. The authors conducted extensive experiments. They also adequately addressed the concerns raised by the reviewers (e.g., theoretical foundation and computational cost).
Decision: Accept (Poster)