Wavefront Parallelization for Efficient Learned Image Compression

2026-07-21

AI Summary

This work addresses the slow inference speed of autoregressive context models in learned image compression, a limitation exacerbated by existing acceleration methods that require architectural modifications and retraining, thus hindering compatibility with pretrained models. For the first time, this study adapts the wavefront parallelization concept from video coding to this domain, proposing a training-free inference acceleration algorithm. By reordering computations according to an interleaved wavefront schedule, the method substantially reduces serial steps while strictly preserving autoregressive dependencies and rate-distortion performance. The approach is plug-and-play, applicable to any pretrained autoregressive model without modification. Experiments on Cheng et al.'s models demonstrate over 13× decoding speedup, with further gains achievable through controlled relaxation of context dependencies.

Abstract

Autoregressive context models are foundational for learned image compression,but they suffer from slow serial inference. Existing acceleration methods such as checkerboard context require architectural changes and retraining, thus are inapplicable to pre-trained models. We propose a completely training-free inference-time acceleration algorithm inspired by wavefront parallelism in video coding standards. Our method reorganizes inference into an optimal ``staggered'' wavefront order, minimizing sequential steps while maintaining exact autoregressive dependencies. Experimental results show our approach accelerates pre-trained autoregressive models (e.g., Cheng et al.) by more than $13\times$ while preserving the original rate-distortion performance. We also demonstrate that faster decoding is possible by trading off precise context dependencies. Source code will be available at https://github.com/tokkiwa/compressai-wavefront.

Problem

• learned image compression
• autoregressive context models
• inference acceleration
• pre-trained models
• serial inference

Innovation

• wavefront parallelization
• learned image compression
• autoregressive context
• inference acceleration
• training-free