Stability Regularization for Discrete Representation Learning

1 INTRODUCTION Neural networks are universal approximators of continuous functions . Often , however , discrete computations are desirable , whether for the intermediate neurons and their representations ( Oord et al. , 2017 ) , the parameters ( Courbariaux et al . ) , or the outputs . Current methods for training neural networks require differentiability which means that it is not straightforward 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 et al. , 2017 ; Pervez et al. , 2020 ) with various trade-offs 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 stability of Gaussian functions , we propose a novel regularization strategy on representations that yields precise and hassle-free discrete representations . Several approaches have been introduced in the literature for learning discrete representations with backpropagation . The simplest approach is the Straight-Through 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 et al. , 2017 ; Jang et al. , 2017 ) , 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 objective . When employed with complex architectures , Straight-Through and Gumbel-Softmax estimators often underperform due to this bias , as in figure 1 . The reason is that with continuous relaxation methods there is a 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 as well is the pressure towards non-discrete solutions , thus increasing bias further . Adding to the complexity of obtaining discrete representations , with current methods there is no direct incentive for the optimization procedure to obtain discrete representations : with Gumbel-Softmax the extent of how close-to-discrete a representation is obtained is controlled by a temperature variable , which must be manually tuned . Unbiased estimators like REINFORCE ( Williams , 1992 ) – and reduced variance extensions like REBAR ( Tucker et al. , 2017 ) and RELAX ( Grathwohl 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 a need for alternative methods for obtaining hassle-free discrete representations with neural networks , especially when increasing their complexity . In this work we present a regularization procedure for discrete representations , which can be used either as a standalone method , or in combination with existing continuous relaxation or straightthrough estimators . In its standalone form , the method replaces a discrete variable by a parameterized continuous function whose output ( say a sigmoid or softmax function ) corresponds to the discrete variable , and which is then regularized to produce discrete outputs . In combination with continuous relaxations such as Gumbel-Softmax , 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 of a Gaussian function is a measure of its resilience to noise . Given a Gaussian function f and correlated Gaussian variables , ′ ∈ Rd , noise stability of f is defined as Stab = E , ′ [ f ( ) 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 to 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 . 2 NOISE STABILITY OF GAUSSIAN FUNCTIONS . 2.1 STABILITY AND GAUSSIAN ISOPERIMETRY . Noise Stability of a Gaussian function f : Rn → R is defined for a noise parameter ρ ∈ ( 0 , 1 ) as Stabρ [ f ] = E , ′ [ f ( ) f ( ′ ) ] , ( 1 ) ′ = ρ + √ 1− ρ2 ′′ ( 2 ) where , ′ are called ρ-correlated Gaussian pairs and , ′′ ∼ N ( 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 = 1A , the indicator function of a set A , stability measures the probability that both and ′ remain within A. Stabρ [ f ] = P [ ∈ A ∧ ′ ∈ 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. , E [ f ] for f the object ’ s indicator function ) are half spaces . Theorem 1 ( Borell Isoperimetric Theorem ( Borell , 1985 ) ) . For fixed ρ ∈ ( 0 , 1 ) and f ∈ L2 ( Rn ) in Gaussian space with range [ 0 , 1 ] and fixed volume E [ f ] = α , Stabρ [ f ] is maximized by f = 1H where 1H is an indicator function of a half space with volume α . As a consequence , given a parameterized bounded Gaussian function , maximizing the stability makes the function f approach the indicator function of some half space as illustrated in figure 2 . For further details on noise stability in the context of Gaussian analysis of functions we refer the interested reader to ( O ’ Donnell , 2014 ) . 3 STABILITY REGULARIZATION . Stability regularization can be used either as a standalone method or in combination with GumbelSoftmax style continuous relaxation . 3.1 REGULARIZATION FOR DISCRETE VARIABLES . We start from a continuous model , say a neural network with L layers or modules , f = fL◦· · ·◦f1 ( x ) . Next , we describe how to employ stability regularization so that any arbitrary intermediate function fl ( z ; θ ) with bounded output learns to output discrete variables . Given input z for fl we estimate the stability of fl and , thereafter , maximize it by adding it to the loss objective as a regularizing term . For a single input vector z ∈ Rk , following the definition of stability in equation ( 1 ) , we sample ρ-correlated Gaussian variables , ′ as in equation ( 2 ) . We then evaluate fl twice : once for z + and once for z + ′ . The expectation of their product , E , ′ [ fl ( z + ) fl ( z + ′ ) ] is the stability , Stabρ [ fl ( z ) ] . Maximizing the stability for a single input z , Stabρ [ fl ( z ) ] for a fixed ρ ∈ ( 0 , 1 ) , the function fl approaches an indicator as described by Borell ’ s theorem . In a batch setting , we compute a Monte Carlo estimate of the expected stability over the input , that is Ez [ Stabρ [ fl ( z ) ] ] , by sampling one ρ-correlated Gaussian pair per batch element . Given z , , ′ ∈ Rn×k , so that i and ′i are ρ-correlated Gaussian , the estimate is computed as Ez [ Stabρ [ fl ( z ) ] ] ≈ 1 n ∑ i fl ( zi + i ) fl ( zi + ′ i ) , ( 4 ) where n is the batch size and the arithmetic operations are done element-wise . To maximize stability we sum or average the estimate in equation ( 4 ) across dimension and add the result as an additional regularization term to the loss function with which we train the model . 3.2 MEAN-CENTERED STABILITY REGULARIZATION . The regularization makes the function stable relative to correlated Gaussian noise by moving the inputs zi 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 fl given the input zi ensuring that the induced separation remains limited . The idea behind mean-centered stability regularization is compute the stability of fl − E [ fl ] so that maximizing stability reorients the separating hyperplane without changing the Gaussian volume of the corresponding halfspace . According to Borell ’ s theorem maximizing stability of fl for any fixed E [ fl ] causes it to approach a Gaussian halfspace indicator ( figure 2 ) . Since it can be expensive to compute the expectation of a neural network , we maximize the difference of two stability computations : Given parameters ρ1 , ρ2 , ρ2 < ρ1 we optimize Stabρ1 [ fl ] − Stabρ2 [ fl ] . We can show that this objective is equal to Stabρ1−ρ2 [ fl − E [ f ] ] to first order with an error that is quadratic in ρ1 , ρ2 . Proposition 1 . Given a Gaussian function f : Rn → [ 0 , 1 ] and parameters ρ1 , ρ2 ∈ ( 0 , 1 ) , Stabρ1−ρ2 [ f − E [ f ] ] = Stabρ1 [ f ] − Stabρ2 [ f ] +O ( ρ2 ( ρ2 − ρ1 ) ) . Algorithm 1 Stability Regularization Require : Input z ∈ Rn×k ; stability layer fl with range ( 0 , 1 ) m ; noise parameter ρ ∈ ( 0 , 1 ) ; stability constraint α ∈ ( 0 , 1 ) 1 : Sample , ′ ∈ Rn×k ρ-correlated Gaussian vectors . 2 : Compute y1 = fl ( z + ) , y2 = fl ( z + ′ ) 3 : Estimate average stability over batch per dimension as S = 1n ∑ i y1 , iy2 , i . 4 : Apply stability constraint per dimension : S = clip ( S , 0 , α ) . 5 : Sum S across dimensions and optimize by gradient descent Algorithm 2 Stability Regularization with Mean Centering Require : Input z ∈ Rn×k ; stability layer fl with range ( 0 , 1 ) m ; noise parameter ρ1 , ρ2 ∈ ( 0 , 1 ) , ρ2 < ρ1 1 : Sample ( 1 , ′1 ) ρ1-correlated and ( 2 , ′ 2 ) ρ2-correlated Gaussian vectors from Rn×k . 2 : Compute y1 = fl ( z + 1 ) fl ( z + ′1 ) , y2 = fl ( z + 2 ) fl ( z + ′ 2 ) 3 : Estimate average stability over batch per dimension as S = 1n ∑ i ( y1 , i − y2 , i ) . 4 : Sum S across dimensions and optimize by gradient descent See appendix A for a 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 , a constraint on stability limits how far the points can be from the boundary improving optimization . Borell ’ s theorem guarantees that the function f will converge to a halfspace for any ρ . We did not observe the method to be sensitive to ρ in our experiments . The precise procedures are described in algorithms 2 and 1 . Stability Regularized Layers . The stability regularized neural network layers can 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 as VAEs where it is important to be able to compute log probabilities of obtained samples . Given a block of layers fl with a Gumbel softmax ( or Gumbel sigmoid ) activation , i.e. , fl = GumbelSoftmax ( logits ) , we compute stability using a standard softmax or sigmoid without adding the Gumbel noise as Stabρ [ Softmax ( logits ) ] and use the Gumbel softmax output as input to the downstream network . 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 to 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 is a 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 fl . The rest of the network is unaffected . Depending on the application , there usually exist only a 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 a noticeable increase in computational cost in our experiments .

Summary: The paper presents a new method to train neural networks with stochastic discrete variables called stability regularisation. The method pushes  the outputs of functions of Gaussian random variables to be close to discrete, and unlike other methods used with discrete variables,  such as Gumbel Softmax which requires temperature annealing, it is easy to tune. The method is demonstrated on a very rich variety of tasks and models and shows state of the art performance.