iFlow: Numerically Invertible Flows for Efficient Lossless Compression via a Uniform Coder

1 Introduction . The volume of data , measured in terms of IP traffic , is currently witnessing an exponential yearon-year growth [ 13 ] . Consequently , the cost of transmitting and storing data is rapidly becoming prohibitive for service providers , such as cloud and streaming platforms . These challenges increasingly necessitate the need for the development of high-performance lossless compression codecs . One promising solution to this problem has been the development of a new class of so-called “ neural compression ” algorithms [ 30 , 38 , 18 , 4 , 17 , 37 , 26 , 31 , 6 , 41 ] . These methods typically posit a deep probabilistic model of the data distribution , which , in combination with entropy coders , can be used to compress data with the minimal codelength bounded by the negative log-likelihood [ 29 ] . However , despite reliably improved compression performance compared to traditional codecs [ 14 , 33 , 9 , 34 ] , meaningful commercial applications have been limited by impractically slow coding speed . In this paper , we focus on developing approaches with deep probabilistic models based on normalizing flows [ 11 , 25 , 32 , 27 , 7 , 28 , 23 ] . A normalizing flow admits a learnable bijective mapping between input data and a latent variable representation . In this paradigm , inputs can be compressed by first transforming data to latent variables , with the resulting output encoded by a prior distribution . Compared to other classes of generative models , normalizing flows typically perform best in both tasks of probability density estimation ( as compared with variational autoencoders [ 24 , 15 , 8 ] ) and inference speed ( as compared with autoregressive factorizations [ 35 , 39 , 22 ] ) . This suggests that compression with flows can jointly achieve high compression ratios along with fast coding times . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Unfortunately , lossless compression requires discrete data for entropy coding , and the continuous bijections of normalizing flows would not guarantee discrete latent variables . As such , transformations would require discretization , resulting in a loss of information [ 41 , 3 ] . To resolve this issue , Integer Discrete Flows ( IDF ) [ 18 , 4 ] ( Fig . 1 , left ) proposed an invertible mapping between discrete data and latent variables . Similarly , iVPF [ 41 ] achieves a discrete-space bijection using volume-preserving flows . However , the above models must introduce constraints on flow layers to ensure a discretespace bijection , which limits the expressivity of the transformation . Local Bits-Back Coding ( LBB ) [ 17 ] ( Fig . 1 , middle ) was the first approach to succeed in lossless coding with the flexible family of continuous flows . It resolves the information loss by coding numerical errors with the rANS coder [ 12 ] . However , the coder is extraordinarily slow , making the method impractical . In this paper , we introduce iFlow : a numerically invertible flow-based neural compression codec that bridges the gap between high-fidelity density estimators and practical coding speed . To achieve this we introduce two novelties : Modular Scale Transform ( MST ) and Uniform Base Conversion Systems ( UBCS ) . MST presents a flexible family of bijections in discrete space , deriving an invertible class of flows suitable for lossless compression . UBCS , built on the uniform distribution , permits a highly efficient entropy coding compatible with MST . The main ideas are illustrated in Fig . 1 ( right ) . We test our approach across a range of image datasets against both traditional and neural compression techniques . Experimentally we demonstrate that our method achieves state-of-the-art compression ratios , with coding time 5× quicker than that of the next-best scheme , LBB . 2 Numerically Invertible Flows and Lossless Compression . Let f : X → Z be our normalizing flow , which we build as a composition of layers such that f = fL ◦ ... ◦ f2 ◦ f1 . Defining d-dimensional y0 = x ∈ X and yL = z ∈ Z , the latents can be computed as yl = fl ( yl−1 ) , with pX ( x ) calculated according to log pX ( x ) = log pZ ( z ) + L∑ l=1 log |det Jfl ( yl−1 ) | , ( 1 ) where Jfl ( x ) is the Jocabian matrix of the transformation fl at yl−1 . In this paper , we use log to denote base 2 logarithms . Training the flow should minimize the negative log-likelihood− log pX ( x ) . The inverse flow is defined as f−1 = f−11 ◦ ... ◦ f −1 L , such that f −1 ( z ) = x . There are many types of invertible bijections , fl . Popular choices include element-wise , autoregressive ( e.g . coupling layers ) [ 41 ] and 1× 1 convolution transformations [ 25 ] . To perform lossless compression with normalizing flows , the input data should be transformed to latent variables , which are then encoded with a prior distribution . As data and encoded bits should be in binary format , x and z must be discrete , with a bijection between discrete inputs and outputs established . However , as discussed , this operation is usually intractable with popular classes of transformations used in continuous flows due to the numerical errors induced by discretization . In this section , we introduce a novel algorithm derived from continuous flows that allows us to achieve an exact bijiective mapping between discrete x and discrete z . We present the general idea of a flow layer ’ s numerical invertibility in discrete space in Sec . 2.1 . We then introduce our flow layers in Sec . 2.2 and 2.3 , with their application to compression discussed in Sec . 2.4 and 2.5 . 2.1 Invertibility of Flows in Discrete Space . To discretize the data , we follow [ 41 ] and adopt k-precision quantization to assign floating points to discretization bins such that x̄ = b2k · xc 2k , ( 2 ) where b·c is the floor function . The associated quantization error is bounded by |x̄− x| < 2−k . Any d-dimensional x can be quantized into bins with volume δ = 2−kd . The probability mass of x̄ can then be approximated by P ( x̄ ) = p ( x̄ ) δ , such that the theoretical codelength is − log ( p ( x̄ ) δ ) . Denote our proposed numerically invertible flow ( iFlow ) according to f̄ = f̄L ◦ ... ◦ f̄1 . The input and output should be subject to k-precision quantization , where each layer should be appropriately invertible such that ȳl−1 ≡ f̄−1l ( f̄l ( ȳl−1 ) ) . Denoting ȳ0 = x and z̄ = ȳL , it is further expected that z̄ = f̄ ( x ) and x ≡ f̄−1 ( z̄ ) . We will show in the following subsections that f̄ can be derived from any continuous flow f in such a way that the error between f̄ ( x ) and f ( x ) is negligible . 2.2 Numerically Invertible Element-Wise Flows . We begin by describing our approach for element-wise flows z = f ( x ) , where f is a monotonic element-wise transformation . For notation simplicity , we assume that the flow contains just one layer . In most cases , f does not present a unique inverse in discrete space , such that f−1 ( z̄ ) 6= x . In what follows , we introduce an approach to derive an invertible , discrete-space operation from any element-wise flow transformation . For simplicity , we denote the inputs and output pairs of a continuous flow as x and z = f ( x ) ; and we denote the k-precision quantized input-output pairs of iFlow as x̄ and z̄ = f̄ ( x̄ ) . 2.2.1 Scale Flow . Algorithm 1 Modular Scale Transform ( MST ) : Numerically Invertible Scale Flow f ( x ) = R/S · x . Forward MST : z̄ = f̄ ( x̄ ) . 1 : x̂← 2k · x̄ ; 2 : Decode rd using U ( 0 , R ) ; ŷ ← R · x̂+ rd ; 3 : ẑ ← bŷ/Sc , re ← ŷ mod S ; 4 : Encode re using U ( 0 , S ) ; 5 : return z̄ ← ẑ/2k . Inverse MST : x̄ = f̄−1 ( z̄ ) . 1 : ẑ ← 2k · z̄ ; 2 : Decode re using U ( 0 , S ) ; ŷ ← S · ẑ + re ; 3 : x̂← bŷ/Rc , rd ← ŷ mod R ; 4 : Encode rd using U ( 0 , R ) ; 5 : return x̄← x̂/2k . We begin with the scale transformation defined as z = f ( x ) = a · x , ( a > 0 ) , from which we are able to build more complex flow layers . The input can be converted to an integer with x̂ = 2k · x̄ , and the latent variable z̄ can be recovered from integer ẑ by z̄ = ẑ/2k . Inspired by MAT in iVPF [ 41 ] , we first approximate a with a fractional such that a ≈ RS , where R , S ∈ N. Denote ŷ = R · x̂+ rd where rd is sampled uniformly from { 0 , 1 , ... , R− 1 } , we can obtain ẑ with ẑ = bŷ/Sc , and the remainder re = ŷ mod R can be encoded to eliminate the error . It is clear that ( x̂ , rd ) ↔ ŷ ↔ ( ẑ , re ) are bijections . Let U ( 0 , A ) be the uniform distribution such that P ( s ) = 1A , s ∈ { 0 , 1 , ... , A− 1 } . The numerically invertible scale flow f̄ is displayed in Alg . 1 . As the modular operation is essential to Alg . 1 , we name it the Modular Scale Transform ( MST ) . Setting S andR = round ( S ·a ) to be large , we observe that the following propositions hold . Proposition 1 . [ 41 ] |z̄ − f ( x̄ ) | < O ( S−1 , 2−k ) . Proposition 2 . The codelength of MST is Lf ( x̄ , z̄ ) = logS − logR ≈ − log |a| = − log |f ′ ( x̄ ) | . Proposition 1 establishes that the error is small if S and k are large . Proposition 2 demonstrates that the codelength is almost exactly the log-Jocabian of f . The above properties and the correctness of the algorithm are discussed in the Appendix . Further note that , with the exception of the usual encoding and decoding processes , MST ’ s operations can be parallelized , resulting in minimal additional overhead . In contrast , MAT ’ s operation in iVPF [ 41 ] only deals with volume-preserving affine transform and is performed sequentially along dimensions , limiting its usage and efficiency . 2.2.2 General Element-wise Flow . Algorithm 2 Numerically Invertible Element-wise Flows Forward : z̄ = f̄ ( x̄ ) . 1 : Get x̄l , x̄h , z̄l , z̄h given x̄ , so that x̄ ∈ [ x̄l , x̄h ) ; 2 : Set large S ; get R with Eq . ( 4 ) ; x̄′ ← x̄− x̄l ; 3 : Get z̄′ with forward MST in Alg . 1 , given input x̄′ and coefficients R , S ; 4 : return z̄ = z̄′ + z̄l. . z̄ ∈ [ z̄l , z̄h ) . Inverse : x̄ = f̄−1 ( z̄ ) . 1 : Get x̄l , x̄h , z̄l , z̄h given z̄ , so that z̄ ∈ [ z̄l , z̄h ) ; 2 : Set large S ; get R with Eq . ( 4 ) ; z̄′ ← z̄ − z̄l ; 3 : Get x̄′ with inverse MST in Alg . 1 , given input z̄′ and coefficients R , S ; 4 : return x̄ = x̄′ + x̄l. . x̄ ∈ [ x̄l , x̄h ) . For simplicity , we assume the non-linear f is monotonically increasing . For a monotonically decreasing f , we simply invert the sign : −f . MST in Alg . 1 works with linear transformations , and is incompatible with non-linear functions . Motivated by linear interpolation , we can approximate the non-linear flow with a piecewise linear function . In general , consider input x̄ to be within an interval x̄ ∈ [ x̄l , x̄h ) 1 and z̄l = f ( x̄l ) , z̄h = f ( x̄h ) . The linear interpolation of f is then finp ( x ) = z̄h − z̄l x̄h − x̄l ( x̄− x̄l ) + z̄l . ( 3 ) It follows that z̄ can be derived from MST with input x̄−x̄l followed by adding z̄l . Furthermore , z̄ and x̄ can be recovered with the corresponding inverse transformation . To preserve monotonicity , z̄ must be within interval z̄ ∈ [ z̄l , z̄h ) . Denoting the invertible linear flow derived from finp as f̄inp , it must hold that f̄inp ( x̄l ) ≥ z̄l , f̄inp ( x̄h−2−k ) ≤ z̄h−2−k . As we use MST in Alg . 1 , we observe that when z̄h−z̄l x̄h−x̄l is approximated withR/S , the minimum possible value of f̄inp ( x̄l ) is z̄l ( when rd = 0 in Line 2 ) and the maximum possible value of f̄inp ( x̄h−2−k ) is b ( R·2k ( x̄h−x̄l−2−k ) +R−1 ) / ( S ·2k ) c+z̄l ( when rd = R− 1 in Line 2 ) . Given large S , the largest possible value of R should be R = ⌊ ( 2k · ( z̄h − z̄l ) − 1 ) · S + 1 2k · ( x̄h − x̄l ) ⌋ . ( 4 ) Another difficulty is in determining the correct interpolation interval [ x̄l , x̄h ) . One simple solution is to split the input domain into uniform intervals with length 2−h such that x̄l = b ( 2h · x̄ ) /2hc , x̄h = x̄l + 2 −h . However , for the inverse computation given z̄ , it is not easy to recover the interpolation interval as f−1 ( z̄ ) may not be within [ x̄l , x̄h ) . Instead , as z̄ ∈ [ z̄h , z̄l ) , the interval [ x̄l , x̄h ) can be obtained via a binary search in which x̄l is obtained with z̄ ∈ [ f ( x̄l ) , f ( x̄h ) ) . Another approach is to split the co-domain into uniform intervals such that z̄l = b ( 2h · z̄ ) /2hc , z̄h = z̄l + 2−h , with x̄l = f−1 ( z̄l ) , x̄h = f−1 ( z̄h ) . In this case , determining [ z̄l , z̄h ) during the inverse computation is simple . While for the forward pass , [ z̄l , z̄h ) should be determined with a binary search such that x̄ ∈ [ f−1 ( z̄l ) , f−1 ( z̄h ) ) . In practice , we have simpler tricks to determine the correct interval [ x̄l , x̄h ) , [ z̄l , x̄h ) for both the forward and inverse computation , which can be found in the Appendix . The general idea of the non-linear flow adaptations are summarized in Alg . 2 . Note that Alg . 2 can be used in any element-wise flow including linear flow . In Alg . 2 , Proposition 2 holds such that Lf ( x̄ , z̄ ) = − log ( R/S ) ≈ − log z̄h−z̄lx̄h−x̄l ≈ − log |f ′ ( x̄ ) | . It is possible to arrive at a similar conclusion in Proposition 1 where the corresponding error is |z̄ − f ( x̄ ) | < O ( S−1 , 2−k , 2−2h ) . 1All intervals must partition the domain of f .

Summary: This paper proposes a lossless image compression method based on normalizing flows. In order for this class of flows to be useful for compression, they need to be working in the quantized domain, and (of course) have to be invertible. The authors present an extension to a much more general class of functions that can be used in the composition formulation of the flow, therefore being able to achieve higher rates of compression than what was possible before. Additional contributions include using a variant of bits-back coding to achieve higher compression rates, and a much faster entropy coding method that vastly outperforms rANS.