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 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.   The 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.