Finite-State Autoregressive Entropy Coding for Efficient Learned Lossless Compression
Authors: Yufeng Zhang
Hang Yu
Jianguo Li
Weiyao Lin
Venue: ICLR 2024 spotlight (2024)
Keywords: Lossless Compression
Autoregressive Model
Acceleration
Entropy Coding
Autoencoder
Primary area: general machine learning (i.e., none of the above)

Abstract: Learned lossless data compression has garnered significant attention recently due to its superior compression ratios compared to traditional compressors. However, the computational efficiency of these models jeopardizes their practicality. This paper proposes a novel system for improving the compression ratio while maintaining computational efficiency for learned lossless data compression. Our approach incorporates two essential innovations. First, we propose the Finite-State AutoRegressive (FSAR) entropy coder, an efficient autoregressive Markov model based entropy coder that utilizes a lookup table to expedite autoregressive entropy coding. Next, we present a Straight-Through Hardmax Quantization (STHQ) scheme to enhance the optimization of discrete latent space. Our experiments show that the proposed lossless compression method could improve the compression ratio by up to 6\% compared to the baseline, with negligible extra computational time. Our work provides valuable insights into enhancing the computational efficiency of learned lossless data compression, which can have practical applications in various fields. Code is available at https://github.com/alipay/Finite_State_Autoregressive_Entropy_Coding.

Published as a conference paper at ICLR 2024
FINITE-STATE AUTOREGRESSIVE ENTROPY CODING FOR EFFICIENT LEARNED LOSSLESS COMPRESSION
Yufeng Zhang 2 Hang Jianuu Li Weiyao Lin
Shanghai Jiao Tong University, PAn Group
ABSTRACT
Learned lossless data compression has garnered significant attention recently due to its superior compression ratios compared traditional compressors. However, the computational efficiency of these models jeopardizes their practicality. This paper proposes novel system for improving the compression ratio while maintaining computational efficiency for learned lossless data compression. Our approach incorporates two essential innovations. First, we propose the Finite-State AutoRe- gressive (FSAR) entropy coder, an efficient autoregressive Markov model based entropy coder that utilizes lookup table expedite autoregressive entropy coding. Next, we present a Straight- Through Hardmax Quantization (STHQ) scheme to enhance the optimization of discrete latent space. Our experiments show that the proposed lossless compression method could improve the compression ratio by up to 6% compared to the baseline, with negligible extra computational time. Our work provides valuable insights into enhancing the computational efficiency of Icarned lossless data compression, which can have practical applications in vari- com/ alipay/Finite Finite
ous fields. Code is available at httpss://ttthttii A.t.rrrgeeesiiie.............. SttpessivveEEntttreeeeee Beto
INTRODUCTION
Lossless data compression has been a prominent area of interest in both academic research and industry. By reducing the amount of space required for data storage and minimizing the bandwidth necessary for data communication, lossless compression codecs hold great potential for applications in broad range of computing and communication systems.
practical lossless compression codec requires both superior compression ratios and computational ef- ficiency, but most existing methods fail to achieve both simultaneously. Traditional methods (Boutell 1997 webmproject] 2023) rely on hand-crafted codecs that offer efficiency but yield suboptimal compression ratios. Conversely, recent advances in machine learning have introduced generative
models, including Autoencoders (Kingma & Welling 2013) Burda et al. 2013), GANs (Mirza
& Osindero 2014 (Arjovsky (207)). Flow models (Dinh et al. [0166] Kingma & Dhariwal
2018) and Autoregressive models (van den Oord et al.) 2016b Parmar et al.
2018). which exhibit
remarkable potential in modeling data likelihood. These models automatically learn codecs for different domains given sufficient training data. However, while these models have demonstrated improved compression ratios compared to the traditional methods, they suffer from high time complexity, rendering them impractical for general-purpose computation devices such as CPUs.
In this paper, we target improving the compression ratio while maintaining the computational efficiency for lossless compression. Our focus lies on latent space models, specifically autoencoders. due their ease of optimization and efficient computation (Mentzer et al. 2018 Townsend et al., 2019b; Ballé et al., 2018 Minnen et al. 2018). Most existing frameworks in this direction could be built upon the Asymmetric Numeral System (ANS) (Duda 2009) for latent entropy coding, and can be categorized into three groups: continuous random latent space, discrete deterministic latent space, and discrete autoregressive latent space. The decompression workflows of these frameworks are illustrated in Figure la Figure abl Figure Ic| respectively. The first group, based on continuous latent spaces (Townsend et al. 2019a Kingma et al. 2019 Theis & Ho 2021 Ruan et al.) (021). typically employs
coding( Townsend al.] 2019a)
(cf. Figure lb). Although the continuous
"This work was done when Yufeng Zhang was a research intern at Ant Group.
Corresponning Authors.

---

Published as a conference paper at ICLR 2024
ptPDe Comto Modes
AGG Ensode
[AN5 Coulin
Ausonogrsstie AANS
FSAMAAAS
Pl
AMS Decade
ANS Decrde
ANG
ANS
Decode
Decane
Diccade
Decrade
(a) Bits-back ANS
(b) tANS
(c) Autoregressive TANS
(d) FSAR-ANNS
Figure 1: Different ANS-based decompression workflows for latent space models. (a) Bits-back ANS (Twnsend
al. 2019a for continuous random latent space. (b) table ANS (LANS) (Duda 2013) coder for discrete
deterministic latent space. (c) Autoregressive range ANS (rANS) (Duda 2013) coder for discrete autoregressive latent space. (d) The proposed FSAR-tANS coder for discrete autoregressive latent space.
latent space facilitates optimization and achieves a favorable compression ratio, the decompression process becomes complex and slow due to the additional ANS encode step via the inference model. By contrast, frameworks (Ballé al. 2018 Mentzer et al. 2018) utilizing discrete deterministic latent spaces enable fast decompression using low-complexity table ANS (tANS) coder (Duda 2013) (cf. Figure lb). However, the representation power of independent discrete space is limited, and optimizing the discrete latent space is notoriously challenging, resulting in an inferior compression ratio. To mitigate this issue, the third group of works (Minnen et al.) 2018 Lee et al. [2018] Cheng ett al. 2020) improves the compression ratio by assuming flexible autoregressive dependence among the discrete latent variables instead of strict independence. However, similar to the continuous latent space, this improvement comes the cost of decompression speed, as autoregressive models require lengthy sequential decompression process that involves a complex context model in each iteration (cf. Figure lc). Ultimately, existing frameworks fail to simultancously achieve a satisfactory compression ratio and computational efficiency.
In contrast to the aforementioned works, our proposed method combines the best of both worlds. It features an efficiently coded discrete autoregressive latent space with flexible latent model that can fit the input data well while robustly optimizing this latent space. We first develop an efficient autoregressive model called the Finite-State AutoRegressive (FSAR) model as the latent prior, which can be easily implemented with a lookup table, as illustrated in Figure Id Compared with other autoregressive models, the FSAR implementation has a similar time complexity to the low-complexity tANS, only adding an extra lookup step to the original tANS implementation. Thus, it achieves an excellent compression ratio with the help of autoregressive models while keeping the computation as efficient as tANS. Moreover, to better optimize the discrete latent space as required by the FSAR, propose Straight-Through Hardmax Quantization (STHQ) based the straight-through estimator (Bengio et al., 2013). STHQ further strengthens the representation ability of the latent space and improves the compression ratio. To summarize, our contributions are listed follows:
The Finite-State AutoRegressive (FSAR) entropy coder is proposed for flexible discrete latent space coding. It combines low-complexity autoregressive Markov model with fast entropy coder to achieve efficient latent coding.
An optimization scheme called Straight-Thruugh Hardmax Quantization (STHQ) is proposed for robust optimization of the discrete latent space. enables gradient descent via the deterministic quantization process while maximizing the latent entropy for better likelihood estimation.
Extensive experiments demonstrate that our proposed method improves the compression ratio by up to 6% compared to commonly used discrete deterministic autoencoders, with negligible additional computational time. Moreover, it is over 300 times faster than commonly used discrete autoregressive models while achieving similar compression ratio.
2 RELATED WORKS
In this section, we review learned compression methods based on discrete deterministic latent space models and discrete autoregressive latent space models, as they directly relevant to our work. For a brief survey on continuous random latent space models, we refer the readers to Appendix [A]
Discrete Deterministie Latent Space Numerous compression methods employ discrete deterministic autoencoders due to their simplicity in latent coding. However, quantizing continuous latent space to discrete counterparts while enabling gradient calculation for optimization poses challenge.
2

---

Published as a conference paper at ICLR 2024
Ja) Finite State Autoragressive Entropy Coding
Adoetive State Number
Fintt State Markov Model
Laok-up
Latent Coder
Table
0
Latent Stream
Codebook
Latent Indices
bok-rodue
(x(22
Generative
Network
Data
44000 Penzzna
8y/6
Likeli-
hood
Data ANS Coder
Quant
Dequant
Data Stream
[b] Straight-thruugh Quantization
Figure 2: Overview of the proposed method. Here x represents the input data, z, y represents latent variables and samples respectively, ANS are employed for entropy coding of both the latent space and the data. Our major contributions are highlighted in red dashed boxes: (a) Finite-State Autoregressive Entropy Coding for latent entropy coding. (b) Straigh-TThouugh Hardmax Quantization for latent space optimization scheme.
To address this, universal quantization and soft scalar quantization have been applied respectively in (Ballé et al.] (2018) and (Mentzer et 2018), allowing for gradient pass-through while simulating quantization. However, these methods overlook the complexity of the high-dimensional latent space, resulting in poor approximation. Alternatively, VQ-VAE (van den Oord et al. 2017) introduces vector quantization (VQ) and improves latent modeling through an additional codebook. Subsequent works, such (Razavi et al., 2019; Esser et al., 2020). adopt this approach, employing trainable codebooks to enhance the quality of generative models. However, optimizing VQ-VAE is difficult as it requires auxiliary loss function similar to the K-means algorithm. To counteract this problem, Sønderby (017)) proposed the continuous relaxation of VQ, representing the VQ process as sampling from categorical distribution and enabling optimization using the ELBO similar VAEs. This idea was further extended by utilizing hierarchical relaxed VQ- VQ-SAEs to model more complex latent spaces (Zhu
(022). However, this relaxation breaks the deterministic nature of VQ. Instead, Agustsson et al. (2017) introduce soft VQ with deterministic annealing to enable optimization through hard quantization, but the practical implementation of the annealing schedule requires careful tuning. Moreover, Takida et al. (2022) propose self-annealed stochastic quantization (SQ-VAB) to make the stochastic VQ-VAE VAE converge to deterministic quantization. Unfortunately, the convergence from SQ to VQ relies on the collapse of the likelihood term (i.c., the generative model) (Takida et al. 20022). which is impractical for lossless compression as the likelihood is required for entropy coding.
Discrete Autoregressive Latent Space Autoregressive models have been widely applied for image likelihood modeling and generation (van den Oord et al. 2016bb a Razavi et al. 2019). However, most methods exploit masked convolution (van den Oord et al.) 2016b) to implement autoregressive models in the observation space for images, incurring heavy floating-poii operations (FLOPs) and increasing the computational cost. Furthermore, the sequential nature of autoregressive models, where each observation depends previous ones, hinders parallel computations and further diminishes efficiency. Although recent works (Ruan et al. 2021 Ryder et al., 2022; Guo et al., 2022; 2023)
develop parallelizable autoregressive models for efficient coding on parallel computation devices like GPUs, their theoretical time complexity remains high, resulting in suboptimal efficiency on general computation devices such as CPUs. One potential solution is to apply autoregressive methods (e.g., masked convolutions) to the lower-dimensional latent space (Minnen et al., 2018) Lee et al., 2018 Cheng al. 2020), the latent space typically smaller than the observation space and subsequently reducing the number of sequential steps. However, the number of FLOPs in masked convolution still remains large.
In this work, our objective is to improve the compression ratio of autoencoder- codecs while maintaining computational efficiency. To this end, we prefer the discrete autoregressive latent space due to its superior compression ratio, and we aim to further boost its efficiency.
3 THE PROPOSED LEARNED LOSSLESS COMPRESSION ARCHITECTURE
In this section, we propose an effective and efficient framework for lossless data compression.
Overall Architecture Our proposed lossless compression system combines latent space models with autoregressive models, as depicted in Figure 2, The framework is composed of three main parts: 1) backbone network that consists of the inference network and generative network, 2) latent space that incorporates an efficient autoregressive Markov model, and 3) a quantization method for the discrete latent space optimization. In this system, the data is initially processed by an inference
3

---

Published as a conference paper at ICLR 2024
Bits
Bits
(t--
St
S/-1
(ab1)(5t-1-11)
Yi
Table
Table
Yi
(a)
(b)
Figure 3: Different Designs of the ANS Coder, where Table represents the state transition lookup table of table ANS (tANS), while s;-1 and s; correspond to the previous and current state of ANS, respectively. Additionally, p denotes the prior probability utilized by ANS to construct I able. (a) Non-Autoreeressive IANS Coding. (b) Proposed Finite-State Autoregressive (Order-2 Channel) LANS Coding. Here, L signifies the lookup table.
network to obtain discrete latent variables z using codebook, which are then quantized to y (see Figure (b)) and encoded by an autoregressive latent ANS coder (see (Figure 2 (a)). Based on these latent variables, the generative network generates the data likelihood, which is essential for data ANS coding. More details on this framework are provided Appendix B
Backbone Network The backbone network serves as crucial component of the latent space model. Generally, deeper networks yield a more compact latent space and better generative results, thus improving the compression ratio. However, deeper networks demand more computational resources, making shallow networks preferable for efficient compression. Moreover, for lossless compression, the generative results should be random rather than deterministic. This is because the distribution of the observation space is necessary for entropy coding of the original data. Thus, autoencoders with collapsed generative networks (Takida et al. 2022) cannot be employed for lossless compression.
Latent Space The latent space, serving condensed representation of the original data, also needs to be encoded in the final bitstream. As discussed in Section [] encoding continuous random latent spaces is fraught with difficulties, motivating our preference for efficient frameworks that leverage discrete deterministic latent spaces. In line with previous findings that vector quantized latent spaces yield superior results for generative tasks (Kang et al.) (022), we adopt this format as the latent space representation in our framework. Furthermore, to enhance the compression of the latent space without relying on deep networks as the backbone, we introduce an efficient autoregressive model, namely the FSAR model as compensatory measure (see Figure (a)), which will be further elaborated upon in Scction The FSAR model can be easily implemented with lookup table and boasts similar time complexity to modern entropy coders such ANS (Dudal 2009). Morcover, we propose Learnable State Number (LSN) to further reduce the space complexity associated with the lookup table.
Optimization As discussed in Section 2 optimizing discrete deterministic latent spaces can be challenging due to the non-differentiable quantization process. Therefore, we analyze existing works and propose an improved guantization scheme, namely STHQ, based on the straight- estimator (Bengio et al. 2013) Figure (b)). This scheme facilitates rapid and robust optimization of the vector quantized latent space, which will be discussed Section 5 Thanks to the STHQ quantization scheme, the entire framework, encompassing the backbone network, latent space, and vector quantization codebook, can be collectively optimized in an end-to-end manner.
4 FINITE-STATE AUTOREGRESSIVE ENTROPY CODING
Autoregressive models have shown promise for improving compression ratios in lossless compression. However, their limited efficiency poses significant restrictions. this section, we propose a novel and compact approach to autoregressive modeling based on finite-state Markov models. This model acts as the prior for a discretized latent space within deterministic autoencoders, thereby shortening the sequential process and enabling compact coding. Importantly, the implementation of this model using a lookup table showcases low complexity comparable to modern entropy coders. We refer to the combination of FSAR and entropy coder as Finite-State Autoregressive Entropy Coding.
The Efficiency of Autoregressive Models Existing autoregressive models are computationally expensive, primarily due to the long sequential process and high time complexity, as detailed in Section [ On the contrary, modern entropy coders, despite their sequential nature, typically demonstrate satisfactory efficiency due to their low complexity at each step. This suggests that by making autoregressive models sufficiently compact, high computational efficiency is still attainable. Indeed, compression methods prioritizing computational efficiency. such as Ballè et al. (2018) and Zhu al, (2022), could utilize efficient entropy coders such tANS (Duda 2013) discrete latent
4

---

Published as a conference paper at ICLR 2024
spaces, relying solely on lookup tables with minimal time complexity. Hence, aim to harness the capabilities of such entropy coders to alleviate the computational burden of autoregressive models.
Finite-State Markov Models In discretized latent space, the latent variables can only take values from a finite set. Consequently, the outcomes of the autoregressive model for a discretized latent space can also be represented by a finite set. To implement such autoregressive model, we introduce an order-N finite-stare state Markov model .M that produces prior distributions based on the previous N variables, each with C possible states. Each iteration step of the autoregressive Markov model can be executed with an N-dimensional lookup table L. similar to IANS (Duda 2013). Denoting samples from the finite set as y, this process can be succinctly expressed as:
(1) As an illustration, Figure 3b shows an order-2 sinite-state Markov model combined with tANS coding, leveraging 2- d-dimensional lookup table L to reduce the time complexity. Here, only two memory operations are necessary to read both dimensions of the lookup table and retrieve the corresponding state transition lookup table for tANS.
p. _M(y -.. ..N) = L(y,, ...y)) V-N E 21,22..... C}.
is important to mention that the training process of the finite-state Markov model does not explicitly optimize the lookup table. Instead, the transition process in the finite-state Markov model M can be approximated using neural networks with discrete input, such as a Multi-Layer Perceptron network. Once the training is complete, the lookup table can be generated by feeding all possible state combinations into the network.
Learnable State Number We notice that the size of the lookup table O(C.)) can become pro- hibitively large, especially when N 2 2. A common practice to optimize C' is through hyperparameter tuning. However, this tuning process
consuminn as requires training the full model multiple times to determine the optimal hyperparameter. To mitigate this issue, we propose an end-to-end adaptive optimization method for selecting C' called Learnable State Number (LSN).
A straightforward approach to adjusting the state number involves applying a masking function to the state probabilities, preventing certain states from being sampled. Additionally, a trainable parameter l controlling the sparsity of the mask can be introduced. To facilitate gradient descent, the mask should be a continuous function of 0. Here, we adopt a-entmax (Peters et al. 2019). a normalizing function similar to the softmax function, that allows the generation of precisely zero probabilities and enables sparse masking of states. It can be defined as:
entmax O0, a) = [(a 1)0
where T is constant that ensures Σ entmax (0, a) 1. The state probabilities, denoted as π,, are then modified by the a-entmax function as shown below:
π (onimax(O,c))
0 e RC
E π ent
Note that a can be manually set to control the reduction rate or set as trainable parameter following (Correia et al., 2019) for more flexible masking during training.
Discussion for FSAR as Latent Prior When compressing natural data vector æ, it is common for dependencies to exist among the variables in æ. Markov models are often employed to simplify these dependencies, assuming that each variable depends only on a few preceding variables. However, even with Markov models, decoding each p(r;| N:-1) still requires the previous N' variables, resulting in a sequential iteration process for every element of æ, which can be time-consuming. Alternatively, latent space models, such as autoencoders, x. By applying Markov models in the latent space, the sequential iteration process can be shortened. Although disentanglement of all latent variables is unattainable under unsupervised optimization targets for lossless compression (Locatello et al. 2208)), the inference network in autoencoders can still disentangle certain dependencies among variables, making compact Markov model sufficient as a prior for the latent space. This approach provides appreciable improvements for compression ratios while preventing overfitting. Consequently, we employ the proposed FSAR model with a low order (typically N < 2) as the latent prior.
transform into a smaller latent vector
Relationship Masked Convolution (van den Oord et al. 2016b) Masked convolution can also be exploited as an autoregressive prior in the discrete latent space. In fact, the masked convolution- autoregressive model essentially a subset of the finite-state Markov model. However, constructing the lookup table becomes nearly impossible due to the high order in most cases. For instance, a 3x3 2D
5

---

Published as a conference paper at ICLR 2024
masked convolution- based autoregressive model corresponds an Order-4 finite-state Markov model in the discrete latent space, denoted as J.s............ 1. - Vij-1.A: Vi.j-1A-1), where i. represent the indices of the channel, height, and width dimensions. respectively. Generating the lookup table for this model requires a space complexity of ((C*), which is generally unacceptable. Furthermore, mask convolution usually requires FLOPs, which can introduce potential inconsistencies in computation between the compression and decompression stages Ballé et al. (2019). On the contrary, the proposed method immune to such inconsistencies it relies solely on memory operations for latent coding.
Relationship ANS (Duda) 2009) Recent implementations of entropy coders have embraced ANS (Duda, (009). as shown in Figure 3al In the i-th step, ANS utilizes the previous state s;-1- the current symbol yt. and the corresponding prior distribution p to determine the state transition process. When the inputs of the state transition process, namely 8;-1 and yi, are drawn from finite set, techniques such as tANS (Duda 2013) can be employed to precalculate all possible inputs and construct lookup table for efficient state transition computation. This concept inspires us to devise similar scheme for coding in the discrete latent space. In fact, the proposed FSAR can seamlessly integrate with tANS, where state transition table is derived from the FSAR lookup table L at each step (see Figure 3b). For a detailed implementation of FSAR-tANS and a discussion on the time complexity of related methods, we direct the readers to Appendix C
5 STRAIGHT THROUGH HARDMAX QUANTIZATION
As mentioned in Section | optimizing discrete latent space models is complicated despite their preference for efficient latent coding. In this section, we examine the limitations of existing methods for optimizing vector quantized (VQ) discrete latent space models and propose robust quantization method called Straight Through Hardmax Quantization (STHQ).
Recall that the original VQ-VAE (van den Oord et al.
2017) is optimized with E-meenns-lkke loss: LvQ = || (x=) - By la + /1 ((x) sg (yy Byd))lz - log por By"),
(4) where ((x) is the output of the inference network, B represents the codebook, y/ denotes the one-hot vector that selects the closest codeword to I [xx)
B, and sg represents the stop gradient operator. However, updating all codewords in each step is infeasible due to the one-hot selection of a single codeword by y. resulting zero gradients for other codewords in the codebook (Sønderby 2017 Roy et al. 2018).
To overcome this limitation, the relaxed VQ (RVQ) approach (Sønderby 2017) introduces a cate- gorical distribution ((z(<) (z) in the latent space. The logits of this distribution are determined by the Euclidean distance between the latent variable and the codewords, that is,
q= x) =
(5)
where Di(((=)) = (z) - B:I12 is the euclidean distance between I [(x) and i-th codeword vector
B, The entire model is then optimized using the negative evidence lower bound (ELBO) formulation:
LavQ - (g(2 (=)) - (log p(2)) log p(=|=),
(6)
where (g((zaa) (z)) is the entropy of q(2|=)) The inclusion of the entropy term facilitates global optimization of all distances D, encouraging close latent variables and codewords to be grouped together. This formulation has been adopted by recent works such as SQ-VAE (Takida et al. 2022).
However, there are two potential issues immediately arising when optimizing the entropy term in RVQ. Firstly, the partial derivative of the entropy term with respect to the distance D; can be negative for all which could be derived from Proposition 5.1 As consequence, the distance may continually increase during the optimization process and cause the codebook to diverge.
Proposition 5.1. It is possible that Vi, OH (g(l)))))) < O.
Proof. See Appendix D.I for proof.
Secondly, the inclusion of the entropy term H (1(2(=)) creates an entropy gap between VQ and RVQ. resulting in an inferior performance of RVQ in generative tasks, as proven in Proposition 5.2
Proposition 5.2. Assume p(2) follows a uniform distribution and give the same sample input y and parameters 0 for the generative model log pe po(x)By")) we have log Pang min, Lxva (x) By"") log Pargming Eva By"".
Proof. See Appendix D.2 for proof and see Appendix D.4 for the corresponding experiments.

---

Published as a conference paper at ICLR 2024
In order to avoid these problems and enable global updates similar to RVQ, we propose STHQ.
Straight Through Hardmax Quantization Inspired by the straight through Gumbel-softmax method (Jang et al.] 2016) which employs categoricat one-hot samples in the forward process but Gumbel-softmax based soft samples in the backward process (backpropagation), we propose a similar gradient estimator called straight through hardmax (st-hardmax). This approach incorporates ahiaapprrrd
a VQ-based hardmax in the forward process while using sampling in the backward process:
y = st-hardmax( ( -(((=))) ++ y = hardmax( - D(T(2))), ( By ST d (-D,(=)(((()))
(7)
aD where GS represents the Gumbel-softma reparameterizatioo (Jang al. 2016) with With temperature T. In this manner, the forward process remains the same as VQ, thereby mitigating the entropy gap issue. Meanwhile, the backward process allows for the global update of z and all B; based on their respective distances D,. Consequently, we can use simple regularization loss term || [|Dmin m similar to VQ,, to update all distances Di:
0D
la.
|| = || arg min Dill2 = (1) - Byy
(8)
By combining the st- -tadddd operator and this 1.2 regularization term, the overall loss function of STHQ can be expressed as:
LSTIQ= = ||D 122 - log p(ær By") - B,y log)(z).
Note that in VQ, p(z) is typically assumed to follow uniform distribution and is omitted from the loss function. However, in STHQ, we adopt the proposed FSAR model as p(z) for better compression. To learn the parameters of the FSAR model, p(z) needs to be included in Eq. (9). comeene
Specifically, following previous VQ methods. we only utilize E. log p(z) for optimizing the FSAR model, while excluding this term when optimizing the backbone network and the codebook.
Relationship to Other VQ-based Methods By assuming p(z) to be uniform in STHQ (9) and ignoring the stop-gradient operations in VQ-VAE (4), we can find that the loss functions in these two methods are identical. It is worth noting that both VQ-VAE and STHQ apply the straight-through e
estimator on the likelihood term log p(x Byy") to enable gradient calculation on the inference model. However, in STHQ, we also apply the ttraight-through estimator to the L2 regularizer term || [mmin ll2. which benefits from the global update property and leads faster convergence. On the other hand, when compared to RVQ- methods such as (Senderby) 2017 Takida et al.) 2022). STHQ omits the entropy term. This term can potentially hinder the optimization process, as indicated in the above two propositions. More details are presented in Appendix D.3
above
6 EXPERIMENTS
6.1 DATASETS AND METRICS
In our experiments, we focused on compressing and decompressing image datasets, specifically CIFARIO (CF10) (Krizhevsky) 2009) and ImageNet32/64 (IN32, IN64) (Deng et al., 2009). We evaluated performance using four criteria: Bits Per Dimension (BPD) for compression ratio, compres- sion speed (CSpd) and decompression speed (DSpd) measured in megabytes per second (MB/s) to assess time complexity, and occupied memory (Mem) in megabytes (MB) to assess space complexity. The measured time or speed was obtained running on CPUs, as our method targets general-purpose computation devices. The detailed experimental setup can be found in Appendix F.
6.2 CAMMAAANN
WITH OTHER LATENT CODING METHODS
To demonstrate the effectiveness and efficiency of the proposed FSAR model as prior for latent modeling, we conducted a comparative analysis with various other latent space coding methods. The methods compared include commonly used continuous latent models with bits-back coding (Townsend et al., 2019a). discrete latent models such as Vector Quantization (VQ) (van den Oord et al. 2017). /. Universal Quantization (UQ) (Ballé et al.) 2018), and McQuic (Zhu et al.] (022). as well as autoregressive-bassee based latent models such as Masked Convolution (MaskConv) (van den Oord et al., (016b) and Checkerboard context model (Checkerboard) (He et al.) 202T). Our proposed FSAR model ancompassed Markov models of different orders, ranging from Order-1 to Order-3. Additionally, we consider the Order-2 model with LSN to evaluate its usefulness. We utilized the proposed STHQ training for all autoregressive methods (MaskConv2D, Checkerboard, FSAR(O1), FSAR(O2), FSAR(O2))LSN, FSAR(O3)). while also conducting an additional experiment in the non- autoregressive latent space (i.c., only LANS (Duda] 2013) entropy coding) as a reference. Furthermore, to Mutorregoe
provide detailed efficiency comparisons, we presented the average decompression time of the conaaatooomaaii
7

---

Published as a conference paper at ICLR 2024
Table I: Experiment results of different latent coding methods. The detailed running time for each module in different methods is obtained by processing a single image from CIFARIO dataset. Order-1, Order-2 and Order-3 denoted as O1, 02. and O3, respectively. For BPD results. bold numbers indicate the best practical BPD for each dataset, "OOM" indicates the out-of-memory error during coding. "ERR" indicates inconsistency between the decompressed data and the original data, and "NAN" indicates failed training process.
BPD (Tecoiiical/Paactaaal)
CF10 Decomp Time (ms)
Mem (MB)
Methods
CF10
IN32
IN64
Net
Coder
Total
CF10
Contimuous
bits-back
Towasend
2019a
4,99/8.99
5:45/9.11
5:05/11.15
32:44
9.85
42.29
0.000
UQ Balls et al
2018
5.16/5,19 19
5.65/5.68
5,36/538
1.60
1.54
3.15
0.015
Discrete
VQ van den Dord et al. 2017
5.505.52
6.16/6.17
5.61/5.62
1.57
0.88
2.45
0.008
MoQuii 2hu et al
2022
5.81/5.00
6.19/9.57
NANINAN
1.57
8:75
10:32
0.082
MaskConv2D van den Oord cl
2016b
5.02/ERR
5/42/RR
5.08/RR
1.75
632.24
633,99
1.127
Autorogressive
Checkerboard (He al.
20211
5.11/5.11
5,49/5.50
5.23/5.33
1,63
4.71
6.34
1,127
IANS (Duda 2013
5,25/5.25
5.72/5.72
5.46/5.48
1.66
0.86
2.52
0.008
PSAR(O1)
5.14/5.17
5.54/5.55
5.165. 18
1.58
0.97
2.55
2,010
Proposed
FSAR(O2)
5.05/5.06
5,4515,46
5,0515.07
1.62
1.04
2.65
516.512
FSAR/O2+LSN FSAR(O3)
5.08/5.09
5.43/5.44
5.10/5.12
1.59
1.06
2.65
276.586
5.00/00M
5.41/0OM
5.08/DOM
00M
OOM
0OM
OOM
backbone networks (Net) and the entropy coders (Coder) for processing a single image in CIFARIO. The results are shown in Table I Compression time are further provided in Appendix EE. for reference. It is apparent that the proposed method achieves comparable or superior performance to the existing methods in terms of both compression ratio (i.e., theoretical/practical BPD) and speed (i.e., run time).
Concretely, in terms of the practical BPD, the proposed FSAR-based methods achieve the best performance on all datasets. Moreover, the theoretical BPD closely aligns with the practical one for all methods, except for bits-back and MaskConv. This discrepancy arises from the initial bits issue encountered in bits-back (Flamich et al. 2020) and the floating-poin inconsistency issue observed in MaskConv (Ballé et al. 2019). Comparing FSAR models of different orders, we find that the order-2 Markov model outperforms the order-1 I model in terms of compression performance, while the order-3 model does not provide a significant improvement over the order-2 model. This observation implies that a compact Markov model can adequately describe the latent space, as the backbone model capable of disentangling dependencies within the data to some extent, resulting in simplified dependency structure in the latent space. Furthermore, considering the weeorry-itennsii nature of order-3 models, the order-2 Markov model is the optimal choice for FSAR
In terms of speed, the bits-back and MaskCony methods are notably slow due to their high time complexity in the backbone network and entropy coder, respectively. The Checkerboard context model employs parallelized computation during latent decoding, resulting in significantly shorter decoding time. However, Checkerboard remains slower than the non-autoregressive IANS, whereas the FSAR model performs as fast as tANS. In addition, the increased time cost of accessing lookup tables in FSAR, which scales with memory size, has minimal impact on the overall decompression time, as the time is dominated by the FLOPs in the backbone networks.
Note that memory occupation is a concern for FSAR models. Although the proposed LSN method successfully reduces the memory requirement of FSAR(O2)) by approximately 50% without causing noticeable changes in compression performance, the memory occupation remains relatively high.
6.3 COMPARISON WITH OTHER COMPRESSION METHODS
In this section, we evaluate the proposed codec by comparing it with state-of state-of-the-arr image compres- sion methods in general setting. No restrictions are imposed on the network backbone or training schedule in this experiment. The benchmark methods include traditional image lossless compression
codecs such as PNG (Boutell 1997). WebP (webmproject, (023) and FLIF (Sneyers & Wuille| 2016). as well as autoencoder based methods including L3C (Mentzer et al. 2018). ,). bitswap (Kingma et al. 2019), and PILC (Kang et al. 2022 Since the proposed method solely employs latent space models and does not restrict the method of observation space data coding, it can be seamlessly combined with PILC, which adopts autoregressive models and entropy coding in the observation space, by replacing the VQ-VAE used in PILC with the proposed method. We calculate the BPD based on the length of the compression bitstream to evaluate practical compression performance. For qualitative understanding of the results, present the BPD versus Speed comparison for different methods on CIFARIO in Figure 1. For complete numerical results, please refer to Appendix E.2
Generally, the proposed learned codec has similar BPD and speed with traditional WebP and FLIF, thereby bridging the gap between learned and traditional methods. Compared to the autoencoder-
8

---

Published as a conference paper at ICLR 2024
1E42
1E42
PILC. PILC+STHO,
PILC, 4.26 <WebR 4.52
PNG. 5.78
1F41
PILCYSTHQ+
PNG.
1F41
PILC-STHQ++
PILCYSTHQ,
CSPO (M//S)
10+0
FSAR,
WebP, 4.52
1610
FSAR, 4.18
4.29
4.29
FLIE 4.29
16-1
LBC, 4.77
DSPO
1C-1
16-2
bitswap,
3.84
16-2
bitswap,
3.84
Autoencoder
Traditional
Proposed
16.3
16.3
3.50
4.00
4.50
4.00
4.50
6.00
BPD
BPD
Figure 4: The BPD versus Speed (CSpd and Dspd) comparison for different methods on CIFARIO.
Table 2: BPD results (with standard deviation) for the ablation studies.
BPD
Saved
BPD
Saved
||
BPD
Saved
VQ van den Oord et al. 2017
5.37(0.02)
0.00%
STHQ
5.23(0.02)
-2.66%
-STH
5.54(0.08)
3.16%
RVO Spnnerby 2017
5.68(0.10)
5.73%
+FSAR(O1)
5.15(0.03)
-4.6%
-12
6.58(0.03)
22.36%
Takida al 2022
5.50(0.1))
2.42%
+FSAR(O2)
5.06(0.02)
-5.80%
+Entropy
5.34(0.06)
-0.67%
based method PII.C the improved version of the proposed method, called "PILC+ "PILC+STHQ+*PARR reduces BPD by 2.2-2.5% but also slightly decreases CSpd by 7% and DSpd by 5%. The reduced BPD can be attributed to the incorporation of the FSAR model. Additionally, the compression performance of our method is similar to that of the deterministic autoencoder-basse method L3C (Mentzer et al. 2018), but is inferior to the variational autoencoder- based method bitswap (Kingma et al. 2019). However, in terms of speed, our method outperforms both L3C and bitswap by large margin, owing the shallow backbone and the efficient entropy coder. It is worth noting that bitswap requires initial bits, rendering it unsuitable for single or small-batch image compression.
6.4 ABLATION STUDIES
Here, we conduct ablation studies to examine the role of each component the proposed method. The theoretical BPD is reported, as previous experiments have shown a negligible disparity between theoretical and practical compression performance. Our baseline is the widely-used VQ-VAE with a discrete latent space ( (van den Oord et al.) 2017). We compare it with RVQ- VAE (Sønderby (017). SQ-VAE (Takida et al., (022)). and the proposed STHQ-based methods. For a fair comparison, we disable the self-annealing mechanism in SQ- VAE for the sake of lossless compression. Different variants of STHQ are considered, including the one with the FSAR prior (+FSAR), without the st-hardmax operator (-STH), without the L2 regularization term [|Dmminll (-L2), and with an entropy term similar to RVQ-VAE (+Entropy). The experiments are performed on CIFARI0 images using the same backbone network and sizes of latent variables and codebooks. The results presented in Table 2 For additional results from the ablation study, including different settings of the backbone network, latent space, and codebook size, we refer the readers to Appendix E.3 E.5
It is obvious that the proposed method surpasses all variants. STHQ reduces the BPD by about 2% compared to VQ-VAE, and the inclusion of FSAR yields an additional 2-4% improvement. RVQ- VAF and SQ-VAE VAE perform worse than VQ-VAE, probably due to the entropy gap discussed in Section The higher variance observed in RVQ- and SQ-VAE suggests instability, which can be explained by the codebook divergence (cf. Proposition 5.1). Regarding the variants of STHQ, we observe that the L2 regularization term achieves 20% BPD decrease by supporting the simultaneous update of all codewords. Moreover, without the st-hardmax operator, the distance updates in STHQ resemble those of SQ- VAE, and so both STHQ by 6%. Recall that SQ-VAE is unsuitable for compression purposes due to its reliance on likelihood degeneration obtain deterministic decoder (see Section 2). In addition, introducing the entropy term can cause instability in the optimization process (cf. |Proposition 5.1). thus degrading STHQ's performance by about 2%. In summary, the st-hardmax operator, L2 regularization term, and exclusion of the entropy term are all vital for the robust optimization of the discrete latent space.
and SQ- VAE perform similarly but worse than
CONCLUSIONS
This paper presents a novel framework for efficient learned lossless compression, which incorporates FSAR for efficient autoregressive latent coding and improved compression, together with STHQ for robust optimization of the discrete latent space and further enhancement of compression performance. Experimental results demonstrate substantial improvement over the baseline, with a 6% increase in compression ratio while maintaining comparable compression and decompression speeds. These findings highlight the potential of approach for advancing the field of lossless compression.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The work was supported by Ant Group and the National Natural Science Foundation of China (No. 62325109, U2IB2013).
REFERENCES
Workshop and challenge on learned image compression. ht.p:///li...copresssio.. 2023. Accessed: 2023-01-08
Kodak lossless true color suite. 2023-11-12.
2023. Accessed:
Eirikur Agustsson, Fabian Mentzer, Michael Tschannen, Lukas Cavigelli, Radu Timofte, Luca Benini, and Lue Van Gool. Soft-to-hard vector quantization for end-to-end learning compressible representations. In NIPS, 2017.
Martín Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan. ArXiv, 2017.
Johannes Ballé, David C. Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston. Variational image compression with scale hyperprior. ArXiv, abs/ 100.0013244 2018
Johannes Ballé, Nick Johnston, and David Minnen. Integer networks for data compression with latent-variable models. In International Conference Learning Representations, 2019. URL
wwt. Azg/CorpusID: 10830082200
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. ArXiv, abs/1308.3433 2013.
Thomas Boutell. Png (portable network graphies) specification version 1.0. RFC, 2083:1-102, 1997.
Yuri Burda, Roger Baker Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. 2015.
CoRR,
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and J. Katto. Learned image compression with discretized gaussian mixture likelihoods and attention modules. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9936-7945, 2020.
Gonçalo M. Correia, Vlad Niculae, and André F. T. Martins. Adaptively sparse transformers. In Conference on Empirical Methods in Natural Language Processing, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, K. Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp.
2009.
Laurent Dinh, Jascha Narain Sohl-Dickster and Samy Bengio. Density estimation using real nvp. ArXiv, abs/ (B0116050880. 2016.
Jarek Duda. Asymmetric numeral systems. ArXiv, abs/0902.0271, 2009.
Jarek Duda. Asymmetric numeral systems: entropy coding combining speed of huffman coding with compression rate of arithmetic coding. arXiv: Information Theory, 2013. URL https
ID:134994:
Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis. 202/ IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp- 12878, 2020.
Gergely Flamich, Marton Havasi, and José Miguel Hernández-Lobato. Compressing images by encoding their latent representations with relative entropy coding. ArXiv, abs/2010.01 185, 2020.
Gergely Flamich, Stratis Markou, and Jos'e e Miguel Hern andezo Fast relative entropy coding with a coding. ArXiv, 00.02201111887, 2022.
10

---

Published as a conference paper at ICLR 2024
Lina Guo, Xinjie Shi, Dailan He, Yuanyuan Wang, Rui Ma, Hongwei Qin, and Yan Wang. Practical learned lossless jpeg recompression with multi-level cross-channel entropy model in the det domain. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5852-5861 2022. URL httph::tttsaapt. semant cscholar org/ Corpus ID:
Lina Guo, Yuanyuan Wang, Tongda Xu, Jixiang Luo. Dailan He, Zhen fang Ji, Shanshan Wang, Yang Wang, and Hongwei Qin. Efficient learned lossless jpeg recompression. 2023. URL
https: //api semanticcchhee ar www.ola.r..rooooeeeeeeoeeeeeeeee wwt. Corpus 266214335
Dailan He, Yaoyan Zheng, Baochen Sun, Yan Wang, and Hongwei Qin. Checkerboard context model for efficient learned image compression. 202/ IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 14766-14775, 2021.
Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2015.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. In Neural Informa- tion Processing Systems, 2017. URL httpt::////// /apsi. semant iescholccccarrr org/CorpusID;
326772
Emiel Hoogeboom, Jom W. T. Peters, Rianne van den Berg, and Max Welling, Integer discrete flows and lossless compression. ArXiv, abs/ /bss/199500776, 2019.
Eric Jang. Shixiang Shane Gu, and Ben Poole. Categorical reparameterization with gumbel-softmal ArXiv, 1/5/1611.01144 2016.
Ning Kang, Shanzhao Qiu, Shifeng Zhang, Zhenguo Li, and Shutao Xia. Pilc: Practical image lossless compression with an end-to-end gpu oriented neural framework. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3729- 3738, 2022.
Diederik P. Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. ArXiv, Mts//8877.00399 2018
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. CoRR, abs/1312.61 14, 2013.
Friso H. Kingma, Abbeel, and Jonathan Ho. Bit-swap: Recursive bits-back coding for lossless compression with hierarchical latent variables. In International Conference on Machine Learning. 2019.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Alina Kuznetsova, Hassan Rom, Neil Gordon Alldrin, Jasper R. R. Uijlings, Ivan Krasin, Jordi Pont-Tuset, Shahab Kamali, Stefan Popov, Matteo Malloci, Alexander Kolesnikov, Tom Duerig. and Vittorio Ferrari. The open images dataset v4. International Journal of Computer Vision, 128: 1956-1981, 2018.
Jooyoung Lee, Seunghyun Cho, and Seunghwa Beack. Context-adaptive entropy model for end- to-end optimized image compression. In International Conference on Learning Representations, 2018.
Fangzheng Lin, Heming Sun, Jinming Liu, and J. Katto. Multistage spatial context models for learned image compression. ArXiv, abs/2302.00263,, 2023. Utthttts::/aap. semanticscho Btt/CorpuuID:2550388
Francesco Locatello, Stefan Bauer, Mario Lucic, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled represen- tations. ArXiv, abs/181 1. abs/1811.12259, 2018.
Fabian Mentzer, Eirikur Agustsson, Michael Tschannen, Radu Timofte, and Luc Van Gool. Practical full resolution learned lossless image compression. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10621-10630, 2018.
11

---

Published as a conference paper at ICLR 2024
David C. Minnen, Johannes Ballé, and George Toderici. Joint autoregressive and hierarchical priors for learned image compression. ArXiv, abs/1809.027333 2018.
Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. ArXiv, abs/1411.170 2014.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam M. Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning. 2018.
Ben Peters, Vlad Niculae, and André F. T. Martins. Sparse sequence-to-sequence models. ArXiv, abs/1905.05702. 2019.
Boris Polyak and Anatoli B. Juditsky. Accelcration stochastic approximation by averag- ing. Siam Journal Control and Optimization, 30:838 855, 1992. URL https:/ //ap
wrtpaanicccccoooaooo/roorrorree 3548822 3548228
Ali Razavi, Aäron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. ArXiv, abs/1906.000466, 2019.
Kenneth M. Rose. Deterministic annealing for clustering, compression, classification, regression, and related optimization problems. Proc. IEEE, 86:2210-222 22299, 1998.
Aurko Roy, Ashish Vaswani, Niki Parmar, and Arvind Neelakantan. Towards a better understanding of vector quantized autoencoders. 2018
Yangjun Ruan, Karen Ullrich, Daniel de Souza Severo, James Townsend, Ashish Khisti, A. Doucet, Alireza Makhzani, and Chris J. Maddison. Improving lossless compression rates via monte carlo bits-back coding. In International Conference on Machine Learning, 2021.
Tom Ryder, Chen Zhang, Ning Kang. and Shifeng Zhang. Split hierarchical variational compression. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 386-395, 2022.
Claude Shannon. A mathematical theory of communication. Bell Syst. Tech. J. 27:623-6566 6666,,666666 1948.
Jon Sneyers and Pieter Wuille. Flif: Free lossless image format based on maniac compression. 2016
IEEE International Conference on Image Processing (ICIP), pp. 66-70, 2016.
Casper Kaae Sønderby. Continuous relaxation training of discrete latent variable image models. 2017.
Yuhta Takida, Takashi Shibuya, Wei-Hsiang Liao, Chieh-Hsin Lai, Junki Ohmura, Toshimitsu Uesaka, Naoki Murata, Shusuke Takahashi, Toshiyuki Kumakura, and Yuki Mitsufuji. Sq-vae: Variational bayes on discrete representation with aenellnd stochastic quantization. ArXiv, abs/2205.07547, 2022.
Lucas Theis and Jonathan Ho. Importance weighted compression. In Neural Compression: From Information Theory to amplications-Workkhheen ICLR 2021, 2021.
James Townsend, Thomas Bird, and David Barber. Practical lossless compression with latent variables using bits back coding. ArXiv, abs/1901 10486,, 2019a.
James Townsend, Thomas Bird, Julius Kunze, and David Barber. Hilloc: Lossless image compression with hierarchical latent variable models. ArXiv, abs/1912.09953, 2019b.
Rianne van den Berg, Alexey A. Gritsenko, Mostafa Dehghani, Casper Kaae Sønderby, and Tim Salimans. Idf++: Analyzing and improving integer discrete flows for lossless compression. ArXiv, abs/2006. 12459, 2020. URL httph:tt////pp/t remanticscho ar. wwt..rr//rooooo/O00 219965949
Aäron den Oord, Nal Kalchbrenner, Lasse Espeholt, Koray Kavukcuoglu, Oriol Vinyals, and Alex Graves. Conditional image generation with pixelenn decoders. ArXiv, abs/1606.05328. 2016a.
URL https::///........... //api./pp.... www.nticscooooe..ooo.. wttp:aar...rrrooouaaaeeeeeeoeee
149899993
12

---

Published as a conference paper at ICLR 2024
Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. ArXiv, abs/1601.06759, 2016b.
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In NIPS, 2017.
webmproject. libwebp. httpa//s/gttttooooooo/////pppppp 2023. Accessed: 2023-01-09.
Feiyu Yang. Luis Herranz, Yongmei Cheng, and Mikhail G. Mozerov. Slimmable compres- sive autoencoders for practical neural image compression. 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4996-5000 2021. URL https:
i
//api. htt.://et.cccccoo.....oooooooooa
Jiahui Yu, L. Yang, N. Xu, Jianchao Yang, and Thomas S. Huang. Slimmable neural networks.
ArXiv, abs/1812.089288 2018. URL http:::tt//att /appi. semant somanticscholasr htt.aissssooooooooooooossss.saa. 56657799
Xiaosu Zhu, Jingkuan Song, Lianli Gao, Fengcai Zheng, and Hengtao Shen. Unified multivariate gaussian mixture for efficient neural image compression. 2022 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 17591-17600, 2022.
13

---

Published as a conference paper at ICLR 2024
A FURTHER REVIEW ON RELATED WORKS
Continuous Random Latent Space Variational autoencoders (VAE) with continuous random latent space (Kingma & Welling 2013) are commonly employed in learned lossless compression due to their ease of optimization using the Evidence Lower Bound (ELBO). However, the practical encoding of their latent space poses challenge, as the random samples are continuous and cannot be directly processed by a simple entropy coder. To tackle this issue, previous approaches have proposed two main solutions: bits-back based latent coding (Townsend et al. 2019a; Kingma et al. 2019; Theis & Ho 202T; Ruan et al.] 2021) and relative entropy coding (Flamich et al. [2020] 2022). Bits-back ANS (Townsend ct al. 2019a) suggests taking random samples directly from the compressed stream. Subsequent works have extended or optimized bits-back for different types of VAE variants, including hierarchical VAEs (Townsend et al. 2019b Kingma et al. 2019). importance weighted autoencoders (Theis & Ho 2021), and Monte- Carlo VAES (Ruan et al. 20211) I). Another approach, relative entropy coding (Flamich et al. 2020 (022), incorporates a shared random number generator between the compressor and decompressor, compressing the sampled index represent latent samples. However, both methods still require additional bits for the random source, and the implementation of the entropy coder remains complex.
B BACKGROUND KNOWLEDGE FOR SECTION 3
Entropy Coding with ANS Entropy coding plays a central role in most compression frameworks, as it compresses messages sampled from a probability distribution into a bitstream with length close to its entropy (Shanmon 1948). expressed as EEp log, (xx). For a single sample from P(æ:), the optimal code length is loga2 P(x). A recently proposed near-optimal entropy coder, ANS (Duda 2009), treats the bitstream as a growing state. It pushes and pops values to and from this state based on their probabilities. One of the ANS implementations, tabled ANS, is illustrated in Figure 3a ANS achieves the optimal code length- log, P(x) for any given samples from P(x))
Compression with Variational Autoencoders The proposed framework draws inspiration from previous compression frameworks based on variational autoencoders. A variational autoencoder models the probability of the input data p(x:) by transforming the likelihood p( 2) based on z. Consequently, x can be compressed or decompressed using p(z|z) by applying an entropy coder. The latent variables follow random distribution conditioned on æ, denoted as (()))). As z is required by p(æ) =), z also needs to be entropy coded with p(z). The entire process of modeling p(æ)) using a variational autoencoder can be represented as:
latent variables and generating
p(z) p(=))(=)))
(10)
q(
The length of the compressed stream produced by entropy coders can be estimated using E. p(z), which can be rewritten
E, log p(z) = Egg log q [z) - log p(z) log p()p(()))))))))))))))(()))))(p [L(((>)=)|l(()) -og logp(z)z). (11)
Here, log p(==)) estimates the compressed length of x. and ((g(\\>)\lp<=)) estimates the compressed length of z. It's important to note that common entropy coders can only process samples rather than distributions. Hence, the true compressed length should satisfy the inequality given by:
- -log p(2) log p(rl2) 2 Feg log ((2 |ær) - log p(z) - log p(gg)(=()) = E. log p(x). (12) As q(>(>) is a discrete distribution, equality is achieved only if q(z|x) is a one-hot distribution. Therefore, we can straightforwardly set q(2 (æz) as categorical distribution with one-hot probabilities, which can be easily parameterized using a quantization function.
As shown in Figure 2] the inference network provides parameters for g(2(z), while the generative network implements p(æ| (x). The one-hot representation of q(>(<²) is obtained by quantizing the output of the inference network using a learnable codebook, employing the proposed STHQ method. To be specific, we calculate the euclidean distances between the inference network output and each codeword in the learnable codebook. The negative euclidean distances are then utilized to determine the sampling probability of each codeword, following similar approach as in (Sonderby 2017). Additionally, we incorporate the sparse mask from the LSN module in FSAR to further prune the codebook. The proposed STHQ produces one-hot latent samples by selecting the codeword with the highest probability through quantization. These one-hot samples can be converted into latent indices for compression into the latent stream and are also used in the dequantization process, where
14

---

Published as a conference paper at ICLR 2024
specific codeword is chosen from the codebook the input for the generative network. p(z) is effectively encoded using an efficient autoregressive entropy coder based on the proposed FSAR model. The optimization process compressed bitstream and achieve superior compression ratios.
Eq. (12) the loss function minimize the length of the
C DISCUSSIONS FOR SECTION
C.I TIME COMPLEXITY COMPARISON
Comparison of Complexity with ANS The utilization of lookup tables has been explored in table ANS (tANS) implementation (Duda 2013). This approach precalculates all possible state transitions and utilizes lookup tables during the coding process expedite the process. Given that autoregressive Markov model can iterated alongside the entropy coder, it straightforward to design similar lookup table- base- autoregressive Markov model for acceleration. To compare the complexity, we provide a breakdown of the number of operations for different ANS implementations. including range ANS (rANS) (Duda 2013) and table ANS (tANS) (Duda 2013), shown in Table evident that the proposed FSAR model requires only a few additional memory operations compared to ANS, thanks to its lookup table-based iteration process.
Comparison of Complexity with Masked Convolution based Autoregressive Models Masked convolution (van den Oord al. 2016b). which relies FLOPs, is commonly employed for implementing autoregressive models. In Table 4] we provide a breakdown of the number of operations per iteration required for masked convolution-based autoregressive models. In contrast, the proposed FSAR model completely avoids expensive FLOPs and exhibits complexity comparable to that of entropy coders, shown in Table
Table 3: Integer, floating-point, and memory operations per symbol in modern entropy coders. Operations are distinguished by color-coded boxes in the accompanying pseudo-code.
PseudoCode
Int
Float
Mem
start, freq
CDF
[x]]
rANS Encode
$ (8:-1
freq) nbit + 8
%
freq start;
8, flushbits (s);
efreq S & bitmask:
V4. start, freq
ICDF [efreq]
rANS Decode
6
0
freq
(8
nbit) cfreq start;
readbits (s)
nbit, delta TransTable [yd]
tANS Encode
flushbits (st-1);
0
8: StateTable [(>-_1 ubit) + delta]
nbii, DecTable [s--1]
tANS Decode
0
readbits (s)
FSAR(O2)+ tANS Encode
TransTable StateTable Table [3j-1] [y]-1]
3
0
4
IANSEncode (TransTable, StateTable);
FSAR(O2)+ ANS Decode
DecTable = Table [y:-1] [g)--1]
1
0
3
IANSDecode (Dec'Table);
15

---

Published as a conference paper at ICLR 2024
Table 4: Integer. floating-Poine Point, and memory operations per symbol in various autoregressive models. Oper- ations are distinguished by color-coded boxes the accompanying formula or pseudo-code. Convolutional layers are assumed to have a single channel. Ma, and Ca, b represent kernel sets, defined as Ma. b = (5, k)]vj, k E 2, -a < < O, [k] bu (0, k)/kk < 0 and Ca, h = (i, k) E 2, (i + k) mod 2 # 0, re- spectively.
Formula/PseudoCode
Int
Float
Mem
Mas3Conv2D3x3
(m, M.+n m.
ma,
2 M122 =
MasSConv-IDSx5
Σ ( (m. Dj Uj+m.kk k t W m.n
2/M,,,al 24
0
Checkerboard2D3333
Σ min
(m,n)ECo, Mjtm,ktn W mn
2|C2,21 =
0
FSAR(OI)
Table [ 1]
0
1
FSAR(O2)
Table [y/ 11 ykk-1
0 2
C.2 IMPLEMENTATION OF
In theory, FSAR has the potential to be combined with any entropy coders. However, for practical implementation efficiency, we have opted for tANS (Duda) 2013), which leverages lookup tables for accelerated processing. The algorithm for finite-state autoregressive entropy coding based on LANS, including the initialization, encoding, and decoding processes, is illustrated in Algorrihh The modified or added steps. which differ from the original LANS algorithm, are highlighted in red. It is obvious that the only additional step during encoding and decoding is the n-dimensional table lookup.
D DISCUSSIONS FOR SECTION 5
D.1 PROOF OF PROPOSITION 5.1
Proposition 5.1 It is possible that Vi, 0H (H(((()))) 0
Proof.
He OH (1(2[=))
0 0E 9: log q: Ög:
(13)
OD,
dq.
aD;
Note that
q: = softmax, (-D) = e D.
(14)
E. e-D
D:
Let S; e D. (0, 1].
we have
θ E, (qi log q)
+ log + log S: - log Σ S,
(15)
0y:
and
dq:
aS, 0g: OD, as
S E3 S, - S
(16)
0D,
E, S,)
Therefore,
(H((zz))) 3D,
E, S. S [, S.)²
S,( + log S: log Σ S.)
(17)
Since
S. Σ 5, S
>0,
(18)
[, S.)2
16

---

Published as a conference paper at ICLR 2024
Algorithm Finite-State Autoregressive tANS Coding Algorithm. The steps that are different from tANS are highlighted in red.
procedure INITALLIZATIONN M, NumStates C, Order n.)
Initialize n-dimensional Lookup Table Lg, LD for yı {0, ..., CJ,..., V/ E {0, 1, C} do
π + softmaax( (My),, Un))
Lg // Vnl BuildT uns EncodeT (bpp(p
LD 1/1 Vnl -- BuildT ans DecodeTTable(m
end for return LE, LD end procedure
procedure (Syools y. EncoderLLookuTTble Lg, Offsets !o.,,.,,,,)) Initialize ANSState BitStream B for E reserve(y) do
ANS requires a LIFO queue
Lgl yi-o.l T + GetEncodeStatel SrsnooddSSaterrrnssttinnn/./ ransition(L. y:) 8, .biis +- (s))
Outpul Bits( B, Lower BiterBBitss,, nbils))
8 + >> nbits
end for Flush(B, return B end procedure
procedure DECODE(BitStream B, Decoderl.ookupTable Lp, Offsets {or y 0n
Initialize Symbols
i =+0 0
+ Bits B, MorBits(s))
while not TsFinished B) do
L t Lpiyi-on Vi-o.,
+ Ded ranssttion
Ví, nbits <- T(s)
8 - (s << nbits) + BitutBits nbits)
i с i
end while return y end procedure
17

---

Published as a conference paper at ICLR 2024
to guarantee that (H((()))))/OD) O, the following inequality must hold:
log S; log Σ S 0.
(19)
Equivalently, we have
D. = = St < 1 log Σ S = log Σe DDD
(20)
Therefore, we only need to find case such that:
Vi,
D; log S: 2 1 1-lloo Σ S.
(21)
As D; 2 0, we only need:
log E5 5.
(22)
which can be easily satisfied when the codebook size is large or the codewords are close to each other. This completes the proof.
D.2 PROOF OF PROPOSITION 5.2
Proposition 5.2 Suppose that p(=) follows a uniform distribution. Given the same sample input y
and parameters 0 of the generative model log po (xx we have log Parg ming Lava (x) By")
log Parg ming Cvo (x) By").
Proof. Suppose that
θ, B arg min Lvę arg min LRVQ
(23)
0,B
It follows from 4 that
i = arg max log PogPo(aa (x: Byy) and, T(x) = Byy
(24)
For RVQ.
δLnvo (H(g(z)x)) 0D
CRVQ = (q(=(z)) - (z) (x"",,,
(25)
0B 3D
0B
According to the proof of Proposition 5.1
can obtain (H( 0 only if
Vi, D. log
log Es. S..
(26)
In other words, D1 = D2 = =... = - log E. S,, which means that g(z ()) is a uniform distribution.
Given (24), the above equality is satisfied only when B = B2 ... (wx) that is, the codebook collapses.
Apparently, this is not the global optimal point of the codebook for LRVQ. which contradicts with (23 Therefore,
log Parg ming Cave (x Byy")
log Pa (x) log Parg lvg (al By"")
(27)
This completes the proof.
D.3 RELATIONSHIP TO SQ-VAE (TAKIDA ET AL.
2022
2022) introduces a stochastic quantization and dequantization process into the latent model of the relaxed VQ. The quantized latent variable is denoted as 2, and the dequantized latent variable as z. The inference model can be represented as (=| (g) 2) and the generative model p(z,)()) )p (x) (1))2(/z/)))))))))))))1(() The optimization objective of SQ- VAE is the negative ELBO:
SQ-VA VAE (Taiida et al.
Lso = Eg (x log D2 =z)) Epp log (2 (2) (1,(2, log log,(z,,) log pl 2 22,) xx)
(28)
18

---

Published as a conference paper at ICLR 2024
400
5.5
Perplest
+00
3.3
100
2.7
25
120020
100000
120020
Traning Step
Training Step
-VQ -VVQ -SQVAE - STHQ
VV -RVQ - S2VAA STHQ
(a)
Figure 5: Convergence processes for different quantization methods.
Let us further assume that z is a Gaussian distribution with the same variance dg in both pa and 41- Eq. (28 can be simplified as:
Lsq = || zell2
(=y)(())) log (z(gz))) - Expp(, logp(z)) logg((z)z))
(29)
2
Recall that the loss function of the proposed STHQ is
LSTHQ= = mimmll 1|2 E.g log p(z) - log
(30)
We can tell that SQ-VAE is equivalent to STHQ when 20 1, the entropy term is removed, and the Gumbel-softmax sampling for the categorical distribution is replaced by the proposed st-hardmax sampling for g2(2,lz)) As discussed in Section 2 SQ-VAE requires log (og|(2=),,) to degenerate so ensure de - 0. However, for compression tasks, log p(z]=z) needs be non-degenerate. Consequently, ||=~#,112 approaches infinity 0g + 0. In other words, the stochastic quantization in
20 SQ- D-EV not applicable to our problem, and thus, we propose STHQ as remedy.
D.4 EMPIRICAL ANALYSIS OF DIFFERENT QUANTIZATION METHODS
In order to further substantiate the preceding discussions, we conducted optimization experiments employing different quantization methods, namely VQ, RVQ. SQ-VAE, and the proposed STHQ. The experimental setup used in Section 6.2 was adopted for training. Following the methodology outlined in (Takida al. 2022) for VAEE we employed codebook perplexity as a measure of codebook utilizztiinn, and the likelihood (log p(z) z2y) as an indicator of the quality of reconstruction throughout the entire training process. The results are presented in Figure
It is obvious that STHQ achieves the best reconstruction result and the codebook utilization is very close to RVQ and VAE. Analysis of the codebook perplexity results reveals that RVQ and SQ- VAE achieve more representative codebook compared to VQ, owing to the global update property facilitated by the entropy term. However, their likelihood remains inferior to that of VQ, as confirmed by Proposition 5.2 In contrast, the proposed STHQ approach achieves perplexity value comparable to that RVQ and SQ-VAE, while exhibiting faster convergence of likelihood than VQ. This notable improvement can be attributed to the global update capability and the mitigation of the entropy gap problem enabled by STHQ, resulting in accelerated convergence of both the codebook and the likelihood.
E EXTRA EXPERIMENTAL RESULTS
E. FULL EXPERIMENT RESULTS FOR SECTION 6.2
Here we present the compression time the backbone network and entropy coders for each method in Table 1. The results, displayed in Table 5 indicate that for most methods, the compression time
19

---

Published as a conference paper at ICLR 2024
Table 5: Experimental results of latent coding Methods. Expands upon Table 1 decompression times for each method CF10.
including compression and
CF10 Comp Time (ms)
CF10 Decomp Time (ms)
Methods
Net
Coder
Total
Net
Coder
Total
Continuous
bits-back
Townsend al. 2019a)
32.07
4.54
36.61
32.44
9.85
42.29
UQ (Ballé et al. 2018
0.98
1.53
2.51
1.60
1.54
3.15
Diserete
VQ van den Uord al. 2017
0.98
1.59
1.57
2.45
McQuie Zhu et al. 2022
1.03
11.25
12.28
1.57
8.75
10.32
MaskConv2D D (vnn den Oord et al
60
1.03
4.26
5.28
1.75
632.24
633.99
Autoregressive
Checkerbourd (Hc ct al. 2021
1.03
4.65
5.67
1,63
4.71
6,34
IANS [D13
2013
0.99
0.92
1.91
1.66
0.86
2.52
FSARIOD)
0.98
1.00
1.98
1.58
0.97
2.55
Proposed
FSARIO2)
0.99
1.08
2.07
1.62
1.04
2.65
FSAR(O2)ALSN FSAR(O3)
0.98
1.47
2.45
1.59
1.06
2.65
0OM
OOM
OOM
OOM
OOM
OOM
Table 6: Experimental results lossless image compressors on different datasets. BPD values for IDF and
IDF+ on CF10, IN32, and IN64 datasets are obtained from their respective original papers.
Practical BPD
Speed (MB/s)
Methods
CF10
IN32
IN64
Kodak
CLIC20
CSpd
DSpd
PNG
5.78
6.09
5,42
4.53
4.23
16.559
59.207
Traditional
WebP
4.52
4.95
4.32
3.20
3.08
1.762
26.688
FLIF
4.29
4.78
4.25
2.88
2.82
0.858
2.715
L3C
4.77
4.89
4.29
3.55
3.34
0.036
0.057
Autoencoder
bitswap
3,84
4.54
0.003
0.004
PILC
4.28
4.95
4.51
3.68
3.55
2.659
3.725
IDF
3.34
4.18
3,90
0.010
0.007
Flow
IDF++
3.26
4.12
3.81
0.009
0.007
PILC+STHQ
4.29
4.90
4.50
3.68
3.57
2.566
3.795
Proposed
PILC+STHQ+ESAR
4.18
4.80
4.39
3.66
3.51
2.378
3,590
closely aligns with the decompression time. However, the MaskConv2D autoregressive method demonstrates notable distinction, benefitting from its parallel computation during the compression process.
E.2 FULL EXPERIMENT RESULTS FOR SECTION 6.3
We present comprehensive results in Figure encompassing experiments conducted on Ima- geNet32/64, well as high-resolution benchmark utilizing the Kodak (Kod 2023) dataset and CLIC 2020 (CLT) 2023) validation set (referred as to Kodak and CLIC20 respectively). The Kodak dataset consists of images with a resolution of 512 x 768, while the CLIC dataset ranges from 720 x 439 to 2048 x 1370 in resolution. It is important to note that our model is trained following the methods described in (Mentzer al.) (018)) and (Kang al. 2022), using a subset of the preprocessed Open Images training dataset (Kuznetsova et al. 2018). Furthermore, our evaluation includes lossless codecs based flow models. specifically integer discrete flows (IDF) (Hoogeboom et al.) 2019) and its enhanced version IDF++ (van den Berg et al. 2020). The corresponding results are presented
Table
The results of the PILC+STHO+FSAR method on IN32/64 demonstrate similarity to CIFAR10, exhibiting similar BPD values WebP and FLIF, and slightly inferior performance compared to PILC. In the case of high-resolution images, the PILC+STHQ+FSAR method performs marginally worse than WebP but outperforms PNG, likely due to the limited versatility the shallow backbone. This implies the necessity for improved backbones when dealing with large-scale images. However, important to note that the proposed STHQ and FSAR techniques enhance the performance PILC across all datasets, validating the effectiveness and robustness of the proposed method for lossless
20

---

Published as a conference paper at ICLR 2024
Table 7: BPD results (with standard deviation) for using different autoregressive Markov models different spaces. The standard deviations are presented in the brackets.
Observation
Latent
None
8.00
5.23 (0.02)
MaskConv
5.37(0.00)
5.05(0.04)
FSAR(O1))
6.19(0.00)
5.15(0.03)
FSAR(O2)
5.49(0.01)
5.06(0.02)
FSAR(O3)
5.26(0.01)
4.96(0.03)
image compression tasks. Regarding flow models, IDF and IDF+ surpass all other methods in terms of BPD, although their slower processing speed remains critical concern.
E.3 COMPARISON OF DIFFERENT USAGES OF AUTOREGRESSIVE MARKOV MODELS
As discussed in Section 4 the use of autoregressive Markov models as the latent prior considered an optimal solution. This is because the latent space is typically smaller than the observation space, and the dependency between variables relatively weak in the latent space. As result, a more compact autoregressive Markov model, such the FSAR model, can be employed. In this section, we empirically demonstrate the validity of this assertion by comparing the performance of various autoregressive Markov models. These models include the proposed FSAR model, higher-order autoregressive Markov model (order-3 FSAR), and a Masked Convolution based model (van den Oord et al. 20) The comparison is conducted for both observation space and latent space scenarios using the CIFARIO dataset. The theoretical BPD is reported for this experiment.
The results demonstrate that, in the observation space, more complex autoregressive Markov models generally yield better compression performance than their simpler counterparts. This suggests that the dependencies between variables are more intricate in the observation space than in the latent space. Conversely, applying higher-order Markov models does not result in significant improvements in the latent space. This highlights the advantages of employing compact, low-order autoregressive Markov models as the latent prior.
E.4 COMPARISON OF DIFFERENT HYPERPARAMETEN FOR OPTIMIZING STHQ
This study aims to establish the superiority of the proposed STHQ method over existing VQ-based methods. To compare these models, various hyperparameters are considered, including backbone network settings, latent space, and codebook size. The experiments are conducted the CIFARIO dataset. Specifically, two backbone networks are evaluated: shallow backbone (consisting one downsampling upsampling layer, and one residual layer with 64 latent channels) used in previous experiments, and a deeper backbone (consisting of two downsampling or upsampling layers, and two residual layers with 256 latent channels). For the shallow backbone, the latent variables are quantized to I, 2, and channels, with corresponding codeword channels of 64, 32, and 16, respectively, For the deeper backbone, the latent variables quantized 4, 8, and 16 channels, with corresponding codeword channels of 64, 32, and 16, respectively. The proposed STHQ method is compared to VQ- VAE (van den Oord et al. 2017) (denoted as VQ), Relaxed VQ-VAE (Sønderby 2017) (denoted
as RVQ) and SQ-VAE (Takida et al. 2022) (denoted as SQ). The results are summarized Table g
Our findings reveal that the proposed STHQ method outperforms all other VQ-based methods in most configurations, except for a few cases where the compression performance of STHQ is comparable to other methods. Generally, VQ-VAE performs relatively worse when quantized to fewer channels (1 or 2). but better when quantized to more channels (4, 8, or 16). Conversely, the relaxed VQ-VAE and SQ- VAE exhibit relatively higher performance when quantized to fewer channels, but lower performance for more channels. In contrast, the proposed STHQ method consistently achieves comparable or superior results across all configurations, demonstrating its robustness in optimizing discrete latent space.
21

---

Published as a conference paper at ICLR 2024
Table 8: BPD results (with standard deviation) for using different hyperparameters for different VQ based
methods. The standard deviations are presented the brackets.
Hyperparameters
Method
Backbone
Latent
Quant
Codeword
VQ
RVQ
SQ
Proposed
64
5.75 (0.04)
5.60 (0.04)
5.48 (0.01)
5.44 (0.02)
I-layer
64
2
32
5.34 (0.07)
5.81 (0.18)
5.62 (0.22)
5.26 (0.06)
4
16
5.87 (0.03)
6.67 (0.02)
6.91 (0.05)
5.66 (0.08)
4
64
5.00 (0.05)
5.09 (0.10)
5.35 (0.06)
4.96 (0.03)
2-layers
256
8
32
4.87 (0.02)
5.23 (0.03)
6.02 (0.02)
4.79 (0.02)
16
16
5.20 (0.07)
6.30 (0.04)
7.40 (0.15)
5.17 (0.03)
Table 9: Generation task results for different VQ-based methods.
NLL
FID
Perplexity
VQ
2.83
58.00
262.93
RVQ
2.88
64.15
327.68
SQ
2.96
83.09
489.07
EMA-VQ
2.80
48.59
380.45
EM-VQ
2.76
55.76
282.97
Anneal-RVQ
3.12
81.02
61.55
Proposed
2.67
42.63
487.75
E.5 COMPARISON OF VQ-BASED METHODS ON GENERATION TASKS
Finally, we conduct comparative analysis of our proposed STHQ method against existing vector quantization (VQ)-based approaches in simple image generation tasks. In addition to the methods listed in Table codebook update for VQ-VAE (EMA-VQ) (Polyak & Juditsky 1992). expectation maximization optimizer for VAE (EM-VQ) Roy al. 2018). and relaxed VQ with exponential deterministic annealing (Anneal-RVQ) (oose 1998), where the softmax temperature is annealed from 1.0 and halved every 200 epochs. In our evaluation, we employ the negative lieellol (NL..) metric, commonly used for assessment, along with the Frechet Inception Distance (FID) (Heusel et al. 2017)) to gauge the quality of the generated images. Additionally, we utilize perplexity as a measure of codebook utilization, where higher perplexity indicates better codebook training. These experiments
include comparisons with VQ variants, namely exponential moving average-based
are conducted on CIFAR 10, utilizing the same hyperparameters as described in Section 6.2 The results presented in Table 9
The results reveal that the proposed STHQ method outperforms all other approaches across all metrics. Among the VQ variants, EMA-VQ and EM-VQ demonstrate superior performance compared to traditional VQ, while Anneal-RVQ exhibits notably worse results, potentially due to an improper annealing schedule. This suggests that tuning deterministic annealing proves challenging practical applications. On the other hand, the proposed STHQ method achieves better performance than both EMA and EM optimization techniques, while also avoiding the complexities associated with tuning the annealing process.
F EXPERIMENT IMPLEMENTATION DETAILS
F.I HARDWARE SETUP
To ensure fair comparisons, we conduct all experiments under the same hardware configuration. For CPU-based experiments, utilize desktop machine equipped with an Intel i7-6800K CPU. For GPU-based model training, we employ virtual machine on a Kubermaker cluster featuring NVIDIA V100 GPUs.
22

---

Published as a conference paper at ICLR 2024
F.2 NETWORK STRUCTURE
In our experiments, unless otherwise specified, we maintain consistency in the backbone network across all methods to ensure that the compression ratio remains unaffected by variations in the backbone architecture. Our focus on general computation devices, such as CPUs, leads us to adopt shallow backbone network based on ResNet (He et al., 2015). This selection aims to minimize the number of floating-poin operations (FLOPs) involved. The backbone network comprises a single downsampling (or upsampling) convolutional layer and one residual block within the inference (or generative) model. Consequently, the inference model performs 2x downsampling, while the generative model performs 2x upsampling. With regard to the latent space, we set the number of input latent channels to 64 and quantize them into 2 channels, each containing 256 codewords. This configuration leads to a codebook size of 2 x 256 x 32, where 2 represents the quantized latent channels, 256 denotes the number of codewords (also known state number), and 32 signifies the codeword dimension. Consequently, a 3 x 32 x 32 input image corresponds to 2 x 16 X 16 latent variables.
F.3 HYPERPARAMETERS AND CONFIGURATION OF THE PROPOSED METHODS
Regarding FSAR, each variable relies on its previous neighboring variables in the spatial domain. Specifically, the latent prior utilizing the Order-1 Markov model is defined as p(vs.,, p(vs,,xlv.,..e- (VC.,,.-1)) 1), while 1,k, Wi.j.k-).. where j. represent the indices of the channel, height, and width dimensions, respectively. During training. the Markov models are implemented using a 3-layer network consisting of 3 Linear layers and 2 ReLU layers. For the learnable state number, the initial number of states is set to 256, and a 1.5 is used in the a-entmax method. For STHQ, the temperature for Gumbel- GS, is set as constant 0.5.
Order-2 is p(yi.j.k 1,,, -i.j..-)), and Order-3 is (.
F.4 TRAINING PROCEDURE
The training process utilized the Adam optimizer with a learning rate of 10-3, For the CIFAR10 dataset, the training was performed for 1000 epochs, with a batch size of 64 per GPU. Regarding the ImageNet32 dataset, the training was conducted for 50 epochs, with batch size of 64 per GPU. As for ImageNet64, the training setup was identical to ImageNet32, except that the batch size was set to 16 per GPU.
is important to note that in the STHQ-basee based experiments, the Gumbel-softmax reparameterization in the st-hardmax operator maintained constant temperature throughout the training process. In this case, a temperature of 1.0 was employed as we required its gradient for optimization purposes. Additionally, the log p(a:) By term in the loss function (9) was implemented using straight- through estimator, similar to VQ-VAE (van den Oord et al. (017). This approach, rather than directly utilizing the samples in the proposed st-hardmaxx proved be beneficial for effectively optimizing the inference network.
F.5 IMPLEMENTATION DETAILS FOR THE LATENT SPACE MODELS IN TABLE
Here we provide implementation details for latent space models used in Table First, Table 10 illustrates the overall architecture of each method by linking them to Figure I Generally, for discrete latent space (Figure 1b), the variations among different methods primarily stem from the choice of quantization method. On the other hand, for models utilizing discrete autoregressive latent space (Figure icp. the main point of differentiation lies in the context model. Regarding our proposed FSAR- tANS model (Figure ld) we utilize the STHQ method for quantization and conduct experiments using different FSAR models.
Next, we provide detailed explanations for the compared methods. Unless otherwise specified, all models employed in this study adhere to the same training scheme and utilize identical inference and generative networks. The reported values are obtained by averaging the corresponding metrics from three independent experiments. It is worth mentioning that the memory occupation values only take into account statically allocated data, such as the lookup table for tANS and the kernel parameters for MaskConv. In the case of LSN, only valid states are utilized to construct lookup tables. Therefore, the number of reduced states is consistent with the memory savings, as the lookup
23

---

Published as a conference paper at ICLR 2024
Table 10: Architecture and Implementation Details of methods in Table 1
Architecture
Method
Detail
Figure la
bits-back
Townsend et al.
2019a
UQ (Ballé ct al. 2018
quant: uniform
Figure 1b
VQ (van den Oord al.
2017
quant: VQ
McQui et al. 2022
quant: 3-layer VQVAE
MaskConv2D (van den Oord et al.
2016b)
context: MaskConv2D
Figure
Checkerboard (He al. 2021
context: Checkerboard
LANS (Duda 2013)
FSAR-ANS: non-auloreerressive
FSAR(OT)
FSAR-IANS: Order-1 Markov
Figure ld
FSAR(O2)
FSAR-tANS: Order-2 Markov
FSAR(02) FSAR(03)
FSAR-ANS: Order-2 Markov + LSN
FSAR-IANS: Order-3 Markov
table takes O(C')) space. For example, the FSAR(O2)+LSN (C 256, N 2) approach achieves an average of 187.33 states, which can be also inferred from the memory occupation saving from Table 256(276.586/ 58055166511111444441 x 187.33.
STHQ+AANS Duda 2013) For the STHQ+ANS method, we utilize the IANS implementation from the zstd library" in accordance with the conventions followed other STHQ-based experiments such FSAR and various autoregressive-basee latent models.
bits-back (Townsend et al. 2019a) The bits-back method is implemented using the craystack library? provided by the authors.
UQ (Ballé et al. 2018) For universal quantization latent coding, we employ the entropy bottleneck proposed in (Ballé et al. 2018) and implemented in Compress The size of the latent channels is set to 64, which is equivalent to the output channels of the inference network.
VQ (van den Oord et al., 2017) The VQ model is implemented using the original implementation of VQ- [AH (van den Oord et al. 2017). with tANS used for constructing the latent coder. The codebook size the same as that used in the STHQ-based experiments, which is 2 256 32.
McQuic Zhu et al., 2022) We adhere to the official McQuic repository for implementation, replacing the entropy coder with tANS for fair speed comparison. We adopt the recommended hyperparameters m = 2 and K = [8192, 2048, 512]. and the size of the latent channels remains consistent with other experiments, which is 64.
MaskConv2D (van den Oord et al. 2016b) We employ a single masked convolution layer, using mask as described in (van den Oord et al., 2016b), as an autoregressive model for latent vari- ables. The entropy coder used is the same as in STHQ+ANSS. During latent entropy coding, the autoregressive model iterated along the spatial dimensions.
Checkerboard (He et al. 2021 Similar to MaskConv2D, we utilize a single masked convolution layer as the autoregressive model, adopting checkerboard-fome mask as described (He et al. 2021).
The latent entropy coding is iterated in two subsequent steps.
F.6 IMPLEMENTATION DETAILS FOR THE COMPRESSION METHODS IN TABLE 6
PNG, WebP To compress images in PNG and lossless WebP formats, we utilize the pillow library
htteer//gthuub.ccm//cccoo/////
httass//ppi..nr//poeeee///111
24

---

Published as a conference paper at ICLR 2024
FLIF For FLIF compression, we the imageio-fil python library which binds the official C code and we exccute iaangcco...imwrriee(bttttss image, format=" "FIIF" disable_come _cclc_coloor__eeeeeeeeeeeeeeec buckets=True) to compress images.
L3C Mentzer et al., 2018 To test using the L3C method, we employ the official code and the provided pretrained models for different datasets.
bitswap Kingma et al. 2019) For bitswap compression, we use the official code and the provided pretrained model (labeled "nz2") ") for testing CIFARIO and ImageNet32 (ImageNet64 is skipped as no pretrained models is provided). It should be noted that, unlike other experiments, all images are stored in the same bitstream, which is one reason for the method's low BPD.
PIL. (Kng et al., 2022) Since the code for PILC is currently unavailable, we implement this method in our benchmark. This includes the three-way autoregressive (TWAR) model for image prediction and the AI-ANS coder for entropy coding, as described in (Kang et al., 2022). Specifically, we adopt the recommended backbone network from the authors, consisting of one downsampling (or upsampling) convolutional layer and four residual blocks in the inference (or generative) model. It is worth mentioning that the code for PILC not publicly available, and therefore we have independently implemented this method in our benchmark. Our implementation includes the proposed three-way autoregressive (TWAR) model for image prediction and the AT-ANS coder (on CPU) for entropy coding. We use the same backbone network architecture our experiments, which consists of one downsampling (or upsampling) convolutional layer and one residual block in the inference (or generative) model. The proposed method combined with PILC also follows this backbone architecture.
G DISCUSSIONS AND FUTURE WORKS
Despite the promising outcomes achieved through the proposed finite-state autoregressive entropy coding approach, there are still certain limitations that require attention. Notably, our experimental utilization of the lookup table still demands approximately 300MB of storage, even after reducing the number of states using the proposed Learnable State Number approach. This presents a significant hurdle for implementing the approach on portable devices and emphasizes the necessity for future research on maximizing the advantages of the autoregressive Markov model while keeping the space complexity limited.
One potential solution is to restrict the number of probability distributions employed for tANS coding. thereby limiting the corresponding state transition lookup tables used by tANS. By doing so, we only need to store pointers to these restricted distributions the FSAR lookup table L. Assuming we limit the distributions to D and each corresponding state transition lookup table contains S states, for order-2 FSAR models, the minimal memory requirement is DS $1 + C2 [log250 D1 bytes. In contrast, unlimited distributions would necessitate S[oggge S1 bytes. In our experiments, setting D = 256 reduced memory usage from approximately 512MB to around 2MB. We leave determining the optimal selection of these limited distributions to achieve minimal performance degradation to future work.
Furthermore, it would be advantageous for future studies to consider adaptive selection of the Markov model structure. For instance, an intriguing approach has been explored in (Lin et al. 2023), where the authors attempt different decoding orders for parallel autoregressive models and select the optimal one during training. Similar methodologies could potentially be applied to the structure learning of FSAR Markov models.
Finally, it is also interesting to investigate the acceleration of the backbone network by means of methods such as neural architecture search. Indeed, Yang et al. (202T) have already delved into this direction in compression codecs using slimmable neural networks (Yu et al. 2018).
ittoss//qtthbbbcoo/////l11/////1
25


Reviews:
[{"review_id": "hT8De7T20y", "reviewer": "Reviewer_8pa1", "summary": "This manuscript proposed a finite-state autoregressive entropy coding method for efficient learned lossless compression, which utilizes a lookup table to expedite autoregressive entropy coding. Specifically, a straight-through hardmax quantization scheme is proposed to enhance the optimization of discrete latent space. Experimental results show that the proposed lossless compression method could improve the compression ratio by up to 6% compared to the baseline, with negligible extra computational time.", "questions": "Please refer to the Weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper addresses the shortcomings of existing autoencoder-based codecs and makes targeted improvements from several aspects. The concept of autoregressive modeling based on finite-state Markov model via look-up table is reasonable to alleviate the computational burden of autoregressive models. Besides, an end-to-end adaptive optimization method for selecting learnable state number is proposed to reduce the look-up table size. Furthermore, the straight through hardmax quantization method is proposed for optimizing vector quantized discrete latent space models.\n\nThis manuscript is well structured. The technique descriptions is detailed and easy to follow. \n\nExperimental results show that the proposed lossless compression method could improve the compression ratio by up to 6% compared to the baseline, with negligible extra computational time.", "weaknesses": "1) Some technique details in Fig. 2 are not clear enough. For example, how to obtain quantize latent y from discrete latent variables z and adaptive state number module?  How to obtain the codebook? There are two bitstream in Fig. 2. Is there any operations proposed for optimizing the data steam?\n\n2) Experimental results are conducted on CIFAR 10 and ImageNet32/64. How about the results on Kodak dataset  and Tecnick dataset? Besides, only the decoding time are provided in Table 1, how about the encoding time? \n\n3) The proposed method achieve RD performance gain and computational efficiency at the cost of memory consumption. Will the size of memory change with different datasets?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "05qRuEzbNj", "reviewer": "Reviewer_QqM2", "summary": "This paper studies lossless image compression and proposes a finite-state autoregressive model to model the data, relying on a lookup table to predict likelihoods. They also present a soft relaxation of VQ to improve optimization.", "questions": "What is the state number you use? Ie if you use LSN, what is predicted?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper focuses on improving compression rates without increasing computational needs. In Table 2, we see that the authors succeed at this, improving over PLIC without significantly slowing the method down.\n\nThe soft VQ variatn (STHQ) is sensible.\n\nThe idea to parameterize a lookup table with an MLP and rely on the discrete nature of the inputs is smart.", "weaknesses": "Presentation of Results\n- I had a very hard time following the results. Some suggestions: Please clarify what architecture was used for Table 1. I did not find it in the text and initially was confused why the IN32 results are so much higher than what is typical in the literature\n- Can you consider replacing Table 2 with a three figures like in PLIC? Where we have one figure for each dataset, and see DSpd vs. BPD? This would make it much much easier to understand the contribution of this paper. Also, some of the methods that have been published for a long time are missing in the table (IDF, IDF++), and I don't know where the 0.057 MB/s for L3C is coming from, interpolating from Table 54 in their paper, I arrive at 0.66MB/s (I was confused because I remember they were only 3x slower than FLIF)\n\nRelated work\nIn the context of STVQ, I am missing a citation to Agustsson et al, 2017, \"Soft-to-Hard Vector Quantization for End-to-End Learning Compressible Representations\", https://arxiv.org/abs/1704.00648. IIUC, they applied a similar Gumbal Softmax trick, but also used it in the forward pass (whereas STVQ uses argmax for the forward). Still, this seems related.\n\nMinor\n- Please remove the negative vspace around tables, it's very hard to know what is table and what is text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qg5gjjdZPR", "reviewer": "Reviewer_DfKd", "summary": "## Summary\n* The authors propose a accelerated auto-regressive entropy coding system for discrete autoregressive coding. It utilizes a look-up-table (LUT) that somewhat reminds me of n-grams markov model. Further, it proposes a hard-max version of gumbel-softmax VQ, with the sampling step replaced by an arg-max operation. The empirical results support their claims well.", "questions": "## Questions\n* The FSAR seems to be a general approach that is not necessarily used with rANS. Can it be used with any FIFO coding approach?\n* The hardmax gumbel VQ seems to converge to softmax gumbel VQ as temperature cools down. The current form of hardmax gumbel VQ can be viewed as a softmax VQ with forward temperature 0. Is that possible to tune the annealing of softmax gumbel VQ to achieve the same effect of hardmax gumbel VQ?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "## Strength\n* The finite state MC model seems to be a very interesting middle point between vanilla neural network based autoregressive model and vanilla n-gram model. It achieves a good balance between compression ratio and speed.\n* The STHQ training of VQ-VAE looks interesting and experimental results show that the improvement of nll (bpd) is quite consistent over many hyperparameter setting. The proposed approach and result is interesting to general readers beyond compression community, as VQ-VAE are widely used in generative modeling (Latent-Diffusion, Dalle).", "weaknesses": "## Weakness\n* I think the FSAR that authors propose is an entropy model, not entropy coder. Usually I only regard general coders such as AC, ANS, RANS / RANGE as entropy coder. And I tend to agree that autoregressive model, FSAR are only entropy model. Though they can be deeply coupled with entropy coder, I still think they are only entropy model. \n* As the authors have setup VQ-VAEs with various architectures, I suggest the authors test more variant of VQ-VAE training, such as expoential moving average codebook, and expectation maximization VAE, beyond the current original VQ, gumbel-softmax VQ and SQ. Furthermore, subtle variants of gumbel-softmax VQ, e.g., soft-gumbel / ST-gumbel, penalize KL-term or not can also be studied. Furthermore, it is also interesting to evaluate the sample quality of those VQs, in terms of FID / IS. As the models with better nll (bpd) might not have better sample quality, and vice versa [A note on the evaluation of generative models]. This can provide a more comprehensive understanding of the proposed hardmax-gumbel training strategy. And make this paper more interesting to density estimation community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qtl3xLD2g6", "reviewer": "Reviewer_7HPB", "summary": "This paper investigates efficient learnt lossless data compression, aiming to improve the compression performance while maintaining computational efficiency.  \nThe model is based on autoencoder-based learnt codecs, with discrete autoregressive latent space.  \nThe paper proposes Finite-State AutoRegressive (FSAR) entropy coder, which combines a low-complexity autoregressive Markov model with a fast entropy coder to achieve efficient latent coding.  \nThe paper proposes Straight-Through Hardmax Quantization (STHQ) to opotimize the discrete latent space.  \nThe proposed method improves the compression ratio by up to 6% compared to commonly used discrete deterministic autoencoders, with negligible additional computational time.", "questions": "It is not clear how the operation number in Table 4 is obtained, will this be included in the opensource code?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Most related works are properly cited and discussed, thus the proposed methods are well motivated.  \n\nThe manuscript is well written and friendly to readers, with in-depth analysis of previous learnt compression architectures.  \n\nThe experiments are thorough and extensive. I like the informative appendix.   \n\nThough the proposed FSAR and STHQ are improvement based on previous techniques as already explained in the manuscript, the technical novelty is enough from my point of view.", "weaknesses": "The memory consumption shown in Table 1 is significantly larger than previous methods. \n\nIt is great to compare the net and coder latency separately in Table 1, however, the coder of each methods are not explained. It is important to mark the coder of each methods since the latency of network is very close.\n\nOnly CIFAR10, ImageNet 32/64 are used in the experiments. It is better to show the scalability of the proposed methods on larger image sizes.\n\nMissing citations. For efficient lossless image compression, I noticed there exist two very efficient design [R1, R2]. Although comparison is not required, these two paper should be discussed to better reflect recent progress regarding efficient learnt lossless image compression.\n[R1] Guo, Lina, et al. \"Practical Learned Lossless JPEG Recompression with Multi-Level Cross-Channel Entropy Model in the DCT Domain.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n[R2] Guo, Lina, et al. \"Efficient Learned Lossless JPEG Recompression.\" arXiv preprint arXiv:2308.13287 (2023).\n\nMinor: \nIn F.5 and F.6, the implementation details of checkerboard are not provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)