Compressing image encoders via latent distillation

TL;DR

The paper tackles the practicality of bulky end-to-end learned image codecs in hardware-constrained settings by introducing latent-space distillation to train compact encoders while preserving the decoder. It employs an asymmetric knowledge distillation framework: a lightweight encoder is trained to mimic the teacher's latent representations, with Factorized Prior distilling only the latent y and Hyperprior distilling both y and z, leaving the decoder fixed. The approach yields better rate-distortion performance than training lightweight encoders from scratch, achieving near-original results at substantial encoder reductions even with limited data, and it delivers significant reductions in MACs and parameter counts. This work enables deployment of learned image codecs on edge devices and IoT contexts, with future work aiming to extend the method to encoder and decoder reductions together.

Abstract

Deep learning models for image compression often face practical limitations in hardware-constrained applications. Although these models achieve high-quality reconstructions, they are typically complex, heavyweight, and require substantial training data and computational resources. We propose a methodology to partially compress these networks by reducing the size of their encoders. Our approach uses a simplified knowledge distillation strategy to approximate the latent space of the original models with less data and shorter training, yielding lightweight encoders from heavyweight ones. We evaluate the resulting lightweight encoders across two different architectures on the image compression task. Experiments show that our method preserves reconstruction quality and statistical fidelity better than training lightweight encoders with the original loss, making it practical for resource-limited environments.

Figures (4)

• Figure 1: Direct encoder reduction does not preserve performance. Green curves show models with encoder widths reduced by $\div 4$, and $\div 2$, trained on $10\%$ of the dataset, while the pink curve denotes the full model. The reduced models underperform, confirming that simple encoder compression limits model capacity.
• Figure 2: Original ($\mathcal{T}$) and reduced ($\mathcal{S}$) Factorized prior architectures. The fire icon denotes the trainable part of $\mathcal{S}$, while the snowflake indicates the frozen part.
• Figure 3: After encoder reduction, our models $\mathcal{M}_{r,KD}^{\rho}$ achieve lower FID and higher PSNR than $\mathcal{M}_{r,Frozen}^{\rho}$. PSNR results are shown in figures (a) and (d) for the Factorized Prior, and in figures (b) and (e) for MS-ILLM, while FID results for MS-ILLM are presented in figures (c) and (f). All evaluations are conducted on CLIC2020 with $\rho \in {0.1, 10.0}$. For $\rho = 10.0$, the reduced architectures $r \in {2,4}$ achieve performance comparable to $\mathcal{M}_{Orig}$. Some $\mathcal{M}_{r,Frozen}^{01}$ points were omitted due to low performance.
• Figure 4: The $\mathcal{M}_{r,KD}^{\rho}$ models produce higher-quality reconstructions. We present examples of MS-ILLM with the lowest bit-rate.