LC-TMNet: learned lossless medical image compression with tunable multi-scale network

Likelihood-based generative models

Arithmetic coding requires obtaining the probability of each symbol; hence, early lossless compression commonly employed autoregressive models. In such models, each pixel is compressed based on previously encoded pixels. For instance, PixelRNN (Van Den Oord, Kalchbrenner & Kavukcuoglu, 2016) processes image pixels sequentially using a recurrent neural network (RNN), while PixelCNN (Van den Oord et al., 2016) models pixel values using a masked convolutional neural network (CNN). Both methods predict pixels based on the conditional distribution of all preceding pixels, modeling pixels as the product of conditional distributions p(x) = ∏p(x_i∣x₁, …, x_i−1), where xi represents a single pixel.

PixelCNN++ (Salimans et al., 2017), an improved version of PixelCNN, reduces dependence on previous channels by modeling the joint distribution of each pixel. Additionally, PixelCNN++ introduces discrete logistic mixture likelihood, multi-scale downsampling, and extra shortcut connections, which enhance compression performance and reduce processing time. Despite these improvements, PixelCNN++ retains the inherent limitation of autoregressive models, requiring network computation for each pixel, resulting in a time complexity of O(W × H), where W is the width and H is the height, leading to prolonged inference times.

Multi-scale entropy models

To optimize compression efficiency and address time constraints. L3C has made pioneering contributions in this field. L3C uses a hierarchical model to predict all pixels of an entire image. Specifically, L3C considers the image X itself as the zeroth layer and utilizes the first latent space Z₁ to construct a probability model P(X|Z₁). Similarly, compressing Z₁ requires the second latent space Z₂ , with the probability expression P(Z₁|Z₂), and so on. However, since the last layer of features Z₃ does not have a subsequent layer of features to serve as a prior, P(Z₃) is set to a fixed value. This fixed value usually deviates significantly from the actual probability, resulting in poor compression performance for the last layer of features Z₃.

To address the lack of prior knowledge, residual compressor (RC) (Mentzer, Van Gool & Tschannen, 2020) and deep lossy plus residual (DLPR) (Bai et al., 2024) approach the problem from the perspective of lossy compression. They use a lossy image as prior knowledge to predict the residual probability of the image, with the probability model P(X_residual|X_lossy). By combining the residual with the lossy image, they achieve lossless compression. These experiments provide crucial insights: a low bit-rate lossy image offers prior knowledge, transforming the original image into a residual image. Leveraging the concentrated nature of residuals improves the accuracy of probability prediction P.

In LC-FNet (Rhee et al., 2022), the authors deeply investigated residual characteristics from a frequency domain perspective and proposed a frequency decomposition network. This approach assumes that residual probability estimation in low-frequency regions is relatively straightforward, whereas high-frequency regions are more challenging. Thus, LC-FNet first compresses the low-frequency regions before processing the high-frequency regions, leveraging prior information from the low-frequency regions to improve the probability estimation of high-frequency regions. The probabilistic model is formulated as P(X_{high−frequency}|X_{low−frequency}). Wang et al. (2023b) proposed a novel decomposition approach based on image storage bytes, contrasting LC-FNet’s frequency method. They split the image into the most significant byte (MSB) subimage for higher-order bits and the least significant byte (LSB) subimage for lower-order bits. The MSB, requiring lower compression bitrate, is encoded traditionally and used as input to a neural network to predict and encode the LSB.

In this paper, we are primarily inspired by LC-FNet, PixelRNN, and RC. While these methods show distinct advantages in specific scenarios, they still face limitations regarding computational complexity, compression efficiency, and flexibility. Building on these insights, this paper proposes a novel multi-scale compression framework that introduces subimage-level lossless priors and binary-tree decomposition strategy to overcome some of the bottlenecks of existing approaches. The specific improvements are summarized in Table 1.

LC-H architecture

The Fig. 3 shows the LC-H architecture. We will detail it in three parts: Preprocessing, Residual Image, and Residual Probability Distribution.

LC-H-Extend architecture

The Fig. 4 shows the LC-H architecture.The The LC-H-Extend module is an extension of the LC-H architecture, designed to retain the even columns of $X_{residual}\in R^{1\times \frac{W}{2}\times \frac{H}{2}}$ while masking the odd columns. Specifically, the module uses a $Mask\in R^{1\times \frac{W}{2}\times \frac{H}{2}}$ to keep the even columns of X_residual unchanged and set the odd columns to a fixed value 0. The mask is given by Eq. (1): (1)${\displaystyle Mask=\left\{\begin{array}{cc}1\hfill & \text{if}W\text{is even}\hfill \\ 0\hfill & \text{if}W\text{is odd}\hfill \\ \hfill \end{array}.\right.}$

By applying X_residual⋅Mask = X_{left−residual}, where all odd columns are 0. Consequently, experiment set the probability at the 0th position (out of 511 positions) in the odd columns of the probability matrix $P_{residual}\in R^{511\times \frac{W}{2}\times \frac{H}{2}}$ to 1, indicating that the probability of odd column values being zero is 100%. Simultaneously, positions 1 to 510 in the odd columns are set to 0, indicating the probability of odd column values being non-zero is 0%. This yields the new probability matrix P_{left−residual}, as defined in Eq. (2). (2)${\displaystyle P_{left-residual}=\left\{\begin{array}{cc}{P}_{residual}\hfill & \text{if}W\text{is even}\hfill \\ 1\left(100\text{\%}\right)\hfill & \text{if}W\text{is odd and at position 0}\hfill \\ 0\left(0\text{\%}\right)\hfill & \text{if}W\text{is odd and at positions 1 to 510}\hfill \\ \hfill \end{array}.\right.}$

LC-L architecture

Figure 5 shows the LC-L architecture. We will detail it in three parts: Preprocessing, Residual Image, and Residual Probability Distribution.

Attention-enhanced UNet architecture

As shown in Fig. 6, the UNet architecture (Ronneberger, Fischer & Brox, 2015) is enhanced with SE attention (Hu, Shen & Sun, 2018) to improve sensitivity to key information and probability estimation. After four downsampling and upsampling stages, the Softmax function ensures the predicted probabilities sum to 1.

During the upsampling process, we introduce the SE attention mechanism. Firstly, apply global average pooling F_sq to the original matrix X ∈ R^C×W×H to obtain the initial weight matrix Z ∈ R^C×1×1. The mathematical expression for F_sq is given by Eq. (3) (3)${\displaystyle Z_c=\frac{1}{W\times H}\sum _{i=1}^W\sum _{j=1}^H{X}_{c,i,j}}$ where Z_c is the channel-wise average of the original matrix X. Next, the initial weight matrix Z is transformed into the final weight matrix S through F_ex, which is a fully connected layer. The mathematical expression for F_ex is provided in Eq. (4). (4)${\displaystyle S=F_{ex}\left(Z\right)=\sigma \left(W_2\cdot \text{ReLU}\left(W_1\cdot Z\right)\right)}$ where W₁ and W₂ are weight matrices of the fully connected layers, ReLU denotes the ReLU activation function, and σ denotes the sigmoid activation function. Finally, the final weight matrix S is used in the scaling operation F_scale, where it is multiplied channel-wise with the original matrix X to obtain the weighted matrix $\hat{X}$. The mathematical formulation for F_scale is shown in Eq. (5). (5)${\displaystyle {\hat{X}}_{c,i,j}=F_{scale}\left(X,S\right)=S_c\cdot X_{c,i,j}}$ where $\hat{X}$ is the weighted matrix obtained by scaling the original matrix X with the final weight vector S.

Loss function

The loss functions for LC-H and LC-L include the subimage prediction loss (SPL) and the probability prediction loss (PPL). LC-H-Extend does not include any loss functions, as it is an extension of LC-H and does not require separate loss functions.

Experimental setup

Performance comparisons

Table 2 presents the comparative results on the specified evaluation set. Evidently, our method, Propose (fast), outperforms both existing conventional codecs and learning-based codecs while Propose (slow) performs even better than Propose (fast). Compared to the traditional algorithm JPEG-XL, Propose (slow) shows an average improvement of 9.0%; compared to the learning-based algorithm DLPR, Propose (slow) averages a 4.5% improvement.

Furthermore, we found that in ultrasound and natural image datasets, our method, Propose (slow), significantly outperforms DLPR. However, in CT and X-ray datasets, the advantage of Propose (slow) is relatively smaller. This is because imaging methods like ultrasound produce images with more complex texture structures. The LC-L module used by Propose (slow) has more accurate inference capabilities for complex texture areas, hence showing superior performance in datasets with complex textures. In contrast, CT and X-ray imaging methods produce smoother images with fewer complex textures, which limits the scope for compression improvements by Propose (slow).

Time comparisons

Table 3 displays the timing comparison results on the specified evaluation set, with experiments conducted on an Nvidia 2080Ti and Intel 12700. In tests across various types of datasets, Propose (fast) consistently demonstrated the fastest compression speed, highlighting its high practicality for real-world applications.

The speed advantage of Propose (fast) is attributed to the use of the LC-H module, which can obtain all residuals and residual probability distributions of subimages in a single inference. In contrast, Propose (slow) employs both LC-L and LC-H-Extend for more detailed subimages compression, resulting in Propose (slow) requiring twice the number of inferences as Propose (fast). Although Propose (slow) is slower than Propose (fast), it maintains a reasonable speed while offering state-of-the-art compression efficiency. Thus, Propose (slow) remains a highly efficient mode worth adopting.

The experiments did not include comparisons with traditional algorithms because traditional algorithms rely on CPU computations, whereas learning-based algorithms primarily depend on GPU computations, leading to a lack of comparability between the two. Additionally, the L-Infinite algorithm does not support GPU operations, to ensure fairness, it was excluded from the comparisons.

Probability heatmap analysis

According to Shannon’s information entropy theory, the optimal number of bits is −log₂P(j), where j is the residual value to be compressed, and P(j) is the probability of j predicted by the neural network. A higher P(j) leads to better compression. In Fig. 7, different colors correspond to different values of the residual probability P. Red indicates P is close to 0, meaning poor compression, while green indicates P is close to 1, meaning excellent compression.

Figure 7A shows the probability heatmap of an ultrasound image. Due to the unique characteristics of ultrasound imaging, the texture is highly intricate. It can also be observed that red spots are widely distributed, indicating that the neural network often struggles with probability estimation in complex textured images. However, after applying the LC-H-Extend and LC-L modules, the number of red spots significantly decreases. This suggests that our Propose (slow) compression scheme can achieve more accurate probability estimation in regions with complex textures, effectively addressing the issue of inaccurate probability estimation in these areas.

Figure 7B shows the probability heatmap of an X-ray image. Since the lungs are primarily filled with air, their texture is relatively simple. In Fig. 7B, the number of red spots is very low, regardless of whether LC-L or LC-H is used. This demonstrates that the neural network is highly efficient in reasoning when dealing with smooth images. After using the LC-H-Extend and LC-L combination for compression, the green areas become even more prominent, further proving that the Propose (slow) compression scheme can still enhance performance on smooth images.

Using the ultrasound image in Fig. 7A as an example, the experiment qualitatively analyzed residual probability values by dividing them into three ranges: greater than 0.1, between 0.01 and 0.1, and less than 0.01. Table 4 shows that after applying the LC-L and LC-H-Extend modules, the proportion of pixels with residual probability greater than 0.1 increases, while those below 0.01 decrease. This confirms that our Propose (slow) compression scheme improves probability estimation, especially in complex textured regions.

LC-L ablation experiment

The ablation experiment is analyzed from two perspectives: performance improvement for a single pixel and performance improvement for the entire subimage. As shown in Fig. 8, the experiment selects the residual point at the coordinates (788,919) as an example. The residual obtained by LC-H is 195, with a corresponding probability P_LC−H(195) = 0.00257. The residual obtained by LC-L is 220, with a corresponding probability P_LC−L(220) = 0.00779. According to Shannon’s information entropy formula, bit =  − log₂P(j), the compression bit count for LC-H is 8.6 bits, and for LC-L, it is 7.4 bits. At (788,919), LC-L saves approximately 14% of the space compared to LC-H.

To evaluate the overall bits per pixel (bpp) improvement of LC-L on subimages (D, B, C), the experiment compares Model 1 using LC-H and Model 2 using LC-L. The network is evaluated on the Breast Ultrasound Images Dataset, excluding the JPEG-XL part (1.05 bpp). Table 5 shows the ablation results for LC-L. Model 2 shows the greatest improvement on subimage D and the smallest improvement on subimage C. The reason for this difference lies in the input context of Model 1 during compression. When compressing subimage C, Model 1′s input includes subimages A, D, and B, providing ample prior knowledge to accurately estimate probability in complex texture regions. However, when compressing subimage D, Model 1′s input only includes subimage A, lacking sufficient prior knowledge, resulting in inaccurate probability inference. In this scenario, Model 2′s input includes both subimages A and D-Left, with D-Left providing strong prior knowledge for D-Right, significantly enhancing probability estimation accuracy. Overall, using LC-L improved the model’s performance by 5.1%.