Lattice Representation Learning

1 PRELIMINARIES . With a few notable exceptions , the majority of the practical research in representation learning assumes that the representation of the objects of interest ( sentences , images , audio signals , etc . ) are vectors of real numbers , as this allows us to use powerful optimization algorithms such as variants of gradient descent in the training of computational networks which encode objects into such representations and then use those representations in downstream tasks . Yet , the idea of representing objects using discrete structures ( for example , through categorical variables or through the use of quantization of otherwise real valued representations ) is rather enticing : sometimes we might inherently believe discrete representations to be the right way to model objects , or may want to use such representations in settings such as reinforcement learning and planning , where discrete actions are important . A classical result by Lindsay ( 1983 ) for maximum likelihood learning of mixture models tells us that the optimal mixing distribution can be chosen to be discrete ( and in fact , the discrete set need not be larger than the amount of training data ) ; this result implies also that the optimal associated “ approximate posterior ” ( when seen as a variational inference problem ) in fact can be chosen so that it produces discrete representations . The main difficulty associated with discrete representations is that it is not straightforward to train networks that produce and use them because either there is no meaningful sense in which differentiation can be used directly ( e.g . in the case of categorical variables ) , or in the case of quantization , the associated gradient is zero almost everywhere . In spite of these difficulties , notable progress has been made . For example , for categorical variables , Jang et al . ( 2017 ) , and Maddison et al . ( 2017 ) proposed essentially the same idea , under the names of Gumbel-Softmax and the Concrete Distribution , respectively . This idea , further improved by Tucker et al . ( 2017 ) , uses a continuous approximation to a “ one-hot ” encoded categorical distribution which can be learned by appropriately setting a parameter controlling the fidelity of the approximation . For the setting where a continuous representation is quantized to obtain the discrete one , an important development is the idea of straight-through estimation ( Bengio et al. , 2013 ) , in which quantization is applied in the “ forward ” direction , but replaces quantization with the identity operator when performing the “ backwards ” differentiation step - see also recent work attempting to provide theoretical justification of straight-through estimation ( Yin et al. , 2017 ) . Of particular relevance to our work is VQ-VAE ( van den Oord et al. , 2017 ) in which general vector quantization is used , together with straight-through estimation , to train the network and the vector quantizers . In fact , the entirety of this paper can be seen as a type of VQ-VAE where the vector quantizers are highly structured instead of being free parameters in a network ; as we will see , the advantages of introducing structure are that we can bring to bear analytical tools that give insight not only on how to train these networks with an equivalent continuous version , but also shed light on how we can approximate the performance of continuous Gaussian VAEs in a systematic way . In this article , we study the possibility of using lattices for object representations by borrowing from theoretical developments in information and coding theory ( Zamir et al. , 2014 ) , in particular as they apply to lossy source coding problems ( Shannon , 1959 ) - see Figure 1 for an example of two lattices . In our article , a lattice is defined as the set of all integral linear mixtures of a given set of basis vectors for an Euclidean space . Lattices have been long studied in information theory as they provide powerful means for building structured codes which are analytically tractable and have good space packing and covering properties , which are useful in channel coding and lossy source coding applications , respectively ( Zamir et al. , 2014 ) . We note that it has long been known that machine learning and lossy compression are related to each other ; for example , we refer the reader to the classical work on the information bottleneck ( Tishby et al. , 1999 ; Tishby & Zaslavsky , 2015 ; Shwartz-Ziv & Tishby , 2017 ; Slonim & Weiss , 2002 ) , the connection between maximum likelihood estimation and rate distortion explored by Banerjee et al . ( 2004 ) , Lastras-Montaño ( 2019 ) , and implicitly , Rose ( 1998 ) , Neal & Hinton ( 1998 ) and a line of research on autoencoders ( Giraldo & Príncipe , 2013 ; Alemi et al. , 2018 ; Higgins et al. , 2017 ) . Our work adds to this general line of research by adopting more specific structures - in this case , lattices - from the information theory field and applying them to machine learning problems . Our primary contribution will be to provide the necessary theory to train computational networks that employ lattices and dithered quantization , leading to a class of algorithms which we call dithered stochastic gradient descent algorithms . Dithering refers to the act of adding random noise to a quantized signal ( e.g . an image or a time series ) for the purposes of diminishing the effect that quantization has on it . The main reason dithering is important to our problem can be traced to a fundamental result in the study of lattices in information theory called the “ Crypto-Lemma ” ( see Zamir & Feder ( 1996 ) , G. D. Forney ( 2004 ) ) which to our knowledge , had not been used before to optimize , with gradient techniques , computational networks that employ quantization in a provably correct way . We will additionally connect the task of lattice representation learning with a well known technique in generative modeling using variational autoencoders ; specifically , Gaussian approximate posteriors ( Kingma & Welling , 2014 ) . In our experimental section , we will demonstrate how we can use our results to train a continuous VAE that can be re-interpreted , through our main result , as a discrete VAE that uses a finite dimensional lattice . 2 PROBLEM SETUP . To make our work concrete , we will illustrate it in the context of a standard encoder/decoder architecture ; for a given object d ∈ D ( whereD is the space of all possible objects ) , it ’ s representation is produced using an encoder : e ( d ) which is then decoded using a decoder : g ( e ( d ) ) . We pause to remark that there is a subtle difference between the way in which the information theory and machine learning communities use the words “ encoder ” and “ decoder ” ; the reader needs to be aware of this difference so as to not get lost in this article . In information theory , an encoder/decoder is almost always a means to represent and retrieve information digitally , in particular , there is a significant emphasis on the efficiency of such representations . In machine learning , an encoder ’ s output is generally thought of as a continuous vector ( of course , as we stated earlier , there are notable exceptions to this ) ; the fact that such a continuous vector is represented digitally ( in the form of a computer representation of a floating point number ) is in a sense an afterthought . Significantly more important in machine learning is the capability of learning parameters for an encoder and decoder through training examples , in which case the assumption of continuity , and more generally , differentiability is much more important . Our work lies in the middle of the two fields . We are interested in a machine learning application , but want to emphasize the representation cost for the objects being encoded as a first class metric to be optimized for . In our work , the representations are discrete , this is for any d ∈ D , e ( d ) belongs to a discrete space , which we denote as X . To each element x ∈ X we assign a probability p ( x ) , so that we can talk about a code for the representation space X . In this setting , a code is a mechanism for mapping x ∈ X to l ( x ) bits ( where l ( x ) stands for the code bit length ) so that x can be recovered , without loss , from the code for x of length l ( x ) . It is well known ( Cover & Thomas , 2006 ) that for any given code one can construct a probability distribution using the equation p ( x ) = 2−l ( x ) . It is also known that for a given distribution { p ( x ) } , the ideal choice for l ( x ) ( in units of bits ) is − log2 p ( x ) , in the sense of minimizing the number of expected bits under the assumption that the x are being generated with a distribution { p ( x ) } . It is also known that such ideal code length in general can only be achieved with block codes which compress multiple elements of X simultaneously . In our work , we will be indifferent to this - we will simply use l ( x ) = − log p ( x ) for the code length ; in fact for our theoretical development we will be using the unit of nats ( natural logarithms ) as it is more convenient . Assume that D is a random quantity uniformly distributed over the training data { d1 , · · · , dn } . Then the average representation cost for the training data is −ED logP ( e ( D ) ) ( 1 ) This is one of the metrics in our work . The other metric is free to be specified by the designer of the architecture . Typically , this will be of the form ED [ ` ( D , g ( e ( D ) ) ) ] ( 2 ) where ` is some loss function . This covers both unsupervised and supervised settings ; in the latter , we assume that the loss function also includes the label for D. It also includes variational autoencoders when they are trained using the Evidence Lower Bound ( ELBO ) . One way to create discrete representations e ( D ) is to quantize an otherwise continuous representation ec : e ( D ) = K ( ec ( D ) ) where K denotes a quantization operation ; for example , K could use the uniform quantizer with spacing ∆ applied to every one of the m dimensions . The encoder e ( ) , the decoder g ( ) and the representation code − log p ( ) all participate in the end-to-end objective function ED [ ` ( D , g ( K ( ec ( D ) ) ) ) ] + λED [ log 1 P ( K ( ec ( D ) ) ) ] ( 3 ) where λ > 0 is a parameter that controls the importance of the representation cost for the overall optimization process . In a nutshell , the problem to be solved is to find a theoretically sound basis for the design of good quantizers K that can be trained in an end-to-end computational network . As stated earlier , this can be seen as an attempt to bring additional mathematical rigor to the VQ-VAE concept .

Summary: This paper proposes a framework for gradient descent optimization of latent variable models where the latent code lies on a (n-dimensional) lattice. After a brief overview of relevant results in lattice theory, a method ("Dithered gradient descent") is proposed to differentiate through lattice quantization by means of an additive "dither" noise variable and a lower bound on the total reconstruction+representation loss using a prior distribution. After elaborating on the relationship between their framework and the VAE literature, the authors go on to give an explanation of the link between the covering property of a lattice and the reconstruction loss. Finally, empirical evidence is given for the better compression and reconstruction properties of a sentence autoencoder with lattice latent codes.