Lossless Compression with Probabilistic Circuits

1 INTRODUCTION . Thanks to their expressiveness , modern Deep Generative Models ( DGMs ) such as Flow-based models ( Dinh et al. , 2014 ) , Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) , and Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) achieved state-of-the-art results on generative tasks such as creating high-quality samples ( Vahdat & Kautz , 2020 ) and learning low-dimensional representation of data ( Zheng & Sun , 2019 ) . However , these successes have not been fully transferred into lossless compression . Specifically , GANs can not be used for lossless compression due to their inability to assign likelihoods to observations . Latent variable models such as VAEs rely on rate estimates obtained by lower-bounding the marginal likelihood of the data , i.e. , the quantity which is theoretically optimal for lossless compression ; they furthermore rely on sophisticated schemes such as bits-back coding ( Hinton & Van Camp , 1993 ) to realize these rates , oftentimes resulting in poor single-sample compression ratios ( Kingma et al. , 2019 ) . Therefore , good generative performance does not imply good compression performance for lossless compression , as the model needs to support efficient algorithms to encode and decode close to the model ’ s theoretical rate estimate . While both Flow- and VAE-based compression algorithms ( Hoogeboom et al. , 2019 ; Kingma et al. , 2019 ) support efficient and near-optimal compression under certain assumptions ( e.g. , the existence of an additional source of random bits ) , we show that Probabilistic Circuits ( PCs ) ( Choi et al. , 2020 ) are also suitable for lossless compression tasks . This class of tractable models has a particular structure that allows efficient marginalization of its random variables–a property that , as we show , enables efficient conditional entropy coding . Therefore , we introduce PCs as backbone models and develop ( de ) compression algorithms that achieve high compression ratios and high computational efficiency . Similar to other neural compression methods , the proposed lossless compression approach operates in two main phases — ( i ) learn good PC models that approximate the data distribution , and ( ii ) compress and decompress samples x with computationally efficient algorithms . The proposed lossless compression algorithm has four main contributions : A new class of entropy models . This is the first paper that uses PCs for data compression . In contrast to other neural compression algorithm , we leverage recent innovations in PCs to automatically learn good model architectures from data . With customized GPU implementations and better training pipelines , we are the first to train PC models with competitive performance compared to deep learning models on datasets such as raw MNIST . A new coding scheme . We developed a provably efficient ( Thm . 1 ) lossless compression algorithm for PCs that take advantage of their ability to efficiently compute arbitrary marginal probabilities . Specifically , we first show which kinds of marginal probabilities are required for ( de ) compression . The proposed algorithm combines an inference algorithm that computes these marginals efficiently given a learned PC and SoTA streaming codes that use the marginals for en- and decoding . Competitive compression rates . Our experiments show that on MNIST and EMNIST , the PC-based compression algorithm achieved SoTA bitrates . On more complex data such as subsampled ImageNet , we hybridize PCs with normalizing flows and show that PCs can significantly improve the bitrates of the base normalizing flow models . Competitive runtimes . Our ( de ) compressor runs 5-20x faster compared to available implementations of neural lossless compressors with near SoTA performance on datasets such as MNIST.1 Notation We denote random variables by uppercase letters ( e.g. , X ) and their assignments by lowercase letters ( e.g. , x ) . Analogously , we use bold uppercase ( e.g. , X ) and lowercase ( e.g. , x ) letters to denote sets of variables and their joint assignments , respectively . The set of all possible assignments of X is denoted val ( X ) . 2 TRACTABILITY MATTERS IN LOSSLESS COMPRESSION . The goal of lossless compression algorithms is to map every input sample to an output codeword such that ( i ) the original input can be reconstructed from the codeword , and ( ii ) the expected length of the codewords is minimized . Practical ( neural ) lossless compression algorithms operate in two main phases — learning and compression . In the learning phase , a generative model p ( X ) is learned from a dataset D : = { x ( i ) } Ni=1 . According to Shannon ’ s source coding theorem ( Shannon , 1948 ) , the expected codeword length is lower-bounded by the negative cross-entropy between the data distribution D and the model distribution p ( X ) ( i.e. , −Ex∼D [ log p ( x ) ] ) , rendering it a natural and widely used objective to optimize the model ( Hoogeboom et al. , 2019 ; Mentzer et al. , 2019 ) . In the compression phase , compression algorithms take the learned model p and samples x as input and generate codewords whose expected length approaches the theoretical limit ( i.e. , the negative cross-entropy between D and p ) . Although there exist various close-to-optimal compression schemes ( e.g. , Huffman Coding ( Huffman , 1952 ) and Arithmetic Coding ( Rissanen , 1976 ) ) , a natural question to ask is what are the requirements on the model p such that compression algorithms can utilize it for encoding/decoding in a computationally efficient manner ? In this paper , we highlight the advantages of tractable probabilistic models for lossless compression by introducing a concrete class of models that are expressive and support efficient encoding and decoding . To encode a sample x , a standard streaming code operates by sequentially encoding every symbol xi into a single bitstream b , such that xi occupies approximately− log p ( xi|x1 , . . . , xi−1 ) bits in b . As a result , the length of b is approximately− log p ( x ) . For example , Arithmetic Coding ( AC ) encodes the symbols { xi } Di=1 ( define D : = |X| as the number of features ) sequentially by successively refining an interval that represents the sample , starting from the initial interval [ 0 , 1 ) . To encode xi , the algorithm partitions the current interval [ a , b ) using the left and right side cumulative probability of xi : li ( xi ) : = p ( Xi < xi | x1 , . . . , xi−1 ) , hi ( xi ) : = p ( Xi≤xi | x1 , . . . , xi−1 ) . ( 1 ) Specifically , the algorithm updates [ a , b ) to the following : [ a + ( b−a ) ·li ( xi ) , a + ( b−a ) ·hi ( xi ) ) , which is a sub-interval of [ a , b ) . Finally , AC picks a number within the final interval that has the shortest binary representation . This number is encoded as a bitstream representing the codeword of x . 1Note that there exists compression algorithms optimized particularly for speed by using simple entropy models ( Townsend et al. , 2019 ) , though that also leads to worse bitrates . See Sec . 3.3 for detailed discussion . Upon decoding , the symbols { xi } Di=1 are decoded sequentially : at iteration i , we decode variable Xi by looking up its value x such that its cumulative probability ( i.e. , li ( x ) ) matches the subinterval specified by the codeword and x1 , . . . , xi−1 ( Rissanen , 1976 ) ; the decoded symbol xi is then used to compute the following conditional probabilities ( i.e. , lj ( x ) for j > i ) . Despite implementation differences , computing the cumulative probabilities li ( x ) and hi ( x ) are required for many other streaming codes ( e.g. , rANS ) . For these streaming codes , the main computation cost of both the encoding and decoding process comes from calculating li ( x ) and hi ( x ) . The main challenge for the above ( de ) compression algorithm is to balance the expressiveness of p and the computation cost of { li ( x ) , hi ( x ) } Di=1 . On the one hand , highly expressive probability models such as energy-based models ( Lecun et al. , 2006 ; Ranzato et al. , 2007 ) can potentially achieve high compression ratios at the cost of slow runtime , which is due to the requirement of estimating the model ’ s normalizing constant . On the other hand , models that make strong independence assumptions ( e.g. , n-gram , fully-factorized ) are cheap to evaluate but lack the expressiveness to model complex distributions over structured data such as images.2 This paper explores the middle-ground between the above two extremes . Specifically , we ask : are there probabilistic models that are both expressive and permit efficient computation of the conditional probabilities in Eq . ( 1 ) ? This question can be answered in its affirmative by establishing a new class of tractable lossless compression algorithms using Probabilistic Circuits ( PCs ) ( Choi et al. , 2020 ) , which are neural networks that can compute various probabilistic queries efficiently . In the following , we overview the empirical and theoretical results of the proposed ( de ) compression algorithm . We start with theoretical findings : the proposed encoding and decoding algorithms enjoy time complexity O ( log ( D ) · |p| ) , where |p| ≥ D is the PC model size . The backbone of both algorithms , formally introduced in Sec . 3 , is an algorithm that computes the 2D conditional probabilities { li ( x ) , hi ( x ) } Di=1 given any x in O ( log ( D ) · |p| ) time , as justified by the following theorem . Theorem 1 ( informal ) . Let x be a D-dimensional sample , and let p be a PC model of size |p| , as proposed in this paper . We then have that computing all quantities { li ( xi ) , hi ( xi ) } Di=1 takes O ( log ( D ) · |p| ) time . Therefore , en- or decoding x with a streaming code ( e.g. , Arithmetic Coding ) takes O ( log ( D ) ·|p|+D ) = O ( log ( D ) ·|p| ) time . The properties of PCs that enable this efficient lossless compression algorithm will be described in Sec . 3.1 , and the backbone inference algorithm with O ( log ( D ) ·|p| ) time complexity will later be shown as Alg . 1 . Table 1 provides an ( incomplete ) summary of our empirical results . First , the PC-based lossless compression algorithm is fast and competitive . As shown in Table 1 , the small PC model achieved a near-SoTA bitrate while being ∼ 15x faster than other neural compression algorithms with a similar bitrate . Next , PCs can be integrated with Flow-/VAE-based compression methods . As illustrated in Table 1 ( right ) , the integrated model significantly improved performance on sub-sampled ImageNet compared to the base IDF model . 3 COMPUTATIONALLY EFFICIENT ( DE ) COMPRESSION WITH PCS . In the previous section , we have boiled down the task of lossless compression into calculating conditional probabilities { li ( x ) , hi ( x ) } Di=1 given p and x . This section takes PCs into consideration and demonstrates how these queries can be computed efficiently . In the following , we first introduce 2Flow-model-based neural compression algorithms adopt p defined on mutually independent latent variables ( denoted Z ) , and improve expressiveness by learning bijection functions between Z and X ( i.e. , the input space ) . This is orthogonal to our approach of directly learn better p. Furthermore , we can naturally integrate proposed expressive p with the bijection functions and achieve better performance as demonstrated in Sec . 5. relevant background on PCs ( Sec . 3.1 ) , and then proceed to introduce the PC-based ( de ) compression algorithm ( Sec . 3.2 ) . Finally , we empirically evaluate the optimality and speed of the proposed compressor and decompressor ( Sec . 3.3 ) . 3.1 BACKGROUND : PROBABILISTIC CIRCUITS Probabilistic Circuits ( PCs ) are an umbrella term of a wide variety of Tractable Probabilistic Models ( TPMs ) . They provide a set of succinct definitions for popular TPMs such as Sum-Product Networks ( Poon & Domingos , 2011 ) , Arithmetic Circuits ( Shen et al. , 2016 ) , and Probabilistic Sentential Decision Diagrams ( Kisa et al. , 2014 ) . The syntax and semantics of a PC are defined as follows . Definition 1 ( Probabilistic Circuits ) . A PC p ( X ) represents a probability distribution over X via a parametrized directed acyclic graph ( DAG ) with a single root node nr . Similar to neural networks , every node of the DAG defines a computational unit . Specifically , each leaf node corresponds to an input unit ; each inner node n represents either a sum or a product unit that receives inputs from its children , denoted in ( n ) . Each node n encodes a probability distribution pn , defined as follows : pn ( x ) : = fn ( x ) if n is an input unit , ∑ c∈in ( n ) θn , c · pc ( x ) if n is a sum unit , ∏ c∈in ( n ) pc ( x ) if n is a product unit , ( 2 ) where fn ( · ) is an univariate input distribution ( e.g. , Gaussian , Categorical ) , and θn , c denotes the parameter that corresponds to edge ( n , c ) . Intuitively , sum and product units encode weighted mixtures and factorized distributions of their children ’ s distributions , respectively . To ensure a PC models a valid distribution , we assume the parameters associated with any sum unit n are normalized : ∀n , ∑c∈in ( n ) θn , c=1 . We further assume w.l.o.g . that a PC alternates between sum and product units before reaching an input unit . The size of a PC p , denoted |p| , is the number of edges in its DAG . This paper focuses on PCs that can compute arbitrary marginal queries in time linear in their size , since this is necessary to unlock the efficient ( de ) compression algorithm . In order to support efficient marginalization , PCs need to be decomposable ( Def . 2 ) ,3 which is a property of the ( variable ) scope φ ( n ) of PC units n , that is , the collection of variables defined by all its descendent input units . Definition 2 ( Decomposability ) . A PC is decomposable if for every product unit n , its children have disjoint scopes : ∀c1 , c2 ∈ in ( n ) ( c1 6= c2 ) , φ ( c1 ) ∩ φ ( c2 ) = ∅ . All product units in Fig . 1 are decomposable . For example , each purple product unit ( whose scope is { X1 , X2 } ) has two children with disjoint scopes { X1 } and { X2 } , respectively . In addition to Def . 2 , we introduce another structural property , structured decomposability ( abbreviated as SD ) , which is the key to guaranteeing computational efficiency of the proposed ( de ) compression algorithm . Definition 3 ( Structured decomposability ) . A PC is structured-decomposable if ( i ) it is decomposable and ( ii ) for every pair of product units ( m , n ) with identical scope ( i.e. , φ ( m ) = φ ( n ) ) , we have that |in ( m ) | = |in ( n ) | and the scopes of their children are pairwise identical : ∀i ∈ { 1 , ... , |in ( m ) | } , φ ( cmi ) =φ ( cni ) , where cmi and cni are the ith child unit of m and n. The PC shown in Fig . 1 is SD because for all three groups of product units ( grouped by their colors ) with the same scope , their children divide the variable scope in the same way . For example , the children of both orange units decompose the scope { X1 , X2 , X3 , X4 } into { X1 , X2 } and { X3 , X4 } . As a key sub-routine in the proposed algorithm , we describe how to compute marginal queries given a smooth and ( structured- ) decomposable PC in O ( |p| ) time . First , we assign probabilities to every input unit : for an input unit n defined on variable X , if evidence is provided for X in the query 3According to Poon & Domingos ( 2011 ) , another property termed smoothness is also required to compute marginals efficiently . However , since enforcing smoothness on any decomposable PC only imposes at most almost-linear increase in its size ( Shih et al. , 2019 ) , we omit introducing it here . ( e.g. , X = x , X < x ) , we assign to n the corresponding probability ( e.g. , p ( X = x ) , p ( X < x ) ) according to fn in Eq . ( 2 ) ; if evidence of X is not given , probability 1 is assigned to n. Next , we do a feedforward ( children before parents ) traverse of inner PC units and compute their probabilities following Eq . ( 2 ) . The probability assigned to the root unit is the final answer of the marginal query . Concretely , consider computing p ( x1 , x2 , x4 ) for the PC in Fig . 1 . This is done by ( i ) assigning probabilities to the input units w.r.t . the given evidences x1 , x2 , and x4 ( assign 0 to the input unit labeled X2 and ¬X4 as they contradict given evidences ; all other input units are assigned probability 1 ) , and ( ii ) evaluate the probabilities of sum/product units following Eq . ( 2 ) . Evaluated probabilities are labeled beside the corresponding units , hence the marginal probability is p ( x1 , x2 , x4 ) = 0.056 .

Summary: Probabilistic circuits are a formalism, developed relatively recently, for describing multivariate probability distributions. PCs are represented using directed acyclic graphs (DAGs), with an operational semantics, and it is relatively straightforward to deduce which operations (marginalization, maximisation, estimation of moments, etc.) are tractible based on the structure of the DAG and the locations of input variables.  This paper studies the application of PCs to lossless compression, identifying the specific type of marginalisation which is necessary for auto-regressive, arithmetic coding-style compression, and the necessary conditions a PC must satisfy in order for this marginalisation to be tractible. In section 3, a procedure is described for particularly efficient marginalisation by sharing computation.  Experiments show that the approach is competitive in terms of compression rate with other recent 'neural compression' works, and results are presented showing that PCs can be much faster than some existing methods, although I have serious concerns about these (see below).