Source-linked AI summary
Entroformer: A Transformer-based Entropy Model for Learned Image Compression
Yichen Qian, Ming Lin, Xiuyu Sun, Zhiyu Tan, Rong Jin
TL;DR
Learned image compression entropy models need stronger global-dependency modeling and faster decoding. Entroformer uses a transformer-based entropy model with top-k attention, diamond relative position encoding, and bidirectional two-pass context. It reports improved compression performance over CNN methods and BPG at low bit rates while parallel decoding is more time-efficient without performance degradation.
Problem
CNN-based entropy models are inefficient at capturing long-range dependencies, while unidirectional context models create a performance–decoding-speed trade-off.
Method
Entroformer combines transformer self-attention with top-k selection, diamond relative position encoding, and a parallel bidirectional context model for image compression.
Results
Entroformer outperforms advanced CNN methods by 5.2% and BPG by 20.5% at low bit rates, while parallel decoding is more time-efficient without performance degradation.
Takeaways & Limitations
Transformer-based entropy modeling can capture long-range dependencies for improved learned image compression while the bidirectional two-pass design supports faster GPU decoding.
Takeaways & Limitations
Future work considers hierarchical decoding, fine-grained k assignment across images, and the role of rate control in transformer-based entropy models.
Abstract
from arXiv · showhide
One critical component in lossy deep image compression is the entropy model, which predicts the probability distribution of the quantized latent representation in the encoding and decoding modules. Previous works build entropy models upon convolutional neural networks which are inefficient in capturing global dependencies. In this work, we propose a novel transformer-based entropy model, termed Entroformer, to capture long-range dependencies in probability distribution estimation effectively and efficiently. Different from vision transformers in image classification, the Entroformer is highly optimized for image compression, including a top-k self-attention and a diamond relative position encoding. Meanwhile, we further expand this architecture with a parallel bidirectional context model to speed up the decoding process. The experiments show that the Entroformer achieves state-of-the-art performance on image compression while being time-efficient.
1 Introduction
Entroformer addresses long-range dependency modeling and decoding-speed challenges in learned image compression with a transformer-based entropy model and optimized parallel context design.
- Entropy models estimate latent probability distributions for entropy coding, and improving predictive ability can increase compression rates without increasing distortion.
- CNN-based entropy models remain constrained in capturing long-range spatial dependencies across texture, semantic, and color relationships.
- Unidirectional context models improve prediction but decode symbols through an O(n) raster-scan process that cannot be accelerated by modern GPUs.
- Entroformer combines transformer self-attention with top-k selection and diamond relative position encoding to model dependencies for image compression.
- 5.2% improvement over advanced CNN methods and 20.5% over BPG were reported at low bit rates.
- Bidirectional context with two-pass decoding makes parallel Entroformer more time-efficient than serialized decoding on modern GPUs without performance degradation.
2 Compression with Hyperprior and Context
The compression model combines hyperprior and context-based entropy modeling to predict latent distributions, while optimizing coding rate and reconstruction distortion. Its entropy model uses Gaussian latents conditioned on hyperprior and decoded context.
- Entropy model: Hyperprior methods provide side information, while context models use causal latent context without additional bits; Entroformer combines both approaches.
- Model architecture: The main autoencoder learns the quantized latent representation ˆy and reconstructed image ˆx, while the hyper-autoencoder learns quantized hyper-latents ˆz.During training, quantization is approximated with additive uniform noise.
- Entropy model: A transformer encoder learns the hyperprior, a masked-attention transformer decoder learns context, and a linear component combines them to predict Gaussian parameters µ and σ.The resulting latent probabilities are conditioned on hyperprior and context.
- Latent distribution: The entropy model represents each latent ˆyi with a Gaussian distribution convolved with a unit uniform distribution.The Gaussian mean µi and deviation σi are predicted by the entropy model.
- Optimization: Training minimizes L = R + λD, balancing estimated coding rates for latents and hyper-latents against distortion under a metric such as MSE.λ controls the rate-distortion trade-off.
3 Transformer-based Entropy Model
Entroformer uses a transformer-based entropy model to represent spatial and content dependencies in latent representations, with position-aware and sparse attention designs optimized for image compression. A parallel bidirectional context model reduces serial decoding while retaining rich contextual prediction.
- 3 Transformer-based Entropy Model: The compression model combines a CNN-based autoencoder for latent representation learning with a transformer-based entropy model for latent prediction.The entropy model processes latent sequences and combines hyperprior and context features to generate Gaussian parameters.
- 3.1 Transformer Architecture: Transformer self-attention relates different latent positions to compute representations, while the entropy model uses hyperprior and autoregressive features.A linear layer combines hyperprior and context features to generate the Gaussian parameters used for probability estimation.
- 3.2 Position Encoding: Closer context latents contribute more to bitrate influence, motivating a diamond-shaped relative position boundary for spatially aware attention.The relative position is represented as a 2D coordinate, with the maximum absolute relative position constrained by a 2D diamond boundary with maximum ℓ1 distance h.
- 3.3 Top-k Scheme in Self-Attention: Top-k self-attention selects the most similar keys for each query instead of computing attention over every query-key pair.The scheme suppresses irrelevant context latents and helps stabilize training when inference sequences are much longer than training sequences.
- 3.4 Parallel Bidirectional Context Model: Causal masking restricts each latent to earlier positions, creating a unidirectional context that prevents parallel autoregressive decoding.This limitation motivates splitting latents into two slices for a two-pass decoding framework.
- 3.4 Parallel Bidirectional Context Model: The parallel bidirectional model decodes the first slice from the hyperprior and the second slice in parallel using hyperprior plus decoded context.Bidirectional context introduces future context for the second slice; the reported performance degradation is about 1%, compared with about 4% for a CNN-based accelerated method.
4 Experimental Results
Experiments on Kodak evaluate Entroformer’s rate–distortion performance, component contributions, attention choices, and decoding latency. The results show improved compression and efficient parallel decoding.
- 5.2% better performance than state-of-the-art CNN methods and 20.5% better performance than BPG at low bit rates.
- 12.4% and 11.3% improvements at low bit rates are reported for hyperprior-only and context-only models, respectively.Parallel Entroformer changes rate–distortion performance by about 1% bpp.
- Architecture: The architecture uses six transformer encoder blocks for hyperprior modeling and six transformer decoder blocks for context modeling.The latent embedding dimension is 384, and the two module outputs are combined by a linear layer.
- Position Encoding Methods: Diamond relative position encoding achieves 5.8% bpp savings over the CNN context baseline, while 2D relative encoding saves 4.9%.The best clipping distance for diamond RPE is h = 3.
- Top-k Self-attention: Top-k self-attention improves compression for k ≤64 despite sequence lengths of 576 during training and 1536 during testing.For k larger than 64, top-k attention produces results similar to original self-attention.
- Parallel Bidirectional Context Model: Parallel bidirectional context is more time-efficient than serial unidirectional context on a 16GB Tesla V100 GPU.Its compression performance is reported as on par with serial unidirectional context.
5 Discussion
The discussion presents Entroformer as a transformer-based entropy model with diamond RPE, top-k attention, and bidirectional context. It also identifies several directions for future investigation.
- Entroformer combines a novel positional encoding and top-k self-attention for learned image compression.Diamond RPE uses a 2D relative position encoding with a diamond-shaped boundary.
- Bidirectional context supports two-pass decoding at fast speed without performance degradation.
- Future work includes hierarchical decoding, image-specific assignment of k, and studying rate control in transformer-based entropy models.
A.1 PSNR (RGB) on Kodak dataset
The Kodak PSNR appendix presents rate–distortion curves for comparing image-compression models.
- The curves use PSNR on the Kodak image set, with each point averaging MS-SSIM and bit rate.Models marked ∗ use a deeper main auto-encoder.
A.2 MS-SSIM (RGB) on Kodak dataset
The Kodak MS-SSIM appendix presents rate–distortion curves using a decibel-transformed MS-SSIM measure.
- The curves use MS-SSIM on Kodak, reported in dB as MS-SSIMdB = −10 log10(1 −MS-SSIM).Each point averages MS-SSIM and bit rate; models marked ∗ use a deeper main auto-encoder.
A.3 Position Encoding for various image size
Relative position encoding generalizes directly to larger image sizes, whereas absolute position encoding degrades when processing images larger than those used during training.
- Relative position encoding directly generalizes to larger image sizes on the Tecnick dataset.
- Absolute position encoding degrades when processing image sizes larger than the training size.
A.4 Architecture details
The architecture details specify layer conventions for convolutional, upsampling, and transformer components, while model outputs define latents, hyper-latents, reconstructed images, and Gaussian parameters.
- Table 5 describes each generalized-model layer using conventions for convolutional, deconvolutional, and transformer blocks.“Conv” encodes kernel size, output channels, and stride; “Deconv” denotes upsampled convolutions; “Trans” denotes transformer blocks with inner dimension and head number.
- Table 6 compares the parameter count of different entropy-model variants.
- The encoder outputs latents ˆy, the hyper-encoder outputs hyper-latents ˆz, and the decoder outputs the generated RGB image ˆx.
- The hyperprior and context features are concatenated and passed through two Linear layers to produce Gaussian distribution parameters, means, and deviations.
A.5 Training details
The models use specified optimization and regularization settings, transformer pretraining with masked latents, and fixed training hardware and schedules across rate-distortion tradeoffs.
- Training uses Adam with β1 = 0.9, β2 = 0.999, ϵ = 1×10−8, and a base learning rate of 1×10−4.
- The learning rate warms up for 0.05 of total epochs, then decays every 1/5 epoch proportion by 0.75, with gradient clipping set to 1.0.
- No weight decay is used because the rate constraint provides inherent regularization, and Entroformer weights use truncated-normal initialization.
- Pretraining masks 50% of latents by replacing them with zero; masked latents use only hyperprior prediction, while other latents use hyperprior and context.The models are fine-tuned afterward with regular training.
- Models are trained for 300 epochs with batch size 16 and 384 × 384 patches on a 16GB Tesla V100 GPU.
- The hyperparameter λ is adapted to obtain models covering a range of rate-distortion tradeoffs.