SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yanghao Li, Tongda Xu, Yan Wang* Jingjing Liu, Ya-Qin Zhang* Institute for AI Industry Research (AIR), Tsinghua University liyangha18@mails.tsinghua.edu.cn, wangyan@air.tsinghua.edu.cn
We consider the problem of idempotent learned image compression (LIC). The idempotence of codec refers to the stability of codec to re-compression. To achieve idempotence, previous codecs adopt invertible transforms such as DCT [Wallace, 1991] and normalizing flow [Papamakarios et al., 2021]. In this paper, we first identify that invertibility of transform is sufficient but not necessary for idempotence. Instead, it can be relaxed into right-invertibility. And such relaxation allows wider family of transforms. Based on this identification, we then implement an idempotent codec using our proposed blocked convolution and null-space enhancement. Empirical results show that we achieve state-of-the-art rate-distortion performance among idempotent codecs. Furthermore, our idempotent codec can be extended into near-idempotent codec by relaxing the right-invertibility. And this near-idempotent codec has significantly less quality decay after 50 rounds of re-compression compared with other near-idempotent codecs.
Learned Image Compression (LIC) has been widely studied in recent years [Balle et al., 2017, 2018, Minnen et al., 2018, Cheng et al., 2020, Minnen et al., 2020, He et al., 2021, 2022] and has shown promising rate-distortion (RD) performance. However, the loss caused by re-compression is much more severe in LIC compared with traditional codec, which seriously limits the practical application of LIC [Kim et al., 2020]. In this paper, we study the idempotence of LIC, which refers to the stability of codec to re-compression. More specifically, denote the original image as  $\pmb{x}$ , the encode-then-decode procedure as  $f(\cdot)$ , and the reconstructed image as  $f(x)$ , we say a codec is idempotent if:
$$
f (\boldsymbol {x}) = f (f (\boldsymbol {x})). \tag {1}
$$
For traditional codecs such as JPEG [Wallace, 1991] and JPEG2000 [Taubman et al., 2002], idempotence is easily achieved. This is because those codecs adopt invertible transforms such as Discrete Cosine transform (DCT) and Discrete Wavelet transform (DWT). And the only non-invertible operation is the scalar quantization. As scalar quantization using rounding is naturally idempotent, the idempotence of the whole codec can be easily assured. For LIC, however, neural-network-based transform is introduced for expressiveness. And as most neural networks are non-invertible, the idempotence of LIC is not trivial.
A natural solution to this problem is replacing the non-invertible encoding transform with invertible ones. [Helminger et al., 2021] construct the encoding transform with only invertible normalizing flow [Papamakarios et al., 2021]. However, due to the limited expressiveness of invertible operations, a dramatic RD performance drop is observed. Another line of works are targeted at near-idempotence, which means that they achieve a small  $|f(x) - f(f(x))|$ , but not  $f(x) = f(f(x))$ . These works
adopt partially invertible encoding transform [Cai et al., 2022] or use additional regularization loss [Kim et al., 2020] to constrain re-compression loss, but none of them is able to achieve idempotence like JPEG and JPEG2000.
In this work, we first identify that invertibility is sufficient but not necessary for idempotent codescs. Instead, it can be relaxed into right-invertibility, and such relaxation allows for more flexible and expressive encoding transforms. Based on this observation, we investigate practical implementation of right-inverse, and propose a highly efficient blocked convolution to overcome the forbidding time complexity of right-inverse. Additionally, we leverage the null space decomposition [Schwab et al., 2019, Wang et al., 2022a,b] to further boost the expressiveness of the right-invertible encoding transform. Empirically, we achieve state-of-the-art RD performance among existing idempotent codescs. Further, our idempotent codec can be easily extended into a near-idempotent codec, which also achieves state-of-the-art re-compression performance among near-idempotent codescs.
2 Sufficiency of Right-Invertibility for Idempotence
Most modern LIC [Balle et al., 2017, 2018, Cheng et al., 2020, He et al., 2022] are composed of four components: the encoding transform  $E$ , the decoding transform  $D$ , the quantization  $Q$ , and the entropy model. On the encoder side, the encoding transform  $E$  transform the input image  $x$  to some latent representation, which is then quantized by  $Q$  to give the code  $y$ . The entropy model models the entropy of code  $y$ , and is then used losslessly compress  $y$  to a bitstream. On the decoder side, the received bitstream is losslessly decompressed to code  $y$  with the help of the entropy model. After that, the decoding transform  $D$  transforms the  $y$  to the decompressed  $\hat{x}$ . This compress-decompress procedure can be represented as
$$
y = Q \circ E (x), \hat {x} = D (y). \tag {2}
$$
Note that the code  $y$  is losslessly encoded and decoded, regardless of how good the entropy model is. Therefore entropy model does not influence the distortion, and is thus omitted for simplicity.
Using similar notations, the re-compression cycle can be represented as
$$
y _ {n} = Q \circ E \left(x _ {n - 1}\right), x _ {n} = D \left(y _ {n}\right) = D \circ Q \circ E \left(x _ {n - 1}\right), \tag {3}
$$
where  $x_0$  is the original image,  $y_n$  is the code after  $n$  times' re-compression, and  $x_n$  is the corresponding decompressed image. We say a codec is idempotent if and only if
$$
x _ {n} = x _ {1} = D \circ Q \circ E \left(x _ {0}\right), \forall n \geq 1. \tag {4}
$$
An obvious sufficient condition for idempotence (or for Eq. 4) is
$$
D = E ^ {- 1}. \tag {5}
$$
In other words,  $D$  and  $E$  are inverse of each other. Then, the only difference between  $x_{1}$  and  $x_{n}$  is the number of applications of quantization operation  $Q$ :
$$
x _ {n} = D \circ Q \circ E \circ \bar {D} \circ Q \circ E \left(x _ {n - 2}\right) = \dots = D \circ Q ^ {n} \circ E \left(x _ {0}\right). \tag {6}
$$
This re-compression procedure is idempotent as long as  $Q$  is idempotent. Usually  $Q$  is a scalar quantization implemented by rounding, then this procedure is naturally idempotent. The invertibility of encoding transform is satisfied in traditional image codes like JPEG and JPEG2000, and can also be achieved using normalizing flow [Helminger et al., 2021], thus these codes are idempotent. However, invertibility of  $E$  (Eq. 5) is sufficient but not necessary. We notice that for  $E \circ D$  in Eq. 6 to be canceled out, it is sufficient to have
$$
E \circ D \left(y _ {i}\right) = y _ {i}, \tag {7}
$$
which is exactly the definition of right-invertibility.
Relaxing the requirement for  $E$  from invertibility to right-invertibility brings two advantages. On the one hand, the family of right-invertible functions is much less constrained than the family of invertible functions, which means we have wider choices for  $E$  and its easier to improve the expressiveness of  $E$ . On the other hand, the sufficient condition for  $E$  to be invertible is that  $E$  is bijective, and the sufficient condition for  $E$  to be right-invertible is that  $E$  is surjective [Mac Lane, 2013]. A surjective  $E$  can transform different input images  $x$  into the same code to save the bits for distinguishing them, while bijective  $E$  always transforms different input images into different codes. As later shown in the experiment section 4.3, this bit saving property benefits the RD performance of lossy compression.
Figure 1: Right-Invertible Convolution. (a) Naive convolution matrix (upper) and blocked convolution matrix (lower). (b) Receptive field pattern for the proposed blocked convolution. (c) Null-space enhancement (NE) and coupling enhancement (CE) to improve expressiveness.
3 Practical Design of Right-Invertible Codec
In the previous section, we have shown that if an encoding transform  $E$  and quantization  $Q$  are right-invertible, we only need the decoding transform  $D$  to be the right-inverse of  $E$  to formulate an idempotent codec. Since the success of LIC is, to a large extent, owing to the expressive power of the learned encoding and decoding transforms, our task becomes how to design an expressive yet right-invertible encoding transform. As composition of surjections is still surjective [Mac Lane, 2013], this task can be further decomposed into designing small components of right-invertible transforms and combining them altogether. Specifically, we construct the encoding transform  $E$  in a composition form as  $E = E_{n} \circ E_{n-1} \circ \ldots \circ E_{1}$ , and make each  $E_{i}$  a right-invertible sub-transform.
This section is organised as follows: In Sec. 3.1-Sec.. 3.3, we discuss how to design expressive yet right-invertible atom transforms used in LIC, such as convolution, normalization and quantization. In Sec. 3.4, we discuss how to organize those atom transforms into an idempotent codec. And In Sec. 3.5, we discuss how to relax this idempotent codec into near-idempotent codec.
3.1 Efficient & Expressive Right-Invertible Convolution
Convolution is of great importance to LIC and makes up the majority of computation cost. Here, we discuss how to implement right-invertible convolution with efficiency and expressiveness. The overall design of right-invertible convolution is illustrated in Fig.1(c).
3.1.1 Blocked Convolution for Efficiency
The right-inverse of a convolution can be calculated in serial (if it exists), but the time complexity is forbiddingly high. To see why, consider a 1-d convolution with kernel size 5, padding 2, stride 2, and channel 1. The input and output are  $\pmb{x} = (x_{1}, x_{2}, \dots, x_{12})$  and  $\pmb{y} = (y_{1}, y_{2}, \dots, y_{6})$  respectively.
The serial solution of right-inverse goes as follows: first solve  $(x_{1}, x_{2}, x_{3})$  given  $y_{1}$ , then solve  $(x_{4}, x_{5})$  given  $y_{2}$  and already solved  $(x_{1}, x_{2}, x_{3})$ , and so on till the whole  $\pmb{x}$  is solved. This serial solution cannot be made parallel because solving  $(x_{4}, x_{5})$  needs  $(x_{1}, x_{2}, x_{3})$  to be already solved, and is thus extremely time-consuming. The fundamental reason for the dependency in solving  $\pmb{x}$  is that, some same  $x_{i}$  is involved in the forward calculation of different  $y_{i}$ , i.e., the overlapping receptive field. Therefore, if we make the receptive field non-overlapping, then parallel solution of right-inverse becomes possible.
Inspired by the non-overlapping  $1 \times 1$  convolution for inverse [Kingma and Dhariwal, 2018], we propose blocked convolution for right-inverse. Blocked convolution is also non-overlapping, but extends invertibility to right-invertibility. As shown in Sec. 4.3, this extension boots the R-D performance of idempotent LIC by a large margin.
With this non-overlapping blocked convolution, we can make solution of right-inverse parallel. Following the previous example, using the same input and output with a  $4 \times 2$  blocked-convolution kernel, for now, solving  $(x_{1},\ldots,x_{4})$  only needs to know  $(y_{1},y_{2})$ , and solving  $(x_{5},\ldots,x_{8})$  only needs to know  $(y_{3},y_{4})$ , and these procedures can be made parallel. Matrix multiplication equivalents of normal convolution and blocked convolution are depicted in the upper and lower parts of Fig.1(a) respectively, using the well-known GEMM [GEM] formation. And an analysis of time complexity of 2D convolution can be found in appendix A.1.
3.1.2 Null-Space Enhancement and Coupling Enhancement for Expressiveness
Null-Space Enhancement
To solve the right-inverse of the proposed blocked convolution in parallel, we adopt the widely-used Moore-Penrose pseudo-inverse [Moo] as
$$
X = Y K ^ {+}, \tag {8}
$$
where  $Y \in \mathbb{R}^{b \times d}$  is the output of blocked convolution,  $X \in \mathbb{R}^{b \times D}$  is the solved right-inverse, and  $K^{+} \in \mathbb{R}^{d \times D}$  is the Moore-Penrose pseudo-inverse of the kernel  $K \in \mathbb{R}^{D \times d}$ .  $b$  is the batch size,  $D$  and  $d$  are the input and output dimension of convolution in GEMM formation [GEM].
However, simple Moore-Penrose pseudo-inverse does not have enough expressiveness. This is because right-inverse of a surjection might not be unique [Mac Lane, 2013], and the Moore-Penrose pseudo-inverse might not be the optimal choice to calculate the right-inverse for blocked convolution. Therefore, we leverage the null-space decomposition [Schwab et al., 2019, Wang et al., 2022a,b] to identify a right-inverse that is most suitable. More specifically, Eq. 8 can be extended to
$$
X = Y K ^ {+} + F (I - K K ^ {+}). \tag {9}
$$
Here  $F \in \mathbb{R}^{b \times D}$  can be an arbitrarily chosen variable. Utilizing this property, we propose to learn a function  $f(Y)$  to identify an  $F$  that is the most suitable as:
$$
X = Y K ^ {+} + f (Y) (I - K K ^ {+}). \tag {10}
$$
We call this approach null-space enhancement (NE). In this way, right-inverse  $X$  can be made more than just linear transformation of  $Y$  and thus become more expressive. Ablation study of the proposed null-space enhancement can be found in Sec.4.3. Derivation of Eq.9 and parameterization of kernel  $K$  can be found in appendix A.2.
Coupling Enhancement
The proposed blocked convolution makes the calculation of right-inverse parallel-friendly, but it also restricts the receptive field to a blocked pattern (Fig. 1(b)). This restriction limits the exchange of information across different spatial locations. To overcome this drawback, we propose to introduce a coupling enhancement (CE) after the blocked convolution.
Specifically, we implement a coupling structure [Dinh et al., 2016], and utilize normal convolutions without blocked limitation as its scale and translation functions. This structure is formulated as
$$
\boldsymbol {x} = \left[ \boldsymbol {x} _ {1} \boldsymbol {x} _ {2} \right], \boldsymbol {y} _ {1} = \boldsymbol {x} _ {1}, \boldsymbol {y} _ {2} = \boldsymbol {x} _ {2} \odot \exp (s (\boldsymbol {x} _ {1})) + t (\boldsymbol {x} _ {1}), \boldsymbol {y} = \left[ \boldsymbol {y} _ {1} \boldsymbol {y} _ {2} \right] \tag {11}
$$
Here,  $x$  and  $y$  are the input and output, respectively.  $[\cdot]$  is the split/concatenate operation along the channel dimension, and  $\odot$  is the element-wise multiplication.  $s(\cdot)$  and  $t(\cdot)$  are the scale and translation functions, respectively. This structure is fully invertible as
$$
\boldsymbol {y} = \left[ \begin{array}{l l} \boldsymbol {y} _ {1} & \boldsymbol {y} _ {2} \end{array} \right], \boldsymbol {x} _ {1} = \boldsymbol {y} _ {1}, \boldsymbol {x} _ {2} = \left(\boldsymbol {y} _ {2} - t (\boldsymbol {x} _ {1})\right) / \exp (s (\boldsymbol {x} _ {1})), \boldsymbol {x} = \left[ \begin{array}{l l} \boldsymbol {x} _ {1} & \boldsymbol {x} _ {2} \end{array} \right] \tag {12}
$$
Note that the invertibility of this coupling structure does not require the invertibility of the scale and translation functions  $s(\cdot)$  and  $t(\cdot)$ , thus  $s(\cdot)$  and  $t(\cdot)$  can be arbitrary learned transforms. Since the proposed coupling enhancement utilize normal convolutions as its  $s(\cdot)$  and  $t(\cdot)$ , its receptive field is not restricted to be blocked, and can thus serve as an enhancement of expressiveness.
3.2 Right-Invertible Generalized Divisive Normalization
The widely-used generalized normalization (GDN)[Balle et al., 2015] in LIC is invertible in theory, and is thus qualified as a surjection. However, the inverse of GDN has to be solved for every input
(a) baseline framework
(b) proposed framework
Figure 2: Comparison between (a) baseline framework [Balle et al., 2018] and (b) proposed framework. To make the framework idempotent, we replace conv/deconv with right-invertible convolution (described in Sec. 3.1), GDN/iGDN with c-GDN (described in Sec. 3.2), and quantization  $Q$  with idempotent quantization  $Q_{i}$  (described in Sec. 3.3) Both frameworks use the same mean-scale Gaussian entropy model [Minnen et al., 2018]. AE/AD are arithmetic encoder/decoder, respectively.
image in an iterative manner, and is even not guaranteed to converge in finite steps. Therefore, the original GDN is not suitable for idempotent compression.
We propose a coupling GDN layer (c-GDN) that combines the coupling structure [Dinh et al., 2016] and GDN. The inverse of c-GDN layer can be solved in a much simpler analytical manner. Specifically, we implement a coupling structure [Dinh et al., 2016] with normal GDN as its scale and translation functions  $(s(\cdot)$  and  $t(\cdot))$ . And just like the aforementioned coupling enhancement, the forward and inverse of this c-GDN layer can be calculated according to Eq. 11 and Eq. 12, respectively. We demonstrate empirically in Sec. 4.3 that the proposed c-GDN layer can achieve comparable RD performance with the original GDN.
3.3 Right-Invertible Quantization
As previously discussed, the vanilla scalar quantization using rounding is naturally idempotent. However, the widely adopted mean-shift trick quantization for mean-scale Gaussian entropy model [Minnen et al., 2018] is not guaranteed to be idempotent. As proposed by [Minnen et al., 2020], the mean-shifted quantization can be formulated as
$$
Q (\boldsymbol {y}) = \left\lfloor \boldsymbol {y} - \boldsymbol {\mu} \right\rceil + \boldsymbol {\mu}, \tag {13}
$$
where  $\lfloor \cdot \rfloor$  is scalar quantization and  $\pmb{\mu}$  is the predicted mean of  $\pmb{y}$ . Let  $\pmb{y}_1$  denote the result after the first application of  $Q$ , then idempotence requires that  $Q(\pmb{y}_1) = \pmb{y}_1$ , which in turn requires  $\pmb{y}_1 - \pmb{\mu}$  to be integer. However, no existing method can meet this requirement, thus  $Q$  is not ensured to be idempotent.
We propose two types of circumvention to solve this issue. For the first-type circumvention, we change the quantization into
$$
Q _ {i} (\boldsymbol {y}) = \lfloor \boldsymbol {y} - \lfloor \boldsymbol {\mu} \rceil \rfloor + \lfloor \boldsymbol {\mu} \rceil \tag {14}
$$
By adding  $\lfloor \cdot \rfloor$  around the predicted mean  $\pmb{\mu}$ , we force the quantization result to be integer, thus  $\pmb{y} - \lfloor \pmb{\mu} \rceil$  is guaranteed to be integer from the second application of  $Q_{i}$ , and then  $Q_{i}$  is idempotent.
For the second-type circumvention, we resort to the original definition of mean-scale Gaussian entropy model [Minnen et al., 2018], and calculate the quantized CDF (cumulative distribution function) regarding  $Q_{i}(\pmb{y}) = \lfloor \pmb{y} \rceil$  on the fly during inference.
3.4 Overall Framework of Right-Invertible Codec
The overall framework is depicted in Fig. 2. Following prior works [Balle et al., 2017, 2018, Cheng et al., 2020, He et al., 2022] in LIC, We use the mainstream four-stage framework in our work. Specifically, the encoding transform is divided into four stages, and each stage decreases the resolution by a factor of 2. For the first 3 stages, each stage starts with a right-invertible convolution layer (described in Sec. 3.1) and ends with a c-GDN normalization layer (described in Sec. 3.2). The last stage only consists of one right-invertible convolution layer.
(b)
Figure 3: Performance of different codecs on Kodak. Idempotent codecs are marked as idemp, near-idempotent codec are marked as near id and non-idempotent codec are marked as non id. (a) First-time and re-compression (upto 50 times) RD performance of different codec. Idempotent codec only report first-time RD performance (re-compression RD performance is the same). (b) PSNR drop during re-compression (upto 50 re-compression) of different codec. Note that idempotent codec are straight lines and cover each other in the figure.
The decoding transform is built to be a right-inverse of the encoding transform. Note that not all layer in the decoding transform has its own weights. Specifically, for c-GDN normalization layers, there is no additional weights since they are the inverse of their counterparts in the encoding transform. For right-invertible convolution layers, the only additional weights appear in the null-space enhancement (Sec. 3.1).
Entropy model does not influence idempotence in the proposed framework, thus we use the off-the-shelf mean-scale Gaussian entropy model [Minnen et al., 2018] for its availability and efficiency. Additionally, we use the right-invertible quantization (described in Sec. 3.3) for quantizing the code.
3.5 Extension to Near-Idempotent Learned Image Codec
Idempotent codec can make sure that the RD performance keeps unchanged during any times of re-compression. However, this strict idempotence comes with a price that the decoding transform must be the right-inverse of encoding transform. This limitation reduces the expressiveness of transforms and is empirically harmful to the first-time RD performance.
To adapt to the cases where first-time RD performance also matters, we propose to extend our idempotent codec near-idempotent codec by relaxing the right-invertibility. The relaxation of right-invertibility is simple yet effective: we change the first right-invertible convolution layer (described in Sec.1) to be non-surjective, and keep all the rest layers surjective. This is done by allowing  $D$  to be smaller than  $d$  for the kernel  $K$  in the first right-invertible convolution layer. By keeping the right-invertibility of most layers, our near-idempotent codec is more stable to re-compression than existing near-idempotent codec [Kim et al., 2020, Cai et al., 2022], while achieving comparable or better first-time RD performance.
All the models are trained on the training split of open-images dataset [Kuznetsova et al., 2020], and all the evaluations are conducted on the Kodak dataset [Franzen, 1999].
We sketch the training schedule accordingly from existing literature [Balle et al., 2017, 2018, Minnen et al., 2018, 2020, Cheng et al., 2020]. Images are randomly cropped to  $256 \times 256$  for training, and a batch size of 16 is used. All the models are trained using an Adam optimizer. The learning rate is initially set to  $10^{-1}$ , and decays by a factor of 10 when plateaued.
We choose four bitrate level according to the benchmark setting in [Kim et al., 2020]. Specifically, we set  $\lambda = \{18,67,250,932\} \times 10^{-4}$ , and models trained with these  $\lambda$  reaches average bitrates from 0.2-1.5 on Kodak dataset. Following prior works [Balle et al., 2017, 2018], we use a smaller code channels (192) for lower-bpp points, and use a bigger code channels (320) for higher-bpp points. The learned function  $f(\cdot)$  in Eq.10 is implemented with a residual block.
All the experiments are conducted on a computer with AMD EPYC 7742 64-Core Processor and 8 NividA A30 GPU. All the code is implemented based Python 3.9, Pytorch 1.12 and CompressAI [Bégaint et al., 2020].
4.2 Overall Performance
4.2.1 Results of Idempotent Codec
We compare with [Helminger et al., 2021], which is the only prior LIC that achieves idempotent lossy compression to the best of our knowledge. We also compare with traditional idempotent codecs such as JPEG2000.
We report first-time compression RD performance of the above idempotent codescs in Fig. 3(a), and detailed BD-BR and BD-PSNR are listed in Tab. 1. Multi-time re-compression RD performance does not change for idempotent codescs. From the result we see that, our proposed framework exceeds prior art [Helminger et al., 2021] by a large margin, which clearly validates the superiority of right-inverse over strict inverse on the idempotent lossy compression task.
We also compare the FLOPs and encode-decode time in Tab. 1. The results clearly show that the proposed idempotent framework is also more efficient than [Helminger et al., 2021].
Table 1: The BD-BR, BD-PSNR, FLOPs and encode-decode time of different methods on Kodak dataset. FLOPs and enc-dec time are calculated on an input of shape  $256 \times 256 \times 3$ .
Methods	BD-BR (%) ↓	BD-PSNR (dB) ↑	GFLOPs ↓	time (ms) ↓
Idempotent Codec
JPEG2000	0.00	0.00	-	-
[Helminger et al., 2021]	4.83	-0.21	15.89	185
Proposed Idempotent	-28.75	1.63	8.40	110