Learned Compression for Compressed Learning#
In this chapter, we introduce WaLLoC (Wavelet Learned Lossy Compression), a neural codec architecture that combines linear transform coding with nonlinear dimensionality-reducing autoencoders. WaLLoC sandwiches a shallow, asymmetric autoencoder and entropy bottleneck between an invertible wavelet packet transform. Across several key metrics, WaLLoC outperforms the autoencoders used in state-of-the-art latent diffusion models. WaLLoC does not require perceptual or adversarial losses to represent high-frequency detail, providing compatibility with modalities beyond RGB images and stereo audio. WaLLoC’s encoder consists almost entirely of linear operations, making it exceptionally efficient and suitable for mobile computing, remote sensing, and learning directly from compressed data. We demonstrate WaLLoC’s capability for compressed-domain learning across several tasks, including image classification, colorization, document understanding, and music source separation. Our code, experiments, and pre-trained audio and image codecs are available at https://ut-sysml.org/walloc/.
Introduction#
In the last decade, deep neural networks (DNNs) have rapidly evolved from simple classifiers (Hershey et al., 2017, Krizhevsky et al., 2012) to domain-specific and multi-modal foundation models (Archit et al., 2023, Beyer et al., 2024). With this shift, models are increasingly able to make use of minute and high-frequency signal details. For example, when increasing the resolution of PaliGemma from \(224^2\) to \(896^2\) pixels (Fig. 2), its ability to analyze documents increases from 44% to 85% ANLS (Beyer et al., 2024). However, operating at this increased resolution requires significantly more GPU memory (21 vs 8 GB) and leads to 4\(\times\) higher latency.
Fig. 2 In discriminative models (left), resolution reduction increases training and inference efficiency, but significantly degrades accuracy. Replacing resolution reduction with WaLLoC leads to significantly higher accuracy, while providing the same degree of acceleration. For signal enhancement (right), WaLLoC provides better quality when scaling to high resolutions compared to directly operating on image pixels or audio samples.#
Compressed-domain learning (Ehrlich and Davis, 2019, Park et al., 2023, Rombach et al., 2022) has been proposed to improve the trade-off between model accuracy and compute needs. In this paradigm, the model operates on low-dimensional (lossy) compressed data, thereby enabling dramatic reductions in compute cost and inference latency while maintaining model accuracy. However, existing lossy compression methods, coming from three main categories, are not ideal for compressed-domain learning. (a) Linear transform coding methods (e.g., JPEG, MP3) reduce bitrate via energy-compacting time-frequency transforms, but do not meaningfully reduce dimensionality or increase efficiency of downstream models. (b) End-to-end learned codecs (Ballé et al., 2017) achieve better rate-distortion performance and modestly reduce dimension via nonlinear autoencoders, but high encoding overhead negates the benefits of compressed learning. (c) Generative autoencoders (Evans et al., 2024, Rombach et al., 2022) significantly reduce dimension, but do so by synthesizing rather than preserving details—leading to poor performance in discriminative tasks (Goldblum et al., 2024).
Fig. 3 Comparison of our proposed method (WaLLoC) with other autoencoder designs for RGB Images (Cheng2020 (Cheng et al., 2020), Stable Diffusion 3 (Esser et al., 2024)) and stereo audio (EnCodec (Défossez et al., 2023), Stable Audio (Evans et al., 2024)). Additional metrics are reported in Jacobellis and Yadwadkar (2025).#
In this chapter, we introduce WaLLoC (Wavelet Learned Lossy Compression), an architecture for learned compression that simultaneously satisfies three key requirements of compressed-domain learning:
Computationally efficient encoding to reduce overhead in compressed-domain learning and support resource constrained mobile and remote sensors. WaLLoC uses the computationally cheap and invertible wavelet packet transform (Mallat, 2008) to expose signal redundancies prior to autoencoding. This allows us to replace the encoding DNN with a single linear layer (\(<\)100k parameters) without significant loss in quality. As shown in Fig. 3, WaLLoC incurs less than five percent of the encoding cost compared to other neural codecs.
High compression ratio for storage and transmission efficiency. Lossy codecs typically achieve high compression by combining quantization and entropy coding. However, naive quantization of autoencoder latents leads to unpredictable and unbounded distortion. Instead, we apply additive noise during training as an entropy bottleneck (Ballé et al., 2017), leading to quantization-resilient latents. When combined with entropy coding, WaLLoC achieves nearly 6\(\times\) higher compression ratio compared to the VAE used in Stable Diffusion 3 (Esser et al., 2024), despite offering a higher degree of dimensionality reduction and similar quality.
Dimensionality reduction to accelerate compressed-domain modeling. WaLLoC’s encoder projects high-dimensional signal patches to low-dimensional latent representations, providing a reduction of up to 20\(\times\). This allows WaLLoC to be used as a drop-in replacement for resolution reduction while providing superior detail preservation and downstream accuracy.
Sandwiched Asymmetric Autoencoder via Wavelet Packet Transform#
Fig. 4 WaLLoC’s encode-decode pipeline. The entropy bottleneck and entropy coding steps are only required to achieve high compression ratios for storage and transmission. For compressed-domain learning where dimensionality reduction is the primary goal, these steps can be skipped to reduce overhead and completely eliminate CPU-GPU transfers.#
Fig. 5 Example of forward and inverse WPT with \(J=2\) levels. Each level applies filters \(\text{L}_{\text{A}}\) and \(\text{H}_{\text{A}}\) independently to each of the signal channels, followed by downsampling by a factor of two \(\left(\downarrow 2\right)\). An inverse level consists of upsampling \(\left(\uparrow 2\right)\) followed by \(\text{L}_{\text{S}}\) and \(\text{H}_{\text{S}}\), then summing the two channels. The full WPT \(\stackrel{\sim}{\smash{\textbf{X}}\rule{0pt}{1.1ex}}\) of consists of \(J\) levels.#
WaLLoC’s design aims at achieving three goals: computationally efficient encoding, high compression ratio, and uniform dimensionality reduction. We note several key insights that allow us to address the limitations of previous designs that stand in the way of achieving these goals. Each of these goals, limitations, and insights motivate the core design components of WaLLoC, shown in Fig. 4.
Achieving computationally efficient encoding. Two main barriers stand in the way of efficient encoding. (a) poor scaling of autoencoder performance with resolution, and (b) difficulty in preserving quality with lightweight encoders.
(a) Resolution scaling. In existing autoencoder designs (Ballé et al., 2017, Cheng et al., 2020, Défossez et al., 2023, Evans et al., 2024, Rombach et al., 2022), a hierarchy of DNN layers progressively reduce the spatial or temporal resolution while increasing the channel dimension. However, the initial layers of the encoder and the final layers of the decoder operate at the original resolution, leading to high memory and computational requirements (Beyer, 2024). The wavelet packet transform (WPT), shown in Fig. 5, is a linear and invertible transform that performs an analogous operation. In each level of the WPT, the signal is divided into high- and low-frequency components, then downsampled by a factor of two. By recursively applying this process, the WPT allows spatial and temporal resolution to be traded off for frequency resolution with minimal computation and no loss of information. In WaLLoC, we exploit this property by sandwiching the learnable analysis and synthesis transforms between the WPT and its inverse—allowing all neural network layers to operate at low resolution.
(b) Loss of quality in lightweight encoders. Previous efforts use reduced hidden dimension and distillation to reduce the cost of pixel-based autoencoders but incur a significant loss of detail in the process (Bohan, 2023). However, the WPT’s ability to isolate important signal components from redundancies alleviates this issue. Additionally, it is possible to exploit asymmetry between the encoder and decoder. The decoder objective—disentangling mixed signal components—is difficult and requires a complex DNN-based transform. In contrast, the encoder objective—discarding signal redundancies—becomes trivial after applying the WPT. Thus, WaLLoC sandwiches an asymmetric autoencoder—consisting of a shallow, linear analysis transform and a deep, nonlinear synthesis transform—between the WPT and its inverse.
Achieving high compression ratio. Quantization is the primary mechanism used in lossy compression to reduce bit rate and achieve a high compression ratio. However, the GADR-AEs that provide good dimensionality reduction are not compatible with quantization. For example, quantization of Stable Diffusion’s VAE latents leads to severe distortion (Bohan, 2023) However if quantization is applied, very high compression ratios can be achieved via entropy coding. In WaLLoC, we incorporate an entropy bottleneck—additive noise applied during training that guarantees quantization resilience during inference (Ballé et al., 2017). We optimize the noise scale for 8-bit quantization, allowing us to use standard lossless codecs (e.g PNG or WeBP) for entropy coding. This combination provides an additional compression multiplier of up to \(12\times\) compared to reducing the dimension only.
Achieving uniform dimensionality reduction. In addition to quantization, neural codecs achieve high compression ratios via a loss term that encourages sparse, rather than low-dimensional latents (Ballé et al., 2017). Using this objective, it is possible to drive the energy of many of the latent dimensions to zero (He et al., 2022). However, this type of non-uniform dimensionality reduction is difficult to exploit in compressed-domain learning. In WaLLoC, the analysis transform uniformly reduces the dimension by a fixed rate, making it a suitable replacement for resolution reduction in accelerating downstream models.
Implementation: WPT, Entropy Bottleneck, and Entropy Coding#
WaLLoC’s encoder consists of five stages as shown in Fig. 4: (1) wavelet packet transform (WPT) to trade-off spatial or temporal resolution with channel resolution (2) learned analysis transform to reduce dimensionality (3) companding to whiten the latent distribution (4) entropy bottleneck to provide resilience to quantization and (5) entropy coding to provide high compression ratios. The decoder consists of the reverse operations: (5) entropy decoding, (6) decompanding, (7) learned synthesis transform, and (8) inverse WPT. We now provide detailed explanations for each component.
Wavelet packet transform. Fig. 5 shows the workflow of the wavelet packet transform (WPT) and its inverse. We use the Cohen–Daubechies–Feauveau (CDF) 9/7 wavelet (Mallat, 2008) to construct a dyadic filterbank consisting of highpass analysis (\(\text{H}_{\text{A}}\)), lowpass analysis (\(\text{L}_{\text{A}}\)), highpass synthesis (\(\text{H}_{\text{S}}\)), and lowpass synthesis (\(\text{L}_{\text{S}}\)) filters. The CDF 9/7 wavelet is chosen for its balance between computational efficiency and energy compaction. Since these same filters are used in the JPEG 2000 standard, they are widely supported in software. The WPT reduces the input resolution \(R_{\textbf{x}}\) and increases the input channel count \(C_{\textbf{x}}\) by a factor \(2^J\) for 1D signals (audio) and by \(4^J\) for 2D signals (images), but is linear and invertible. For stereo audio, we use \(J=8\), resulting in \(C_{\stackrel{\sim}{\smash{\textbf{X}}\rule{0pt}{0.8ex}}}=512\) channels after the WPT. For RGB images, we use \(J=3\), resulting in \(C_{\stackrel{\sim}{\smash{\textbf{X}}\rule{0pt}{0.8ex}}}=192\).
Autoencoder and entropy bottleneck. The output of the WPT \(\stackrel{\sim}{\smash{\textbf{X}}\rule{0pt}{1.1ex}}\) is projected to a latent representation \(\textbf{z}\) via a learnable analysis transform \(\mathcal{G}_{\text{A}}\), which consists of a single linear layer. The latent dimension \(C_\textbf{z}\) is a hyperparameter chosen based on the desired degree of dimensionality reduction. To achieve quantization-resilient latent representations, we adopt the entropy bottleneck method from end-to-end learned compression (Ballé et al., 2017), which consists of adding uniform noise \(\mathcal{U}[-0.5,0.5]\) to the latent representation during training. Since the sub-band wavelet coefficients of many natural signals follow a generalized Gaussian distribution (GGD) (Westerink et al., 1991), we apply the Gaussian CDF \(\Phi(\textbf{z})\) as a companding operation prior to the entropy bottleneck. Thus, the final encoder output is \(\hat{\textbf{z}}_{\text{t}}=\Phi(z)+\mathcal{U}\) during training and \(\hat{\textbf{z}}_{\text{c}}=\text{round}\left(\Phi(\textbf{z})\right)\) during the compression pipeline. We scale the inputs and outputs of the companding operation \(\Phi\) to guarantee latents in the range [-127, 127], which in turn guarantees that \(\hat{\text{z}}_{\text{c}}\) does not underflow or overflow when quantized to a signed 8-bit integer. The decoder consists of a learnable synthesis transform \(\mathcal{G}_{\text{S}}\) followed by the IWPT. \(\mathcal{G}_{\text{S}}\) is a convolutional neural network consisting of the same residual blocks used in Stable Audio (Evans et al., 2024) and Stable Diffusion 3 (Esser et al., 2024) for 1D and 2D signals respectively. We use a hidden dimension of \(C_{\text{hidden}}=768\) for both the RGB image decoder and stereo audio decoders. Additional implementation details are available in our public code repositories [1].
Entropy coding. After quantization, an additional lossless compression step can be applied. We performed preliminary tests using zlib (Foundation, n.d.), PNG (Deflate), and the lossless mode of WebP (Clark and others, 2016). We found that WebP’s entropy coding provided the best compression ratio—even for audio signals—while maintaining high throughput and compatibility with ML frameworks like PyTorch. Since WebP expects 24-bit RGB inputs, we rearrange the multi-channel 8-bit latent tensor into groups of three and concatenate channel groups along the temporal or spatial dimensions.
Training. We train four codecs—two for stereo audio (5\(\times\), 20\(\times\)) and two for RGB images (4\(\times\), 20\(\times\))—on The lossless MUSDB18-HQ (Rafii et al., 2017) and LSDIR (Li et al., 2023) datasets. In each case, the training objective is to minimize mean squared reconstruction error when latents are subjected to uniform additive noise in the range [-0.5,0.5].
Evaluation#
We conduct a comprehensive evaluation of WaLLoC to demonstrate its efficacy for compressed domain learning. Our evaluation consists of two main parts. (1) Compression trade-off analysis. We compare WaLLoC against other lossy codecs in terms of the trade-off between dimensionality reduction, compression ratio, distortion, perceptual quality, and computation (Compression trade-off analysis). (2) Compressed learning and resolution scaling. We train and evaluate various machine learning models on representations produced by WaLLoC, and compare their resolution scaling properties to pixel-based and sample-based versions (Compressed learning and resolution scaling).
Compression trade-off analysis#
We compare WaLLoC against other popular conventional and neural codecs (Cheng et al., 2020, Défossez et al., 2023, Esser et al., 2024, Evans et al., 2024) across five key metrics: (1) degree of dimensionality reduction, (2) compression ratio, (3) distortion, (4) perceptual quality, and (5) computation. For images, distortion is measured via PSNR and MS-SSIM (Wang et al., 2004), while perceptual quality is evaluated via LPIPS(Zhang et al., 2018) and DISTS (Ding et al., 2020). For audio, distortion is measured via PSNR, SSDR, and SRDR (Watcharasupat and Lerch, 2024), and perceptual quality is evaluated via CDPAM (Manocha et al., 2021). For both audio and images, the computational cost is measured in terms of average encoding and decoding throughput (megapixels or megasamples per second). Measurements are made on three different platforms: Low-power CPU (Raspberry Pi), High-power CPU (Intel i9), and GPU (RTX 4090).
Results of compression trade-off analysis. Fig. 3 summarizes the trade-offs between rate, distortion, perception, computation, and dimension between different types of compression. For RGB Images, WaLLoC achieves nearly 12\(\times\) higher compression ratio (35:1 vs 6:1) compared to the VAE used in Stable Diffusion 3, despite offering a higher degree of dimensionality reduction (16\(\times\) vs 12\(\times\)) and similar quality (13.9 dB vs 13.8 dB DISTS). Compared to Cheng et al. (Cheng et al., 2020), WaLLoC achieves more than 48\(\times\) higher encoding throughput (14.0 vs 0.29 MPix/sec) and similar quality (19.3 dB vs 16.9 dB DISTS). For stereo audio, WaLLoC achieves significantly higher spatial quality (22.5 dB vs 15.7 dB SSDR) than Stable Audio’s VAE, but with more than \(300\times\) higher encoding throughput. Examples of decoded images from the LSDIR validation set are provided on Hugging Face [2]. Additional results, including GPU and Raspberry Pi throughput, are available in our code repository [3].
Compressed learning and resolution scaling#
Applications, models, and datasets. We evaluate WaLLoC on 4 machine perception tasks: (1) image classification, (2) image colorization, (3) document understanding and (4) music source separation. For classification and colorization, we train ViT-Ti models with conditional position encoding (Tu et al., 2022) on the ImageNet-1k dataset. For music source separation, we train a CNN to separate the vocal track from music segments in MUSDB18-HQ. The CNN consists of 12 identical convolutional layers structured identically to Stable Audio’s mid block (Evans et al., 2024). For document understanding, we use PaliGemma (Beyer et al., 2024) fine-tuned at varying resolution on the DocVQA (Mathew et al., 2021) dataset, and report the average normalized levenshtein similarity (ANLS) on the test set.
Resolution scaling strategy. For image classification, we reduce the input sequence length by 4\(\times\) or 16\(\times\) compared to the baseline of \(256^2\) pixels and \(16^2\) patches, but keep the area of each patch constant (\(1/16^2\)). We report the accuracy of models trained on reduced resolution inputs with models trained on the identically WaLLoC latents. For document understanding, training models on the scale of PaliGemma is outside the scope of this work. Instead, we evaluate on decoded WaLLoC representations using the highest-resolution PaliGemma variant (\(896^2\)). To emulate the effect of resolution reduction with this high-resolution variant, we downsample images to the desired resolution, (\(224^2\) or \(448^2\)), then apply Lanczos resampling to interpolate back to \(896^2\). For Image colorization and music source separation, we increase the input patch size proportionally to the resolution to keep the sequence length—and therefore the required computation—roughly constant.
Results of compressed-domain learning and resolution scaling. Fig. 2 shows the improvement in performance when using WaLLoC-derived representations instead of resolution reduction. Across each of the four tasks, WaLLoC provides superior accuracy to naive resolution reduction while providing the same improvement in latency and memory consumption. For discriminative models, WalloC profoundly increases accuracy of efficient image classification (50.6% vs 23.1% accuracy) and document understanding (81.1 vs 43.7 ANLS). For signal enhancement, WaLLoC provides superior scaling to high resolution and large patches—offering a 16.7% improvement in colorization LPIPS and a 3.1 dB improvement in PSNR for source separation.
References#
Anwai Archit, Luca Freckmann, Sushmita Nair, Nabeel Khalid, Paul Hilt, Vikas Rajashekar, Marei Freitag, Carolin Teuber, Melanie Spitzner, Constanza Tapia Contreras, and others. Segment anything for microscopy. bioRxiv, 2023.
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. In ICLR. 2017.
Lucas Beyer. On the speed of ViTs and CNNs. lb.eyer.be/a/vit-cnn-speed, 2024.
Lucas Beyer, Andreas Steiner, André Susano Pinto, Alexander Kolesnikov, Xiao Wang, Daniel Salz, Maxim Neumann, Ibrahim Alabdulmohsin, Michael Tschannen, Emanuele Bugliarello, and others. Paligemma: a versatile 3b vlm for transfer. arXiv:2407.07726, 2024.
O. B. Bohan. Taesd: tiny autoencoder for stable diffusion. 2023. URL: madebyollin/taesd.
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto. Learned image compression with discretized gaussian mixture likelihoods and attention modules. In CVPR. 2020.
missing journal in clark2016pillow
Keyan Ding, Keda Ma, Shiqi Wang, and Eero P Simoncelli. Image quality assessment: unifying structure and texture similarity. IEEE transactions on pattern analysis and machine intelligence, 2020.
Alexandre Défossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. High fidelity neural audio compression. Transactions on Machine Learning Research, 2023.
Max Ehrlich and Larry S Davis. Deep residual learning in the jpeg transform domain. In ICCV. 2019.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Pino, Axel Hoogeboom, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In ICML. 2024.
Zach Evans, Julian Parker, CJ Carr, Zack Zukowski, Jordan Taylor, and Jordi Pons. Stable audio open. arXiv:2407.14358, 2024.
Python Software Foundation. Zlib — compression compatible with gzip.
Micah Goldblum, Hossein Souri, Renkun Ni, Manli Shu, Viraj Prabhu, Sven Gowal, and Andrew Gordon Wilson. Battle of the backbones: a large-scale comparison of pretrained models across computer vision tasks. NeurIPS, 2024.
Dailan He, Ziming Yang, Weikun Peng, Rui Ma, Hongwei Qin, and Yan Wang. Elic: efficient learned image compression with unevenly grouped space-channel contextual adaptive coding. In CVPR. 2022.
Shawn Hershey, Sourish Chaudhuri, Daniel PW Ellis, Jort F Gemmeke, Aren Jansen, R Channing Moore, Manoj Plakal, Devin Platt, Rif A Saurous, Bryan Seybold, and others. Cnn architectures for large-scale audio classification. In ICASSP. IEEE, 2017.
missing journal in jacobellis2025liveaction
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. NeurIPS, 2012.
Yawei Li, Yulun Fan, Xiaoyu Yu, Joshua Batson, Kai Qian, Eirikur Agustsson, and Radu Timofte. Lsdir: a large scale dataset for image restoration. In CVPR. 2023.
Pranay Manocha, Zeyu Jin, Richard Tarr, and Gautham J Mysore. Cdpam: contrastive learning for perceptual audio similarity. In ICASSP. IEEE, 2021.
Minesh Mathew, Dimosthenis Karatzas, and CV Jawahar. Docvqa: a dataset for vqa on document images. In IEEE/CVF winter conference on applications of computer vision. 2021.
Song Park, Sanghyuk Chun, Byeongho Heo, Wonjae Kim, and Sangdoo Yun. Seit: storage-efficient vision training with tokens using 1% of pixel storage. In CVPR. 2023.
missing journal in rafii2017musdb18
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. CVPR, 2022.
Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, and Yinxiao Li. Maxvit: multi-axis vision transformer. In ECCV. 2022.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, 2004.
Karn N Watcharasupat and Alexander Lerch. Quantifying spatial audio quality impairment. In ICASSP. 2024.
Peter H Westerink, Jan Biemond, and Dick E Boekee. Subband coding of color images. Subband Image Coding, 1991.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In CVPR. 2018.